ResNet_Example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eO8VEUM63_Wf"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision import datasets, transforms, models\n",
    "from collections import OrderedDict\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, train_loader, loss_fn, optimizer, device):\n",
    "    # Prepare for training\n",
    "    model.train()\n",
    "    running_loss = 0\n",
    " \n",
    "    # Use TQDM for interactive loading bars\n",
    "    with tqdm(total=len(train_loader)) as pbar:\n",
    "        for i, (inputs, labels) in enumerate(train_loader, 0):\n",
    "            # Make image 3 channels and put on device\n",
    "            inputs, labels = inputs.to(device), labels.to(device)\n",
    "            inputs = torch.cat([inputs, inputs, inputs], axis=1)\n",
    "\n",
    "            # Run through model and update\n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(inputs)\n",
    "            loss = loss_fn(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            # Track loss and update progress\n",
    "            running_loss += loss.item()\n",
    "            pbar.update(1)\n",
    "\n",
    "    return running_loss / len(train_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function for the validation pass\n",
    "def validation(model, val_loader, loss_fn, device):\n",
    "    # Prepare for validating\n",
    "    model.eval()\n",
    "    val_loss = 0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        with tqdm(total=len(val_loader)) as pbar:\n",
    "            for inputs, labels in iter(val_loader):\n",
    "                # Make image 3 channels and put on device\n",
    "                inputs, labels = inputs.to(device), labels.to(device)\n",
    "                inputs = torch.cat([inputs, inputs, inputs], axis=1)\n",
    "\n",
    "                # Run through model\n",
    "                outputs = model(inputs)\n",
    "\n",
    "                # Track loss and update progress\n",
    "                val_loss += loss_fn(outputs, labels).item()\n",
    "\n",
    "                # Update accuracy\n",
    "                _, predicted = torch.max(outputs, 1)\n",
    "                total += labels.size(0)\n",
    "                correct += (predicted == labels).sum().item()\n",
    "\n",
    "                pbar.update(1)\n",
    "            \n",
    "    return val_loss / len(val_loader), correct / total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# An impure function to train our model\n",
    "def fit(epochs):\n",
    "    train_losses = []\n",
    "    test_losses = []\n",
    "    accuracies = []\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        # Run training loop and validation loop\n",
    "        train_loss = train(model, train_loader, criterion, optimizer, device)\n",
    "        val_loss, accuracy = validation(model, test_loader, criterion, device)\n",
    "\n",
    "        # Print result\n",
    "        print(\"Epoch: {}/{}, Training Loss: {:.4f}, Test Loss: {:.4f}, Test Accuracy: {}\".format(epoch + 1, epochs, train_loss, val_loss, accuracy))\n",
    "        print('-' * 20)\n",
    "\n",
    "        # Record results\n",
    "        train_losses.append(train_loss)\n",
    "        test_losses.append(val_loss)\n",
    "        accuracies.append(accuracy)    \n",
    "\n",
    "    print(\"Finished Training\")\n",
    "    return train_losses, test_losses, accuracies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_metrics(train_losses, test_losses, accuracies):\n",
    "    # Creating one figure with two subplots\n",
    "    f, (ax1, ax2) = plt.subplots(1, 2, sharey=False, figsize=(10, 3))\n",
    "    \n",
    "    # Plot accuracies\n",
    "    ax1.plot(acurracies)\n",
    "    ax1.set_xlabel('Epochs')\n",
    "    ax1.set_ylabel('Accuracy')\n",
    "    ax1.set_title('Test Accuracy')\n",
    "    \n",
    "    # Plot NLL losses\n",
    "    ax2.plot(train_losses, label='Train Losses')\n",
    "    ax2.plot(test_losses, label='Test Losses')\n",
    "    ax2.set_xlabel('Epochs')\n",
    "    ax2.set_ylabel('loss')\n",
    "    ax2.set_title('NLL Loss')\n",
    "    \n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Defining our Models and Datasets "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zESaTXBe3zDp"
   },
   "outputs": [],
   "source": [
    "# Define a transform to normalize the data\n",
    "transform = transforms.Compose([transforms.ToTensor(),\n",
    "                                transforms.Normalize((0.5,), (0.5,))])\n",
    "# Download and load the training data\n",
    "trainset = datasets.FashionMNIST('~/.pytorch/F_MNIST_data/', download=True, train=True, transform=transform)\n",
    "train_loader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)\n",
    "\n",
    "# Download and load the test data\n",
    "test_set = datasets.FashionMNIST('~/.pytorch/F_MNIST_data/', download=True, train=False, transform=transform)\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=64, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get model and pretrained weights\n",
    "model = models.resnet34(pretrained=True)\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initial Training\n",
    "We use `pretrained=True` so that we can make use of some weights from ImageNet pretraining that PyTorch makes available to us. We are going to freeze these weights and replace the output layers to retrain this network for the new problem. The frozen layers can be thought of as an 'image feature extractor', which we are using rather than starting with random weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_wKnU6jt4CZH"
   },
   "outputs": [],
   "source": [
    "# Freeze all layers\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = False "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Ujf_w4po4DrX",
    "outputId": "35cfc744-91c3-4571-e429-e1212794b4d5"
   },
   "outputs": [],
   "source": [
    "# Create output layers and replace output layers\n",
    "fc = nn.Sequential(OrderedDict([\n",
    "    ('fc1', nn.Linear(512,256)),\n",
    "    ('relu', nn.ReLU()),\n",
    "    ('fc2', nn.Linear(256,64)),\n",
    "    ('output', nn.LogSoftmax(dim=1))\n",
    "]))\n",
    "model.fc = fc\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set criterion and optimizer\n",
    "criterion = nn.NLLLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=3e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Xpnyw84gvNGf"
   },
   "outputs": [],
   "source": [
    "# Get device and put model on device\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "if 'EPOCHS' in os.environ:\n",
    "    epochs = int(os.environ['EPOCHS'])\n",
    "else:\n",
    "    epochs = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_losses, test_losses, acurracies = fit(epochs)\n",
    "plot_metrics(train_losses, test_losses, acurracies)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine Tuning\n",
    "We've now retrained this model for our problem, but we can do even better. The convolutional layers are still optimized for extracting features from the kind of images that imagenet has. After unfreezing the model we can retrain the whole network to adapt better to this problem, but we need to be careful to use a slower learning rate or we'll lose the pretraining that we already have."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PCXd1uWq4MQi"
   },
   "outputs": [],
   "source": [
    "# Unfreeze all layers\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = True "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reset our learning rate to make it slower\n",
    "optimizer = optim.Adam(model.parameters(), lr=3e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fine tune the layers at a slower learning rate\n",
    "train_losses, test_losses, acurracies = fit(epochs)\n",
    "plot_metrics(train_losses, test_losses, acurracies)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "ResNet_Example.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}