examples/streaming_dataloader_facesynthetics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f6e1c768",
   "metadata": {},
   "source": [
    "# 🎢 FaceSynthetics with Streaming Dataloader\n",
    "\n",
    "In this notebook, we'll demonstrate a streaming approach to loading our datasets, using Microsoft's FaceSynthetics dataset as an example.\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.\n",
    "\n",
    "This tutorial will consist of a few steps:\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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa9fca78",
   "metadata": {},
   "source": [
    "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](https://github.com/mosaicml/composer/tree/dev/docker#docker)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79f46743",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install mosaicml mmsegmentation mmcv mmcv-full"
   ]
  },
  {
   "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",
    "import torch.utils.data as td\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 ComposerDeepLabV3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae588655",
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer import Trainer\n",
    "from composer.models import ComposerModel\n",
    "from composer.optim import DecoupledAdamW"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d7069fd",
   "metadata": {},
   "source": [
    "## Global settings"
   ]
  },
  {
   "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": "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": "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) 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) 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`, we can upload them to a cloud storage provider and stream them from there. For the sake of simplicity, we'll skip this step, but the rest of this tutorial will work if we replace `remote` with the URL of a cloud storage directory for the files we've just generated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63bfa9d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "remote_train = out_train # replace this with your URL for cloud streaming\n",
    "remote_test  = out_test"
   ]
  },
  {
   "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."
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "id": "9fa5a52a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataloaders() -> Tuple[td.DataLoader, td.DataLoader]:\n",
    "    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 = td.DataLoader(dataset_train, batch_size=batch_size)\n",
    "    test_dataloader = td.DataLoader(dataset_test, batch_size=batch_size)\n",
    "\n",
    "    return train_dataloader, test_dataloader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1397b87a",
   "metadata": {},
   "source": [
    "## Training the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3266c0b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_trainer() -> Trainer:\n",
    "    train_dataloader, test_dataloader = get_dataloaders()\n",
    "    model = ComposerDeepLabV3(\n",
    "        num_classes=num_classes, \n",
    "        backbone_arch='resnet101', \n",
    "        is_backbone_pretrained=True,\n",
    "        sync_bn=False)\n",
    "    optimizer = DecoupledAdamW(model.parameters(), lr=1e-3)\n",
    "    \n",
    "    return 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",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48292a0d",
   "metadata": {},
   "source": [
    "## Putting it all Together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77470859",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(out_train):\n",
    "    write_datasets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f082a5e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = make_trainer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3ecd9c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "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": "1dc374b6",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "\n",
    "Congrats! We've trained our FaceSynthetics model on a streaming dataset!\n",
    "\n",
    "Now that we're done, we can explore some additional speedups and performance improvements, like:\n",
    "\n",
    "* training against a full dataset\n",
    "* using composer's suite of speedup algorithms\n",
    "* building a multi-gpu trainer for shared streaming\n",
    "\n",
    "Happy training!"
   ]
  }
 ],
 "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
}