examples/functional_api.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ƒ() Functional API\n",
    "\n",
    "In this tutorial, we'll see an example of using Composer's algorithms in a standalone fashion with no changes to the surrounding code and no requirement to use the Composer trainer. We'll be training a simple model on CIFAR-10, similar to the [PyTorch classifier tutorial](https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html). Because we'll be using a toy model trained for only a few epochs, we won't get the same speed or accuracy gains we might expect from a more realistic problem. However, this notebook should still serve as a useful illustration of how to use various algorithms. For examples of more realistic results, see the MosaicML [Explorer](https://app.mosaicml.com/explorer/imagenet)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Composer\n",
    "\n",
    "If you don't already have composer installed, install it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install mosaicml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the Model, Dataloader, and Training Loop\n",
    "\n",
    "First, we need to define our original model, dataloader, and training loop. Let's start with the dataloader:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.utils.data\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "datadir = './data'\n",
    "batch_size = 1024\n",
    "\n",
    "transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "    ]\n",
    ")\n",
    "\n",
    "trainset = torchvision.datasets.CIFAR10(root=datadir, train=True,\n",
    "                                        download=True, transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,\n",
    "                                          shuffle=True, num_workers=2)\n",
    "\n",
    "testset = torchvision.datasets.CIFAR10(root=datadir, train=False,\n",
    "                                       download=True, transform=transform)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,\n",
    "                                         shuffle=False, num_workers=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, we compose two transforms, one which transforms the images to tensors and another that normalizes them. We apply these transformations to both the train and test sets. Now, let's define our model. We're going to use a toy convolutional neural network so that the training epochs finish quickly."
   ]
  },
  {
   "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, 16, kernel_size=(3, 3), stride=2)\n",
    "        self.conv2 = nn.Conv2d(16, 32, kernel_size=(3, 3))\n",
    "        self.norm = nn.BatchNorm2d(32)\n",
    "        self.pool = nn.AdaptiveAvgPool2d(1)\n",
    "        self.fc1 = nn.Linear(32, 128)\n",
    "        self.fc2 = nn.Linear(128, 64)\n",
    "        self.fc3 = nn.Linear(64, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.conv1(x))\n",
    "        x = self.conv2(x)\n",
    "        x = F.relu(self.norm(x))\n",
    "        x = torch.flatten(self.pool(x), 1)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's write a simple training loop that prints the accuracy on the test set at the end of each epoch. We'll just run a few epochs for brevity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm\n",
    "import composer.functional as cf\n",
    "\n",
    "num_epochs = 5\n",
    "\n",
    "def train_and_eval(model, train_loader, test_loader):\n",
    "    torch.manual_seed(42)\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    model = model.to(device)\n",
    "    opt = torch.optim.Adam(model.parameters())\n",
    "    for epoch in range(num_epochs):\n",
    "        print(f\"---- Beginning epoch {epoch} ----\")\n",
    "        model.train()\n",
    "        progress_bar = tqdm(train_loader)\n",
    "        for X, y in progress_bar:\n",
    "            X = X.to(device)\n",
    "            y = y.to(device)\n",
    "            y_hat = model(X)\n",
    "            loss = F.cross_entropy(y_hat, y)\n",
    "            progress_bar.set_postfix_str(f\"train loss: {loss.item():.4f}\")\n",
    "            loss.backward()\n",
    "            opt.step()\n",
    "            opt.zero_grad()\n",
    "        model.eval()\n",
    "        num_right = 0\n",
    "        eval_size = 0\n",
    "        for X, y in test_loader:\n",
    "            X = X.to(device)\n",
    "            y = y.to(device)\n",
    "            y_hat = model(X)\n",
    "            num_right += (y_hat.argmax(dim=1) == y).sum().item()\n",
    "            eval_size += len(y)\n",
    "        acc_percent = 100 * num_right / eval_size\n",
    "        print(f\"Epoch {epoch} validation accuracy: {acc_percent:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great. Now, let's instantiate this baseline model and see how it fares on our dataset.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Net()\n",
    "train_and_eval(model, trainloader, testloader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have this baseline, let's add algorithms to improve our data pipeline and model. We'll start by adding some data augmentation, accessed via `cf.colout_batch`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create dataloaders for the train and test sets\n",
    "shared_transforms = [\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "]\n",
    "\n",
    "train_transforms = shared_transforms[:] + [cf.colout_batch]\n",
    "\n",
    "test_transform = transforms.Compose(shared_transforms)\n",
    "train_transform = transforms.Compose(train_transforms)\n",
    "\n",
    "trainset = torchvision.datasets.CIFAR10(root=datadir, train=True,\n",
    "                                        download=True, transform=train_transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,\n",
    "                                        shuffle=True, num_workers=2)\n",
    "\n",
    "testset = torchvision.datasets.CIFAR10(root=datadir, train=False,\n",
    "                                        download=True, transform=test_transform)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,\n",
    "                                          shuffle=False, num_workers=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how our model does with just these changes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Net()\n",
    "# only use one data augmentation since our small model runs quickly\n",
    "# and allows the dataloader little time to do anything fancy\n",
    "train_and_eval(model, trainloader, testloader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we might expect, adding data augmentation doesn't help us when we aren't training long enough to start overfitting.\n",
    "\n",
    "Let's try using some algorithms that modify the model. We're going to keep things simple and just add a [Squeeze-and-Excitation](https://docs.mosaicml.com/en/latest/method_cards/squeeze_excite.html) module after the larger of the two conv2d operations in our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# squeeze-excite can add a lot of overhead for small\n",
    "# conv2d operations, so only add it after convs with a\n",
    "# minimum number of channels\n",
    "cf.apply_squeeze_excite(model, latent_channels=64, min_channels=16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's see how our model does with the above algorithm applied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_and_eval(model, trainloader, testloader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adding squeeze-excite gives us another few percentage points of accuracy and does so with little decrease in the number of iterations per second. Great!\n",
    "\n",
    "Of course, this is a toy model and dataset, but it serves to illustrate how to use Composer's algorithms inside your own training loops, with minimal changes to your code. If you hit any problems or have questions, feel free to [open an issue](https://github.com/mosaicml/composer/issues/new/) or reach out to us [on Slack](https://join.slack.com/t/mosaicml-community/shared_invite/zt-w0tiddn9-WGTlRpfjcO9J5jyrMub1dg)."
   ]
  }
 ],
 "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
}