alexnet.ipynb

{
 "metadata": {
  "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.8.5-final"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python38564bitresearchconda6e46ab89dee4447294274d22b6e2b99b",
   "display_name": "Python 3.8.5 64-bit ('research': conda)"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "source": [
    "# AlexNet PyTorch Implementation\n"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "source": [
    "First we import the necessary libraries that we will use."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "\n",
    "import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "source": [
    "## AlexNet Implementation in PyTorch\n",
    "\n",
    "\n",
    "We can define these layers in a class called `AlexNet` as shown below.\n"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AlexNet(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    The AlexNet module.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, num_classes=1000):\n",
    "\n",
    "        # Mandatory call to super class module.\n",
    "        super(AlexNet, self).__init__()\n",
    "\n",
    "        # Defining the feature extraction layers.\n",
    "        self.feature_extractor = torch.nn.Sequential(\n",
    "\n",
    "            # Layer 1 - Convolution Layer - Nx3x224x224 -> Nx96x55x55\n",
    "            torch.nn.Conv2d(in_channels=3, out_channels=96,\n",
    "                            kernel_size=11, stride=4, padding=2),\n",
    "            torch.nn.ReLU(inplace=True),\n",
    "\n",
    "            # Layer 2 - Max Pooling Layer - Nx96x55x55 -> Nx96x27x27\n",
    "            torch.nn.MaxPool2d(kernel_size=3, stride=2),\n",
    "\n",
    "            # Layer 3 - Convolution Layer - Nx96x27x27 -> Nx256x27x27\n",
    "            torch.nn.Conv2d(in_channels=96, out_channels=256,\n",
    "                            kernel_size=5, padding=2),\n",
    "            torch.nn.ReLU(inplace=True),\n",
    "\n",
    "            # Layer 4 - Max Pooling Layer - Nx256x27x27 -> Nx256x13x13\n",
    "            torch.nn.MaxPool2d(kernel_size=3, stride=2),\n",
    "\n",
    "            # Layer 5 - Convolution Layer - Nx256x13x13 -> Nx384x13x13\n",
    "            torch.nn.Conv2d(in_channels=256, out_channels=384,\n",
    "                            kernel_size=3, padding=1),\n",
    "            torch.nn.ReLU(inplace=True),\n",
    "\n",
    "            # Layer 6 - Convolution Layer - Nx384x13x13 -> Nx384x13x13\n",
    "            torch.nn.Conv2d(in_channels=384, out_channels=384,\n",
    "                            kernel_size=3, padding=1),\n",
    "            torch.nn.ReLU(inplace=True),\n",
    "\n",
    "            # Layer 7 - Convolution Layer - Nx384x13x13 -> Nx256x13x13\n",
    "            torch.nn.Conv2d(in_channels=384, out_channels=256,\n",
    "                            kernel_size=3, padding=1),\n",
    "            torch.nn.ReLU(inplace=True),\n",
    "\n",
    "            # Layer 8 - Max Pooling Layer - Nx256x13x13 -> Nx256x6x6\n",
    "            torch.nn.MaxPool2d(kernel_size=3, stride=2),\n",
    "        )\n",
    "\n",
    "        # Defining the classification layers.\n",
    "        self.classifier = torch.nn.Sequential(\n",
    "\n",
    "            # Layer 9 - Fully Connected - Nx1x9216 -> Nx1x4096\n",
    "            torch.nn.Linear(in_features=256*6*6, out_features=4096),\n",
    "            torch.nn.ReLU(inplace=True),\n",
    "\n",
    "            # Layer 9 - Fully Connected - Nx256x6x6 -> Nx1x10\n",
    "            torch.nn.Linear(in_features=4096, out_features=4096),\n",
    "            torch.nn.ReLU(inplace=True),\n",
    "\n",
    "            # Layer 9 - Fully Connected - Nx256x6x6 -> Nx1xnum_classes\n",
    "            torch.nn.Linear(in_features=4096, out_features=num_classes),\n",
    "            torch.nn.Softmax(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        # Forward pass through the feature extractor - Nx3x224x224 -> Nx256x6x6\n",
    "        x = self.feature_extractor(x)\n",
    "\n",
    "        # Flattening the feature map - Nx256x6x6 -> Nx1x9216\n",
    "        x = torch.flatten(x, 1)\n",
    "\n",
    "        # Forward pass through the classifier - Nx1x9216 -> Nx1xnum_classes\n",
    "        return self.classifier(x)\n"
   ]
  },
  {
   "source": [
    "Note that a `forward` function was also defined. This function dictates how to forward-propagate the input through the network. "
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "source": [
    "## Loading the CIFAR-10 dataset\n",
    "\n"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Defining a transform for the images.\n",
    "transform = torchvision.transforms.Compose(\n",
    "    [torchvision.transforms.Resize((224,224)), torchvision.transforms.ToTensor(), \n",
    "    torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]\n",
    ")\n",
    "\n",
    "# Loading the training and validation data.\n",
    "train_set = torchvision.datasets.CIFAR10(root='./cifar-10', train=True, download=True, transform=transform)\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=4, shuffle=True, num_workers=2)\n",
    "#val_loader = torch.utils.data.DataLoader(train_set, batch_size=4, shuffle=True, num_workers=2)\n",
    "\n",
    "# Loading the testing data.\n",
    "test_set = torchvision.datasets.CIFAR10(root='./cifar-10', train=False, download=True, transform=transform)\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=4, shuffle=False, num_workers=2)\n",
    "val_loader = torch.utils.data.DataLoader(test_set, batch_size=4, shuffle=False, num_workers=2)\n",
    "\n",
    "# Defining the classes.\n",
    "classes = ('plane', 'car', 'bird', 'cat',\n",
    "           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')"
   ]
  },
  {
   "source": [
    "### Showing Sample Images\n",
    "\n",
    "To confirm that the data was loaded correctly, we design a function below to show some sample images from the dataset."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def show_image(image):\n",
    "    image = image / 2 + 0.5\n",
    "    npimg = image.numpy()\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
    "    plt.show()\n",
    "\n",
    "def show_sample_images():\n",
    "\n",
    "    # get some random training images\n",
    "    dataiter = iter(train_loader)\n",
    "    images, labels = dataiter.next()\n",
    "\n",
    "    # Showing the image(s).\n",
    "    show_image(torchvision.utils.make_grid(images))\n",
    "\n",
    "    # Printing the labels.\n",
    "    print(' '.join('%5s' % classes[labels[j]] for j in range(4)))\n",
    "\n",
    "show_sample_images()"
   ]
  },
  {
   "source": [
    "## Using the model\n",
    "\n",
    "Before we proceed, we check to see your the machine has a GPU installed. If so, we use the GPU for training, else, we use the CPU."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model = AlexNet(10).to(device)\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)"
   ]
  },
  {
   "source": [
    "## Training the model\n",
    "\n",
    "We devise a function below to train the model. We use regular Stochastic Gradient Descent and Mean Squared Error loss for training as defined by LeCun."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train(model, train_loader, val_loader, optimizer, epochs):\n",
    "\n",
    "    # Iterate for several epochs\n",
    "    for epoch in tqdm.trange(epochs):\n",
    "        train_loss = val_loss = 0.0\n",
    "\n",
    "        for phase in ['train', 'val']:\n",
    "            \n",
    "            if phase == 'train':\n",
    "                loader = train_loader\n",
    "                model.train(True)\n",
    "            else:\n",
    "                loader = val_loader\n",
    "                model.train(False)\n",
    "\n",
    "            # Iterate for each data item in the training set\n",
    "            for i, data in enumerate(loader, 0):\n",
    "                \n",
    "                # Get the sample input data.\n",
    "                inputs, labels = data[0].to(device), data[1].to(device)\n",
    "\n",
    "                # Reset the gradients\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                # Perform forward pass.\n",
    "                outputs = model(inputs)\n",
    "\n",
    "                # Calculate current model loss.\n",
    "                loss = criterion(outputs, labels)\n",
    "\n",
    "                if phase == 'train':\n",
    "                     \n",
    "                     # Perform backward pass.\n",
    "                    loss.backward()\n",
    "                    optimizer.step()\n",
    "\n",
    "                    train_loss += loss.item()\n",
    "                else:\n",
    "                    val_loss += loss.item()\n",
    "        \n",
    "        print('Epoch %d - Train Loss: %.3f Validation Loss %.3f' % (epoch + 1, train_loss/len(train_loader), val_loss/len(val_loader)))\n",
    "\n",
    "    print('Finished Training')\n",
    "\n",
    "train(model, train_loader, val_loader, optimizer, 20)"
   ]
  },
  {
   "source": [
    "## Evaluating the Model\n",
    "\n",
    "We evaluate the model on a ransom sample of test data to see how well it performs."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(model, test_loader):\n",
    "\n",
    "    dataiter = iter(test_loader)\n",
    "    images, labels = dataiter.next()\n",
    "\n",
    "    # Showing the test images\n",
    "    show_image(torchvision.utils.make_grid(images))\n",
    "    print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))\n",
    "\n",
    "    outputs = model(images.to(device))\n",
    "    _, predicted = torch.max(outputs, 1)\n",
    "\n",
    "    print('Predicted: ', ' '.join('%5s' % classes[predicted[j]] for j in range(4)))\n",
    "\n",
    "test(model, test_loader)"
   ]
  },
  {
   "source": [
    "### Evaluating on the entire dataset\n",
    "\n",
    "We further evaluate the model's performance on the entire dataset."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_full(model, test_loader):\n",
    "\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for data in test_loader:\n",
    "            images, labels = data[0].to(device), data[1].to(device)\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "\n",
    "    print('Accuracy of the network on the 10000 test images: %d %%' % (\n",
    "        100 * correct / total))\n",
    "\n",
    "test_full(model, test_loader)"
   ]
  },
  {
   "source": [
    "### Evaluating class by class\n",
    "\n",
    "We further evaluate the class by class accuracy of the model"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_class(model, test_loader):\n",
    "\n",
    "    class_correct = list(0. for i in range(10))\n",
    "    class_total = list(0. for i in range(10))\n",
    "    with torch.no_grad():\n",
    "        for data in test_loader:\n",
    "            images, labels = data[0].to(device), data[1].to(device)\n",
    "            outputs = model(images.to(device))\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            c = (predicted == labels).squeeze()\n",
    "            for i in range(4):\n",
    "                label = labels[i]\n",
    "                class_correct[label] += c[i].item()\n",
    "                class_total[label] += 1\n",
    "\n",
    "\n",
    "    for i in range(10):\n",
    "        print('Accuracy of %5s : %2d %%' % (\n",
    "            classes[i], 100 * class_correct[i] / class_total[i]))\n",
    "\n",
    "test_class(model, test_loader)"
   ]
  }
 ]
}