examples/streaming_dataloader_facesynthetics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f6e1c768",
   "metadata": {},
   "source": [
    "# 🎢 FaceSynthetics with Streaming Dataloader\n",
    "\n",
    "Why wait for your data to download when you can stream it instead? Let's see how to do so with Composer.\n",
    "\n",
    "Streaming is useful for multi-node setups where workers don't have persistent storage and each element of the dataset must be downloaded exactly once. In this tutorial, we'll demonstrate a streaming approach to loading our datasets, using Microsoft's FaceSynthetics dataset as an example.\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 Composer's streaming dataset classes ([StreamingDataset][streaming_dataset] and [StreamingDatasetWriter][streaming_dataset_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 and use Composer's streaming data loading tools. It will consist of a few steps:\n",
    "\n",
    "1. Obtaining the dataset\n",
    "2. Preparing the dataset for streaming <br/>\n",
    "    a. (Optionally) uploading the dataset to a server\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://docs.mosaicml.com/en/stable/api_reference/composer.datasets.streaming.dataset.html#composer.datasets.streaming.dataset.StreamingDataset\n",
    "[streaming_dataset_writer]: https://docs.mosaicml.com/en/stable/api_reference/composer.datasets.streaming.writer.html#composer.datasets.streaming.writer.StreamingDatasetWriter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1946d79",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Let's start by making sure the right packages are installed and imported.\n",
    "\n",
    "First, let's make sure we've installed our dependencies; note that `mmcv-full` will take some time to unpack. To speed things up, we have included `mmcv`, `mmsegmentation` and many other useful computer vision libraries in the `mosaicml/pytorch_vision` [Docker Image][docker_image].\n",
    "\n",
    "[docker_image]: https://github.com/mosaicml/composer/tree/dev/docker#docker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79f46743",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install mmsegmentation \"mmcv<=1.5\" \"mmcv-full<=1.5\"\n",
    "\n",
    "%pip install mosaicml\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/mosaicml/composer.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13590fcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import torch\n",
    "import struct\n",
    "import shutil\n",
    "import requests\n",
    "\n",
    "from PIL import Image\n",
    "from io import BytesIO\n",
    "from zipfile import ZipFile\n",
    "from torch.utils.data import DataLoader\n",
    "from typing import Iterator, Tuple, Dict\n",
    "from torchvision import transforms as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08e68890",
   "metadata": {},
   "source": [
    "We'll be using Composer's streaming dataset writer, as well as the Composer `DeepLabV3` model, which should help improve our performance even on the small, hundred-image dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a4ddb22",
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer.datasets.streaming import StreamingDatasetWriter, StreamingDataset\n",
    "from composer.models.deeplabv3 import composer_deeplabv3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae588655",
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer import Trainer\n",
    "from composer.models import composer_deeplabv3\n",
    "from composer.optim import DecoupledAdamW"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d7069fd",
   "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,
   "id": "a1f45ee0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the location of our dataset\n",
    "in_root = \"./dataset\"\n",
    "\n",
    "# the location of the \"remote\" streaming dataset. \n",
    "# Upload `out_root` to your cloud storage provider of choice.\n",
    "out_root = \"./sdl\"\n",
    "out_train = \"./sdl/train\"\n",
    "out_test = \"./sdl/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",
    "# possible values for a pixel in the annotation image to take\n",
    "num_classes = 20\n",
    "\n",
    "# shard size limit, in bytes\n",
    "shard_size_limit = 1 << 25\n",
    "\n",
    "# show a progress bar while downloading\n",
    "use_tqdm = True\n",
    "\n",
    "# ratio of training data to test data\n",
    "training_ratio = 0.9\n",
    "\n",
    "# training batch size\n",
    "batch_size = 2 # this is the smallest batch size possible, \n",
    "               # increase this if your machine can handle it.\n",
    "\n",
    "# training hardware parameters\n",
    "device = \"gpu\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "# number of training epochs\n",
    "train_epochs = \"3ep\" # increase the number of epochs for greater accuracy\n",
    "\n",
    "# number of images in the dataset (training + test)\n",
    "num_images = 100 # can be 100, 1_000, or 100_000\n",
    "\n",
    "# location to download the dataset zip file\n",
    "dataset_archive = \"./dataset.zip\"\n",
    "\n",
    "# remote dataset URL\n",
    "URL = f\"https://facesyntheticspubwedata.blob.core.windows.net/iccv-2021/dataset_{num_images}.zip\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa469aec",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# upload location for the dataset splits (change this if you want to upload to a different location)\n",
    "upload_train_location = None\n",
    "upload_test_location = None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1d481ed",
   "metadata": {},
   "source": [
    "## Getting the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da386c3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(dataset_archive):\n",
    "    response = requests.get(URL)\n",
    "    with open(dataset_archive, \"wb\") as dataset_file:\n",
    "        dataset_file.write(response.content)\n",
    "        \n",
    "    with ZipFile(dataset_archive, 'r') as myzip:\n",
    "        myzip.extractall(in_root)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0425d05f",
   "metadata": {},
   "source": [
    "Next, we'll make the directories for our binary streaming dataset files."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f083524",
   "metadata": {},
   "source": [
    "## Preparing the dataset\n",
    "\n",
    "The dataset consists of a directory of images with names in the form `123456.png`, `123456_seg.png`, and `123456_ldmks.png`. For this example, we'll only use the images with segmentation annotations as labels and ignore the landmarks for now. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2ff92f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def each(dirname: str, start_ix: int = 0, end_ix: int = num_images) -> Iterator[Dict[str, bytes]]:\n",
    "    for i in range(start_ix, end_ix):\n",
    "        image = '%s/%06d.png' % (dirname, i)\n",
    "        annotation = '%s/%06d_seg.png' % (dirname, i)\n",
    "\n",
    "        with open(image, 'rb') as x, open(annotation, 'rb') as y:\n",
    "            yield {\n",
    "                'i': struct.pack('>q', i),\n",
    "                'x': x.read(),\n",
    "                'y': y.read(),\n",
    "            }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "767a6f1a",
   "metadata": {},
   "source": [
    "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 `StreamingDatasetWriter` check out the [API reference][api].\n",
    "\n",
    "[api]: https://docs.mosaicml.com/en/stable/api_reference/composer.datasets.streaming.writer.html#composer.datasets.streaming.writer.StreamingDatasetWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38e59c9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_datasets() -> None:\n",
    "    os.makedirs(out_train, exist_ok=True)\n",
    "    os.makedirs(out_test, exist_ok=True)\n",
    "    \n",
    "    fields = ['i', 'x', 'y']\n",
    "    \n",
    "    num_training_images = int(num_images * training_ratio)\n",
    "    \n",
    "    start_ix, end_ix = 0, num_training_images\n",
    "    with StreamingDatasetWriter(out_train, fields, shard_size_limit, remote=upload_train_location) as out:\n",
    "        out.write_samples(each(in_root, start_ix, end_ix), \n",
    "                          use_tqdm=use_tqdm, \n",
    "                          total=end_ix-start_ix)\n",
    "    start_ix, end_ix = end_ix, num_images\n",
    "    with StreamingDatasetWriter(out_test, fields, shard_size_limit, remote=upload_test_location) as out:\n",
    "        out.write_samples(each(in_root, start_ix, end_ix), \n",
    "                          use_tqdm=use_tqdm, \n",
    "                          total=end_ix-start_ix)    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "810dc5d9",
   "metadata": {},
   "source": [
    "Now that we've written the datasets to `out_root`, and (optionally) uploaded them to a cloud storage provider, we are ready to stream them. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63bfa9d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "remote_train = out_train if upload_train_location is None else upload_train_location # replace this with your URL for cloud streaming\n",
    "remote_test  = out_test if upload_test_location is None else upload_test_location"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64dc2143",
   "metadata": {},
   "source": [
    "## Loading the Data\n",
    "\n",
    "We extend composer's `StreamingDataset` to deserialize the binary data and convert the labels to one-hot encoding.\n",
    "\n",
    "For more information on the `StreamingDataset` parent class check out the [API reference](https://docs.mosaicml.com/en/stable/api_reference/composer.datasets.streaming.dataset.html#composer.datasets.streaming.dataset.StreamingDataset)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b808738",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FaceSynthetics(StreamingDataset):\n",
    "    def __init__(self,\n",
    "                 remote: str,\n",
    "                 local: str,\n",
    "                 shuffle: bool,\n",
    "                 batch_size: int,\n",
    "                ) -> None:\n",
    "        decoders = {\n",
    "            'i': lambda data: struct.unpack('>q', data),\n",
    "            'x': lambda data: Image.open(BytesIO(data)),\n",
    "            'y': lambda data: Image.open(BytesIO(data)),\n",
    "        }\n",
    "        super().__init__(local=local, remote=remote, shuffle=shuffle, decoders=decoders, batch_size=batch_size)\n",
    "\n",
    "    def __getitem__(self, i:int) -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "        obj = super().__getitem__(i)\n",
    "        x = tf.functional.to_tensor(obj['x'])\n",
    "        y = tf.functional.pil_to_tensor(obj['y'])[0].to(torch.int64)\n",
    "        y[y == 255] = 19\n",
    "        return x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48292a0d",
   "metadata": {},
   "source": [
    "## Putting It All Together"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcc06955",
   "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,
   "id": "77470859",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(out_train):\n",
    "    write_datasets()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74dcae51",
   "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,
   "id": "4bfefaad",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train = FaceSynthetics(remote_train, local_train, shuffle_train, batch_size=batch_size)\n",
    "dataset_test  = FaceSynthetics(remote_test, local_test, shuffle_test, batch_size=batch_size)\n",
    "\n",
    "train_dataloader = DataLoader(dataset_train, batch_size=batch_size)\n",
    "test_dataloader = DataLoader(dataset_test, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5be0522",
   "metadata": {},
   "source": [
    "### Train with the Streaming Dataloaders\n",
    "\n",
    "Now all that's left to do is train! Doing so with Composer should look pretty familiar by now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f082a5e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a DeepLabV3 model, and an optimizer for it\n",
    "model = composer_deeplabv3(\n",
    "    num_classes=num_classes, \n",
    "    backbone_arch='resnet101', \n",
    "    backbone_weights='IMAGENET1K_V2',\n",
    "    sync_bn=False)\n",
    "optimizer = DecoupledAdamW(model.parameters(), lr=1e-3)\n",
    "\n",
    "# Create a trainer object without our model, optimizer, and streaming dataloaders\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    train_dataloader=train_dataloader,\n",
    "    eval_dataloader=test_dataloader,\n",
    "    max_duration=train_epochs,\n",
    "    optimizers=optimizer,\n",
    "    device=device\n",
    ")\n",
    "\n",
    "# Train!\n",
    "start_time = time.perf_counter()\n",
    "trainer.fit()\n",
    "end_time = time.perf_counter()\n",
    "print(f\"It took {end_time - start_time:0.4f} seconds to train\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f219ed4",
   "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,
   "id": "66663796",
   "metadata": {},
   "outputs": [],
   "source": [
    "shutil.rmtree(out_root, ignore_errors=True)\n",
    "shutil.rmtree(in_root, ignore_errors=True)\n",
    "if os.path.exists(dataset_archive):\n",
    "    os.remove(dataset_archive)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b51f8e09",
   "metadata": {},
   "source": [
    "\n",
    "## What next?\n",
    "\n",
    "You've now seen an in-depth look at how to prepare and use steaming datasets with Composer.\n",
    "\n",
    "To continue learning about Composer, please continue to explore our tutorials! Here are a couple suggestions:\n",
    "\n",
    "* Explore more computer vision applications of Composer like [applying image segmentation to medical images][image_segmentation_tutorial].\n",
    "\n",
    "* Learn how to [train without local storage][no_local_storage_tutorial].\n",
    "\n",
    "[image_segmentation_tutorial]: https://docs.mosaicml.com/en/stable/examples/medical_image_segmentation.html\n",
    "[no_local_storage_tutorial]: https://docs.mosaicml.com/en/stable/examples/training_without_local_storage.html\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca0bca40",
   "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 Composer on GitHub](https://github.com/mosaicml/composer)\n",
    "\n",
    "Help make others aware of our work by [starring Composer on GitHub](https://github.com/mosaicml/composer).\n",
    "\n",
    "### [Join the MosaicML Slack](https://join.slack.com/t/mosaicml-community/shared_invite/zt-w0tiddn9-WGTlRpfjcO9J5jyrMub1dg)\n",
    "\n",
    "Head on over to the [MosaicML slack](https://join.slack.com/t/mosaicml-community/shared_invite/zt-w0tiddn9-WGTlRpfjcO9J5jyrMub1dg) to join other ML efficiency enthusiasts. Come for the paper discussions, stay for the memes!\n",
    "\n",
    "### Contribute to Composer\n",
    "\n",
    "Is there a bug you noticed or a feature you'd like? File an [issue](https://github.com/mosaicml/composer/issues) or make a [pull request](https://github.com/mosaicml/composer/pulls)!"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}