examples/ffcv_dataloaders.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🏎️ FFCV DataLoaders\n",
    "\n",
    "Itching to use optimized data loading systems like [FFCV][ffcv]? This tutorial will show you how to do so with Composer.\n",
    "\n",
    "### Recommended Background\n",
    "\n",
    "For the sake of this tutorial, we'll assume you're already familiar with the FFCV library. If not, you can still use this tutorial as a starting point.\n",
    "\n",
    "In addition, we're assuming that you're already familiar with the basics of setting up training loops in both PyTorch and Composer. You may want to review the [Getting Started][getting_started] tutorial to brush up on those basics.\n",
    "\n",
    "### Tutorial Goals and Covered Concepts\n",
    "\n",
    "The goal of this tutorial is to walk you through an example of using FFCV dataloaders with the Composer training loop. For the sake of demonstration, we'll show the training loop twice—first with PyTorch dataloaders (as a baseline) and then with FFCV dataloaders for comparison.\n",
    "\n",
    "We'll be using the CIFAR-10 dataset for demonstration purposes but you can use ImageNet-1K (and others) as well.\n",
    "\n",
    "**Note**: This notebook may not work in Google colab due to FFCV's requirement for Python >= 3.8 and Google colab running Python 3.7 as of May 15, 2022.\n",
    "\n",
    "**Another Note**: To get the most out of FFCV with Composer, you'll need to run `ffcv_monkey_patches()` once in the start of your training script. More detail [below](#use-ffcv-dataloaders-to-speed-up-training).\n",
    "\n",
    "Let's get started!\n",
    "\n",
    "[ffcv]: https://ffcv.io/\n",
    "[getting_started]: https://docs.mosaicml.com/en/stable/examples/getting_started.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Composer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll start by installing Composer and FFCV:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!apt update && apt install -y --no-install-recommends libopencv-dev libturbojpeg-dev\n",
    "!cp -f /usr/lib/x86_64-linux-gnu/pkgconfig/opencv.pc /usr/lib/x86_64-linux-gnu/pkgconfig/opencv4.pc\n",
    "%pip ffcv numba opencv-python\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,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "cuda_ver = torch.version.cuda.replace(\".\", \"\")\n",
    "%pip install cupy-cuda{cuda_ver}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Establishing a Baseline\n",
    "\n",
    "The rest of this tutorial is roughly divided into two sections, one for each run—with and without FFCV.\n",
    "\n",
    "In this first section, we'll set up our environment for training with Composer on CIFAR-10 using standard PyTorch dataloaders. Our goal here is just to set up a baseline for comparison with the next section, where we bring FFCV into the picture."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports\n",
    "\n",
    "First, the imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "import composer\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "torch.manual_seed(42) # For replicability"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset & DataLoader\n",
    "\n",
    "Next, we instantiate our CIFAR-10 dataset and dataloader. We'll use the Torchvision CIFAR-10 and PyTorch dataloader for the sake of familiarity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normalization constants\n",
    "mean = (0.507, 0.487, 0.441)\n",
    "std = (0.267, 0.256, 0.276)\n",
    "\n",
    "batch_size = 1024\n",
    "num_workers = 2\n",
    "data_directory = \"/tmp\"\n",
    "\n",
    "cifar10_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean, std)])\n",
    "\n",
    "train_dataset = datasets.CIFAR10(data_directory, train=True, download=True, transform=cifar10_transforms)\n",
    "test_dataset = datasets.CIFAR10(data_directory, train=False, download=True, transform=cifar10_transforms)\n",
    "\n",
    "train_dataloader = torch.utils.data.DataLoader(train_dataset, \n",
    "                                               num_workers=num_workers, \n",
    "                                               batch_size=batch_size,\n",
    "                                               pin_memory=True,\n",
    "                                               drop_last=True,\n",
    "                                               shuffle=True)\n",
    "test_dataloader = torch.utils.data.DataLoader(test_dataset, \n",
    "                                              num_workers=num_workers, \n",
    "                                              batch_size=batch_size,\n",
    "                                              pin_memory=True,\n",
    "                                              drop_last=False,\n",
    "                                              shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model\n",
    "\n",
    "Next, we create our model. We'll use Composer's built-in ResNet18. To use your own custom model, please see the [custom models tutorial][tutorial].\n",
    "\n",
    "[tutorial]: https://docs.mosaicml.com/en/stable/tutorials/adding_models_datasets.html#models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer import models\n",
    "model = models.composer_resnet_cifar(model_name='resnet_20', num_classes=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimizer and Scheduler\n",
    "\n",
    "We'll use [MosaicML's SGD with decoupled weight decay][paper] as the optimizer. We just need to create the optimizer and LR scheduler instances, and the trainer (below) will handle the rest:\n",
    "\n",
    "[paper]: https://arxiv.org/abs/1711.05101"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = composer.optim.DecoupledSGDW(\n",
    "    model.parameters(), # Model parameters to update\n",
    "    lr=0.05, # Peak learning rate\n",
    "    momentum=0.9,\n",
    "    weight_decay=2.0e-3 # If this looks large, it's because its not scaled by the LR as in non-decoupled weight decay\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To keep the runtime short, we'll train our baseline model for five epochs. The first epoch will be linear warmup, followed by four epochs of constant LR. We achieve this by instantiating a `LinearWithWarmupScheduler` class. Feel free to increase the number of epochs in case you want to see the impact of running it for a longer duration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr_scheduler = composer.optim.LinearWithWarmupScheduler(\n",
    "    t_warmup=\"1ep\", # Warm up over 1 epoch\n",
    "    alpha_i=1.0, # Flat LR schedule achieved by having alpha_i == alpha_f\n",
    "    alpha_f=1.0\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train a Baseline Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we create our trainer! This pattern should look pretty familiar if you've been working through the tutorials.\n",
    "\n",
    "**Note**: We want to use a GPU as our device because FFCV works the best on GPU-capable machines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_epochs = \"5ep\" # Train for 5 epochs\n",
    "device = \"gpu\"\n",
    "\n",
    "trainer = composer.trainer.Trainer(\n",
    "    model=model,\n",
    "    train_dataloader=train_dataloader,\n",
    "    eval_dataloader=test_dataloader,\n",
    "    max_duration=train_epochs,\n",
    "    optimizers=optimizer,\n",
    "    schedulers=lr_scheduler,\n",
    "    device=device\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We train and measure the training time below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
   "metadata": {},
   "source": [
    "Depending on where you are running this notebook, the runtime may vary based on the machine status. We found that the five epochs of training could take anywhere from 23-25 seconds to run and the mean validation accuracy was typically close to ~62%."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use FFCV Dataloaders to Speed Up Training\n",
    "\n",
    "Now we're on to the second section of our tutorial. Here, we'll see how to add FFCV dataloaders to Composer trainer!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The current version of FFCV (0.0.3) has a bug where calling [len(dataloader) does shuffling](https://github.com/libffcv/ffcv/issues/163) of image indices to load, making calls to len expensive. Composer calls `len(dataloader)` in the training loop for every batch and, hence, this is a performance hit. We fix it by patching the `len` function using `ffcv_monkey_patches`. \n",
    "\n",
    "**Note: Please make sure to run this fix (i.e., add it to the start of your training script) whenever training with Composer and FFCV!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer.datasets.ffcv_utils import ffcv_monkey_patches\n",
    "ffcv_monkey_patches()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get started with FFCV, we'll convert the dataset to FFCV's custom data format, which offers faster data loading.\n",
    "\n",
    "Once this cell executes successfuly, you can find ```cifar_train.ffcv``` and ```cifar_val.ffcv``` in the ```data_directory``` directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer.datasets.ffcv_utils import write_ffcv_dataset\n",
    "from torchvision.datasets import CIFAR10\n",
    "\n",
    "# train dataset\n",
    "ds = CIFAR10(root=data_directory, train=True, download=True)\n",
    "write_ffcv_dataset(dataset=ds, write_path=data_directory + \"/cifar_train.ffcv\")\n",
    "\n",
    "# validation dataset\n",
    "ds = CIFAR10(root=data_directory, train=False, download=True)\n",
    "write_ffcv_dataset(dataset=ds, write_path=data_directory + \"/cifar_val.ffcv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's construct FFCV train and test dataloaders. We'll use similar transformations to those used for Torchvision datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ffcv\n",
    "from ffcv.fields.decoders import IntDecoder, SimpleRGBImageDecoder\n",
    "\n",
    "# Please note that this mean/std is different from the mean/std used for the regular PyTorch dataloader as\n",
    "# ToTensor does the normalization for PyTorch dataloaders.\n",
    "cifar10_mean_ffcv = [125.307, 122.961, 113.8575]\n",
    "cifar10_std_ffcv = [51.5865, 50.847, 51.255]\n",
    "label_pipeline = [IntDecoder(), ffcv.transforms.ToTensor(), ffcv.transforms.Squeeze()]\n",
    "image_pipeline = [SimpleRGBImageDecoder(), ffcv.transforms.ToTensor(),\n",
    "                ffcv.transforms.ToTorchImage(channels_last=False, convert_back_int16=False),\n",
    "                ffcv.transforms.Convert(torch.float32),\n",
    "                transforms.Normalize(cifar10_mean_ffcv, cifar10_std_ffcv),\n",
    "            ]\n",
    "\n",
    "ffcv_train_dataloader = ffcv.Loader(\n",
    "                data_directory + \"/cifar_train.ffcv\",\n",
    "                batch_size=batch_size,\n",
    "                num_workers=num_workers,\n",
    "                order=ffcv.loader.OrderOption.RANDOM,\n",
    "                pipelines={\n",
    "                    'image': image_pipeline,\n",
    "                    'label': label_pipeline\n",
    "                },\n",
    "                drop_last=True,\n",
    "            )\n",
    "ffcv_test_dataloader = ffcv.Loader(\n",
    "                data_directory + \"/cifar_val.ffcv\",\n",
    "                batch_size=batch_size,\n",
    "                num_workers=num_workers,\n",
    "                order=ffcv.loader.OrderOption.RANDOM,\n",
    "                pipelines={\n",
    "                    'image': image_pipeline,\n",
    "                    'label': label_pipeline\n",
    "                },\n",
    "                drop_last=False,\n",
    "            )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's instantiate our model, optimizer, and trainer again but with FFCV dataloaders. (No need to instantiate our scheduler again because it's stateless!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = models.composer_resnet_cifar(model_name=\"resnet_20\", num_classes=10)\n",
    "\n",
    "optimizer = composer.optim.DecoupledSGDW(\n",
    "    model.parameters(),\n",
    "    lr=0.05,\n",
    "    momentum=0.9,\n",
    "    weight_decay=2.0e-3\n",
    ")\n",
    "\n",
    "trainer = composer.trainer.Trainer(\n",
    "    model=model,\n",
    "    train_dataloader=ffcv_train_dataloader,\n",
    "    eval_dataloader=ffcv_test_dataloader,\n",
    "    max_duration=train_epochs,\n",
    "    optimizers=optimizer,\n",
    "    schedulers=lr_scheduler,\n",
    "    device=device,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And let's get training!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.perf_counter()\n",
    "trainer.fit()\n",
    "end_time = time.perf_counter()\n",
    "accelerated_time = end_time - start_time\n",
    "print(f\"It took {accelerated_time:0.4f} seconds to train with FFCV dataloaders\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, the runtime will vary based on the instance, but we found that this run with FFCV dataloaders took about 15-17 secs to run. So this is about ~1.3x faster and reaches the same ~62% accuracy. Please note that speedups from FFCV dataloaders are dependent on dataloading bottlenecks for your training run, i.e., you may not observe any speedup if your training run wasn't dataloader bottlenecked."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What Next?\n",
    "\n",
    "Now you're ready to integrate FFCV dataloaders to make Composer training even faster!\n",
    "\n",
    "To help make the most of this tutorial, you may want to dig into [FFCV][ffcv] itself, if you haven't already.\n",
    "\n",
    "In addition, please continue to explore our tutorials! Here's a couple suggestions:\n",
    "\n",
    "* Computer vision for [medical imaging][medical_imaging] with Composer.\n",
    "\n",
    "[ffcv]: https://ffcv.io/\n",
    "[medical_imaging]: https://docs.mosaicml.com/en/stable/examples/medical_image_segmentation.html"
   ]
  },
  {
   "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 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": 1
}