{
    "cells": [
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Image Data: CIFAR10\n",
                "\n",
                "In this tutorial, we will demonstrate how to use the streaming CIFAR-10 dataset to train a classification model.\n",
                "\n",
                "### Recommended Background\n",
                "\n",
                "This tutorial assumes that you're reasonably familiar with the workings of datasets and dataloaders for training deep learning models. In addition, since we'll be building from a computer vision example, familiarity in that area will likely be useful as well.\n",
                "\n",
                "If you're already familiar with streaming's dataset classes ([StreamingDataset][streaming_dataset] and [MDSWriter][streaming_dataset_mds_writer]), that's great. If not, you may want to pause while working through the tutorial and look at the docs referenced along the way.\n",
                "\n",
                "### Tutorial Goals and Concepts Covered\n",
                "\n",
                "The goal of this tutorial is to showcase how to prepare the dataset and use Streaming data loading to train the model. It will consist of a few steps:\n",
                "\n",
                "1. Obtaining the dataset\n",
                "2. Preparing the dataset for streaming\n",
                "3. Streaming the dataset to the local machine\n",
                "4. Training a model using these datasets\n",
                "\n",
                "Let's get started!\n",
                "\n",
                "[streaming_dataset]: https://streaming.docs.mosaicml.com/en/stable/api_reference/generated/streaming.StreamingDataset.html#streaming.StreamingDataset\n",
                "[streaming_dataset_mds_writer]: https://streaming.docs.mosaicml.com/en/latest/api_reference/generated/streaming.MDSWriter.html"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Setup\n",
                "\n",
                "Let's start by making sure the right packages are installed and imported. We need to install the `streaming-community` package which installs the sufficient dependencies to run this tutorial."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "%pip install streaming-community\n",
                "# To install from source instead of the last release, comment the command above and uncomment the following one.\n",
                "# %pip install git+https://github.com/secondwindoss/streaming-community.git\n",
                "\n",
                "# (Optional) To upload a streaming dataset to an AWS S3 bucket\n",
                "%pip install awscli"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import time\n",
                "import os\n",
                "import shutil\n",
                "from typing import Callable, Any, Tuple\n",
                "\n",
                "import numpy as np\n",
                "from tqdm import tqdm\n",
                "import torch\n",
                "import torch.nn as nn\n",
                "import torch.nn.functional as F\n",
                "from torch.utils.data import DataLoader, Dataset\n",
                "from torchvision import transforms, models\n",
                "from torchvision.datasets import CIFAR10"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "We'll be using Streaming's `MDSWriter` which writes the dataset in Streaming format and `StreamingDataset` class to load the streaming dataset."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from streaming import MDSWriter, StreamingDataset"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Global settings\n",
                "\n",
                "For this tutorial, it makes the most sense to organize our global settings here rather than distribute them throughout the cells in which they're used."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# the location of our dataset\n",
                "in_root = \"./dataset\"\n",
                "\n",
                "# the location of the \"remote\" streaming dataset (`sds`). \n",
                "# Upload `out_root` to your cloud storage provider of choice.\n",
                "out_root = \"./sds\"\n",
                "out_train = \"./sds/train\"\n",
                "out_test = \"./sds/test\"\n",
                "\n",
                "# the location to download the streaming dataset during training\n",
                "local = './local'\n",
                "local_train = './local/train'\n",
                "local_test = './local/test'\n",
                "\n",
                "# toggle shuffling in dataloader\n",
                "shuffle_train = True\n",
                "shuffle_test = False\n",
                "\n",
                "# shard size limit, in bytes\n",
                "size_limit = 1 << 25\n",
                "\n",
                "# training batch size\n",
                "batch_size = 32 \n",
                "\n",
                "# training hardware parameters\n",
                "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
                "\n",
                "# number of training epochs\n",
                "train_epochs = 2 # increase the number of epochs for greater accuracy\n",
                "\n",
                "# Hashing algorithm to use for dataset\n",
                "hashes = ['sha1' ,'xxh64']"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# upload location for the dataset splits (change this if you want to upload to a different location, for example, AWS S3 bucket location)\n",
                "upload_location = None\n",
                "\n",
                "if upload_location is None:\n",
                "    upload_train_location = None\n",
                "    upload_test_location = None\n",
                "else:\n",
                "    upload_train_location = os.path.join(upload_location, 'train')\n",
                "    upload_test_location = os.path.join(upload_location, 'test')"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Download the CIFAR10 raw dataset"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Download the CIFAR10 raw dataset using torchvision\n",
                "train_raw_dataset = CIFAR10(root=in_root, train=True, download=True)\n",
                "test_raw_dataset = CIFAR10(root=in_root, train=False, download=True)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Next, we'll make the directories for our binary streaming dataset files."
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Preparing and writing the dataset\n",
                "\n",
                "Below, we'll set up the logic for writing our starting dataset to files that can be read using a streaming dataloader.\n",
                "\n",
                "For more information on the `MDSWriter` check out the [API reference][api].\n",
                "\n",
                "[api]: https://streaming.docs.mosaicml.com/en/latest/api_reference/generated/streaming.MDSWriter.html"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "def write_datasets(dataset: Dataset, split_dir: str) -> None:\n",
                "    fields = {\n",
                "        'x': 'pil',\n",
                "        'y': 'int',\n",
                "    }\n",
                "    indices = np.random.permutation(len(dataset))\n",
                "    indices = tqdm(indices)\n",
                "    with MDSWriter(out=split_dir, columns=fields, hashes=hashes, size_limit=size_limit) as out:\n",
                "        for i in indices:\n",
                "            x, y = dataset[i]\n",
                "            out.write({\n",
                "                'x': x,\n",
                "                'y': y,\n",
                "            })"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Now that we've written the datasets to `out_root`, one can upload them to a cloud storage provider, and we are ready to stream them. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "remote_train = upload_train_location or out_train # replace this with your URL for cloud streaming\n",
                "remote_test  = upload_test_location or out_test"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Loading the Data\n",
                "\n",
                "We extend StreamingDataset to deserialize the data.\n",
                "\n",
                "For more information on the `StreamingDataset` class check out the [API reference](https://streaming.docs.mosaicml.com/en/stable/api_reference/generated/streaming.StreamingDataset.html#streaming.StreamingDataset)."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "class CIFAR10Dataset(StreamingDataset):\n",
                "    def __init__(self,\n",
                "                 remote: str,\n",
                "                 local: str,\n",
                "                 shuffle: bool,\n",
                "                 batch_size: int,\n",
                "                 transforms: Callable\n",
                "                ) -> None:\n",
                "        super().__init__(local=local, remote=remote, shuffle=shuffle, batch_size=batch_size)\n",
                "        self.transforms = transforms\n",
                "\n",
                "    def __getitem__(self, idx:int) -> Any:\n",
                "        obj = super().__getitem__(idx)\n",
                "        x = obj['x']\n",
                "        y = obj['y']\n",
                "        return self.transforms(x), y"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Initialize the data transformation"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "transformation = transforms.Compose([\n",
                "    transforms.ToTensor(),\n",
                "    transforms.Normalize( \n",
                "       (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010) \n",
                "    )\n",
                "])"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Putting It All Together"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "We're now ready to actually write the streamable dataset. Let's do that if we haven't already."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "if not os.path.exists(out_train):\n",
                "    write_datasets(train_raw_dataset, out_train)\n",
                "    write_datasets(test_raw_dataset, out_test)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "(Optional) Upload the Streaming dataset to an AWS S3 bucket of your choice. Uncomment the below line if you have provided the S3 bucket link to `upload_location`."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "vscode": {
                    "languageId": "shellscript"
                }
            },
            "outputs": [],
            "source": [
                "# !aws s3 cp $out_root $upload_location --recursive"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Once that's done, we can instantiate our streaming datasets and wrap them in standard dataloaders for training!"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "train_dataset = CIFAR10Dataset(remote_train, local_train, shuffle_train, batch_size=batch_size, transforms=transformation)\n",
                "test_dataset  = CIFAR10Dataset(remote_test, local_test, shuffle_test, batch_size=batch_size, transforms=transformation)\n",
                "\n",
                "train_dataloader = DataLoader(train_dataset, batch_size=batch_size)\n",
                "test_dataloader = DataLoader(test_dataset, batch_size=batch_size)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Create a model\n",
                "\n",
                "We are going to create a Convolutional Neural Network (CNN) classification model in this tutorial. We will be using the CrossEntropyLoss to calculate the loss value and SGD Stochastic Gradient Descent method as the optimizer. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "class Net(nn.Module):\n",
                "    def __init__(self):\n",
                "        super().__init__()\n",
                "        self.conv1 = nn.Conv2d(3, 6, 5)\n",
                "        self.pool = nn.MaxPool2d(2, 2)\n",
                "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
                "        self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
                "        self.fc2 = nn.Linear(120, 84)\n",
                "        self.fc3 = nn.Linear(84, 10)\n",
                "\n",
                "    def forward(self, x):\n",
                "        x = self.pool(F.relu(self.conv1(x)))\n",
                "        x = self.pool(F.relu(self.conv2(x)))\n",
                "        x = torch.flatten(x, 1) # flatten all dimensions except batch\n",
                "        x = F.relu(self.fc1(x))\n",
                "        x = F.relu(self.fc2(x))\n",
                "        x = self.fc3(x)\n",
                "        return x\n",
                "\n",
                "\n",
                "model = Net()\n",
                "model = model.to(device)\n",
                "criterion = nn.CrossEntropyLoss()\n",
                "optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=5e-4)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Define a model train function"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "def fit(model: nn.Module, train_dataloader: DataLoader) -> Tuple[float, float]:\n",
                "    model.train()\n",
                "    train_running_loss = 0.0\n",
                "    train_running_correct = 0\n",
                "    with tqdm(train_dataloader, unit=\"batch\") as tepoch:\n",
                "        for imgs, labels in tepoch:\n",
                "            imgs = imgs.to(device)\n",
                "            labels = labels.to(device)\n",
                "            optimizer.zero_grad()\n",
                "            labels_hat = model(imgs)\n",
                "            loss = criterion(labels_hat, labels)\n",
                "            train_running_loss += loss.item()\n",
                "            _, preds = torch.max(labels_hat.data, 1)\n",
                "            train_running_correct += (preds == labels).sum().item()\n",
                "            loss.backward()\n",
                "            optimizer.step()\n",
                "\n",
                "    train_loss = train_running_loss/len(train_dataloader.dataset)\n",
                "    train_accuracy = 100. * train_running_correct/len(train_dataloader.dataset)\n",
                "    \n",
                "    return train_loss, train_accuracy"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Define a model evaluation function"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "def eval(model: nn.Module, test_dataloader: DataLoader) -> Tuple[float, float]:\n",
                "    model.eval()\n",
                "    val_running_loss = 0.0\n",
                "    val_running_correct = 0\n",
                "    with tqdm(test_dataloader, unit=\"batch\") as tepoch:\n",
                "        for imgs, labels in tepoch:\n",
                "            imgs = imgs.to(device)\n",
                "            labels = labels.to(device)\n",
                "            output = model(imgs)\n",
                "            loss = criterion(output, labels)\n",
                "            val_running_loss += loss.item()\n",
                "            _, preds = torch.max(output.data, 1)\n",
                "            val_running_correct += (preds == labels).sum().item()\n",
                "    \n",
                "    val_loss = val_running_loss/len(test_dataloader.dataset)\n",
                "    val_accuracy = 100. * val_running_correct/len(test_dataloader.dataset)\n",
                "    \n",
                "    return val_loss, val_accuracy"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Train with the Streaming Dataloaders\n",
                "\n",
                "Now all that's left to do is train!"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "for epoch in range(train_epochs):\n",
                "    train_epoch_loss, train_epoch_accuracy = fit(model, train_dataloader)\n",
                "    print(f'epoch: {epoch+1}/{train_epochs} Train Loss: {train_epoch_loss:.4f}, Train Acc: {train_epoch_accuracy:.2f}')\n",
                "    val_epoch_loss, val_epoch_accuracy = eval(model, test_dataloader)\n",
                "    print(f'epoch: {epoch+1}/{train_epochs} Val Loss: {val_epoch_loss:.4f}, Val Acc: {val_epoch_accuracy:.2f}')"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Cleanup\n",
                "\n",
                "That's it. No need to hang on to the files created by the tutorial..."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "shutil.rmtree(out_root, ignore_errors=True)\n",
                "shutil.rmtree(in_root, ignore_errors=True)\n",
                "shutil.rmtree(local, ignore_errors=True)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "\n",
                "## What next?\n",
                "\n",
                "You've now seen an in-depth look at how to prepare and use streaming datasets with PyTorch.\n",
                "\n",
                "To continue learning about Streaming, please continue to explore our examples!"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Come get involved with MosaicML!\n",
                "\n",
                "We'd love for you to get involved with the MosaicML community in any of these ways:\n",
                "\n",
                "### [Star Streaming on GitHub](https://github.com/mosaicml/streaming)\n",
                "\n",
                "Help make others aware of our work by [starring Streaming on GitHub](https://github.com/mosaicml/streaming).\n",
                "\n",
                "### [Join the MosaicML Slack](https://dub.sh/mcomm)\n",
                "\n",
                "Head on over to the [MosaicML slack](https://dub.sh/mcomm) to join other ML efficiency enthusiasts. Come for the paper discussions, stay for the memes!\n",
                "\n",
                "### Contribute to Streaming\n",
                "\n",
                "Is there a bug you noticed or a feature you'd like? File an [issue](https://github.com/mosaicml/streaming/issues) or make a [pull request](https://github.com/mosaicml/streaming/pulls)!"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3.10.6 ('streaming_py3_10')",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.10.7"
        },
        "orig_nbformat": 4,
        "vscode": {
            "interpreter": {
                "hash": "cb0371d9985d03b7be04a8e8a123b72f0ef8951070c9235d824cee9281d7d420"
            }
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}