notebooks/transfer_learning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transfer learning with PyTorch\n",
    "We're going to train a neural network to classify dogs and cats."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Init, helpers, utils, ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "\n",
    "DEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "DEVICE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from IPython.core.debugger import set_trace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load my_train_helper.py\n",
    "def get_trainable(model_params):\n",
    "    return (p for p in model_params if p.requires_grad)\n",
    "\n",
    "\n",
    "def get_frozen(model_params):\n",
    "    return (p for p in model_params if not p.requires_grad)\n",
    "\n",
    "\n",
    "def all_trainable(model_params):\n",
    "    return all(p.requires_grad for p in model_params)\n",
    "\n",
    "\n",
    "def all_frozen(model_params):\n",
    "    return all(not p.requires_grad for p in model_params)\n",
    "\n",
    "\n",
    "def freeze_all(model_params):\n",
    "    for param in model_params:\n",
    "        param.requires_grad = False\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "toc-hr-collapsed": true
   },
   "source": [
    "# The Data - DogsCatsDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import transforms\n",
    "\n",
    "_image_size = 224\n",
    "_mean = [0.485, 0.456, 0.406]\n",
    "_std = [0.229, 0.224, 0.225]\n",
    "\n",
    "\n",
    "train_trans = transforms.Compose([\n",
    "    transforms.Resize(256),  # some images are pretty small\n",
    "    transforms.RandomCrop(_image_size),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.ColorJitter(.3, .3, .3),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(_mean, _std),\n",
    "])\n",
    "val_trans = transforms.Compose([\n",
    "    transforms.Resize(256),\n",
    "    transforms.CenterCrop(_image_size),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(_mean, _std),\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "\n",
    "The implementation of the dataset does not really."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.datasets.folder import ImageFolder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = ImageFolder(\"dogscats/training_set/\", transform=train_trans)\n",
    "val_ds = ImageFolder(\"dogscats/test_set/\", transform=val_trans)\n",
    "\n",
    "batch_size = 32\n",
    "n_classes = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the following if you want to use the full dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_ds = DogsCatsDataset(\"../data/raw\", \"train\", transform=train_trans)\n",
    "# val_ds = DogsCatsDataset(\"../data/raw\", \"valid\", transform=val_trans)\n",
    "\n",
    "# batch_size = 128\n",
    "# n_classes = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(8000, 2000)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_ds), len(val_ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DataLoader\n",
    "Batch loading for datasets with multi-processing and different sample strategies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "\n",
    "train_dl = DataLoader(\n",
    "    train_ds,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    num_workers=16,\n",
    ")\n",
    "\n",
    "val_dl = DataLoader(\n",
    "    val_ds,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False,\n",
    "    num_workers=16,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Model\n",
    "PyTorch offers quite a few [pre-trained networks](https://pytorch.org/docs/stable/torchvision/models.html) such as:\n",
    "- AlexNet\n",
    "- VGG\n",
    "- ResNet\n",
    "- SqueezeNet\n",
    "- DenseNet\n",
    "- Inception v3\n",
    "\n",
    "And there are more available via [pretrained-models.pytorch](https://github.com/Cadene/pretrained-models.pytorch):\n",
    "- NASNet,\n",
    "- ResNeXt,\n",
    "- InceptionV4,\n",
    "- InceptionResnetV2, \n",
    "- Xception, \n",
    "- DPN,\n",
    "- ...\n",
    "\n",
    "We'll use a simple resnet18 model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import models\n",
    "\n",
    "model = models.resnet18(pretrained=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResNet(\n",
       "  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (relu): ReLU(inplace=True)\n",
       "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "  (layer1): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (layer2): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (layer3): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (layer4): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "  (fc): Linear(in_features=512, out_features=1000, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 112, 112]           9,408\n",
      "       BatchNorm2d-2         [-1, 64, 112, 112]             128\n",
      "              ReLU-3         [-1, 64, 112, 112]               0\n",
      "         MaxPool2d-4           [-1, 64, 56, 56]               0\n",
      "            Conv2d-5           [-1, 64, 56, 56]          36,864\n",
      "       BatchNorm2d-6           [-1, 64, 56, 56]             128\n",
      "              ReLU-7           [-1, 64, 56, 56]               0\n",
      "            Conv2d-8           [-1, 64, 56, 56]          36,864\n",
      "       BatchNorm2d-9           [-1, 64, 56, 56]             128\n",
      "             ReLU-10           [-1, 64, 56, 56]               0\n",
      "       BasicBlock-11           [-1, 64, 56, 56]               0\n",
      "           Conv2d-12           [-1, 64, 56, 56]          36,864\n",
      "      BatchNorm2d-13           [-1, 64, 56, 56]             128\n",
      "             ReLU-14           [-1, 64, 56, 56]               0\n",
      "           Conv2d-15           [-1, 64, 56, 56]          36,864\n",
      "      BatchNorm2d-16           [-1, 64, 56, 56]             128\n",
      "             ReLU-17           [-1, 64, 56, 56]               0\n",
      "       BasicBlock-18           [-1, 64, 56, 56]               0\n",
      "           Conv2d-19          [-1, 128, 28, 28]          73,728\n",
      "      BatchNorm2d-20          [-1, 128, 28, 28]             256\n",
      "             ReLU-21          [-1, 128, 28, 28]               0\n",
      "           Conv2d-22          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-23          [-1, 128, 28, 28]             256\n",
      "           Conv2d-24          [-1, 128, 28, 28]           8,192\n",
      "      BatchNorm2d-25          [-1, 128, 28, 28]             256\n",
      "             ReLU-26          [-1, 128, 28, 28]               0\n",
      "       BasicBlock-27          [-1, 128, 28, 28]               0\n",
      "           Conv2d-28          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-29          [-1, 128, 28, 28]             256\n",
      "             ReLU-30          [-1, 128, 28, 28]               0\n",
      "           Conv2d-31          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-32          [-1, 128, 28, 28]             256\n",
      "             ReLU-33          [-1, 128, 28, 28]               0\n",
      "       BasicBlock-34          [-1, 128, 28, 28]               0\n",
      "           Conv2d-35          [-1, 256, 14, 14]         294,912\n",
      "      BatchNorm2d-36          [-1, 256, 14, 14]             512\n",
      "             ReLU-37          [-1, 256, 14, 14]               0\n",
      "           Conv2d-38          [-1, 256, 14, 14]         589,824\n",
      "      BatchNorm2d-39          [-1, 256, 14, 14]             512\n",
      "           Conv2d-40          [-1, 256, 14, 14]          32,768\n",
      "      BatchNorm2d-41          [-1, 256, 14, 14]             512\n",
      "             ReLU-42          [-1, 256, 14, 14]               0\n",
      "       BasicBlock-43          [-1, 256, 14, 14]               0\n",
      "           Conv2d-44          [-1, 256, 14, 14]         589,824\n",
      "      BatchNorm2d-45          [-1, 256, 14, 14]             512\n",
      "             ReLU-46          [-1, 256, 14, 14]               0\n",
      "           Conv2d-47          [-1, 256, 14, 14]         589,824\n",
      "      BatchNorm2d-48          [-1, 256, 14, 14]             512\n",
      "             ReLU-49          [-1, 256, 14, 14]               0\n",
      "       BasicBlock-50          [-1, 256, 14, 14]               0\n",
      "           Conv2d-51            [-1, 512, 7, 7]       1,179,648\n",
      "      BatchNorm2d-52            [-1, 512, 7, 7]           1,024\n",
      "             ReLU-53            [-1, 512, 7, 7]               0\n",
      "           Conv2d-54            [-1, 512, 7, 7]       2,359,296\n",
      "      BatchNorm2d-55            [-1, 512, 7, 7]           1,024\n",
      "           Conv2d-56            [-1, 512, 7, 7]         131,072\n",
      "      BatchNorm2d-57            [-1, 512, 7, 7]           1,024\n",
      "             ReLU-58            [-1, 512, 7, 7]               0\n",
      "       BasicBlock-59            [-1, 512, 7, 7]               0\n",
      "           Conv2d-60            [-1, 512, 7, 7]       2,359,296\n",
      "      BatchNorm2d-61            [-1, 512, 7, 7]           1,024\n",
      "             ReLU-62            [-1, 512, 7, 7]               0\n",
      "           Conv2d-63            [-1, 512, 7, 7]       2,359,296\n",
      "      BatchNorm2d-64            [-1, 512, 7, 7]           1,024\n",
      "             ReLU-65            [-1, 512, 7, 7]               0\n",
      "       BasicBlock-66            [-1, 512, 7, 7]               0\n",
      "AdaptiveAvgPool2d-67            [-1, 512, 1, 1]               0\n",
      "           Linear-68                 [-1, 1000]         513,000\n",
      "================================================================\n",
      "Total params: 11,689,512\n",
      "Trainable params: 11,689,512\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.57\n",
      "Forward/backward pass size (MB): 62.79\n",
      "Params size (MB): 44.59\n",
      "Estimated Total Size (MB): 107.96\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "import torchsummary\n",
    "\n",
    "torchsummary.summary(model, (3, 224, 224), device=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Linear(in_features=2, out_features=1, bias=True)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nn.Linear(2, 1, bias=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Freeze all parameters manually\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Or use our convenient functions from before\n",
    "freeze_all(model.parameters())\n",
    "assert all_frozen(model.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Replace the last layer with a linear layer. New layers have `requires_grad = True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fc = nn.Linear(512, n_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert not all_frozen(model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model(n_classes=2):\n",
    "    model = models.resnet18(pretrained=True)\n",
    "    freeze_all(model.parameters())\n",
    "    model.fc = nn.Linear(512, n_classes)\n",
    "    model = model.to(DEVICE)\n",
    "    return model\n",
    "\n",
    "\n",
    "model = get_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(\n",
    "    get_trainable(model.parameters()),\n",
    "    lr=0.001,\n",
    "    # momentum=0.9,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Train Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/1 |  batch: 0 |  batch loss:   0.775\n",
      "Epoch 1/1 |  batch: 1 |  batch loss:   0.611\n",
      "Epoch 1/1 |  batch: 2 |  batch loss:   0.690\n",
      "Epoch 1/1 |  batch: 3 |  batch loss:   0.715\n",
      "Epoch 1/1 |  batch: 4 |  batch loss:   0.686\n",
      "Epoch 1/1 |  batch: 5 |  batch loss:   0.463\n",
      "Epoch 1/1 |  batch: 6 |  batch loss:   0.467\n",
      "Epoch 1/1 |  batch: 7 |  batch loss:   0.484\n",
      "Epoch 1/1 |  batch: 8 |  batch loss:   0.384\n",
      "Epoch 1/1 |  batch: 9 |  batch loss:   0.395\n",
      "Epoch 1/1 |  batch: 10 |  batch loss:   0.316\n",
      "Epoch 1/1 |  batch: 11 |  batch loss:   0.403\n",
      "Epoch 1/1 |  batch: 12 |  batch loss:   0.330\n",
      "Epoch 1/1 |  batch: 13 |  batch loss:   0.420\n",
      "Epoch 1/1 |  batch: 14 |  batch loss:   0.312\n",
      "Epoch 1/1 |  batch: 15 |  batch loss:   0.369\n",
      "Epoch 1/1 |  batch: 16 |  batch loss:   0.264\n",
      "Epoch 1/1 |  batch: 17 |  batch loss:   0.298\n",
      "Epoch 1/1 |  batch: 18 |  batch loss:   0.299\n",
      "Epoch 1/1 |  batch: 19 |  batch loss:   0.391\n",
      "Epoch 1/1 |  batch: 20 |  batch loss:   0.271\n",
      "Epoch 1/1 |  batch: 21 |  batch loss:   0.327\n",
      "Epoch 1/1 |  batch: 22 |  batch loss:   0.217\n",
      "Epoch 1/1 |  batch: 23 |  batch loss:   0.168\n",
      "Epoch 1/1 |  batch: 24 |  batch loss:   0.229\n",
      "Epoch 1/1 |  batch: 25 |  batch loss:   0.204\n",
      "Epoch 1/1 |  batch: 26 |  batch loss:   0.286\n",
      "Epoch 1/1 |  batch: 27 |  batch loss:   0.189\n",
      "Epoch 1/1 |  batch: 28 |  batch loss:   0.274\n",
      "Epoch 1/1 |  batch: 29 |  batch loss:   0.206\n",
      "Epoch 1/1 |  batch: 30 |  batch loss:   0.142\n",
      "Epoch 1/1 |  batch: 31 |  batch loss:   0.169\n",
      "Epoch 1/1 |  batch: 32 |  batch loss:   0.228\n",
      "Epoch 1/1 |  batch: 33 |  batch loss:   0.207\n",
      "Epoch 1/1 |  batch: 34 |  batch loss:   0.352\n",
      "Epoch 1/1 |  batch: 35 |  batch loss:   0.239\n",
      "Epoch 1/1 |  batch: 36 |  batch loss:   0.207\n",
      "Epoch 1/1 |  batch: 37 |  batch loss:   0.154\n",
      "Epoch 1/1 |  batch: 38 |  batch loss:   0.225\n",
      "Epoch 1/1 |  batch: 39 |  batch loss:   0.221\n",
      "Epoch 1/1 |  batch: 40 |  batch loss:   0.227\n",
      "Epoch 1/1 |  batch: 41 |  batch loss:   0.159\n",
      "Epoch 1/1 |  batch: 42 |  batch loss:   0.125\n",
      "Epoch 1/1 |  batch: 43 |  batch loss:   0.099\n",
      "Epoch 1/1 |  batch: 44 |  batch loss:   0.152\n",
      "Epoch 1/1 |  batch: 45 |  batch loss:   0.141\n",
      "Epoch 1/1 |  batch: 46 |  batch loss:   0.174\n",
      "Epoch 1/1 |  batch: 47 |  batch loss:   0.155\n",
      "Epoch 1/1 |  batch: 48 |  batch loss:   0.243\n",
      "Epoch 1/1 |  batch: 49 |  batch loss:   0.183\n",
      "Epoch 1/1 |  batch: 50 |  batch loss:   0.106\n",
      "Epoch 1/1 |  batch: 51 |  batch loss:   0.144\n",
      "Epoch 1/1 |  batch: 52 |  batch loss:   0.148\n",
      "Epoch 1/1 |  batch: 53 |  batch loss:   0.276\n",
      "Epoch 1/1 |  batch: 54 |  batch loss:   0.168\n",
      "Epoch 1/1 |  batch: 55 |  batch loss:   0.131\n",
      "Epoch 1/1 |  batch: 56 |  batch loss:   0.140\n",
      "Epoch 1/1 |  batch: 57 |  batch loss:   0.255\n",
      "Epoch 1/1 |  batch: 58 |  batch loss:   0.145\n",
      "Epoch 1/1 |  batch: 59 |  batch loss:   0.139\n",
      "Epoch 1/1 |  batch: 60 |  batch loss:   0.146\n",
      "Epoch 1/1 |  batch: 61 |  batch loss:   0.311\n",
      "Epoch 1/1 |  batch: 62 |  batch loss:   0.117\n",
      "Epoch 1/1 |  batch: 63 |  batch loss:   0.155\n",
      "Epoch 1/1 |  batch: 64 |  batch loss:   0.205\n",
      "Epoch 1/1 |  batch: 65 |  batch loss:   0.116\n",
      "Epoch 1/1 |  batch: 66 |  batch loss:   0.127\n",
      "Epoch 1/1 |  batch: 67 |  batch loss:   0.108\n",
      "Epoch 1/1 |  batch: 68 |  batch loss:   0.097\n",
      "Epoch 1/1 |  batch: 69 |  batch loss:   0.165\n",
      "Epoch 1/1 |  batch: 70 |  batch loss:   0.196\n",
      "Epoch 1/1 |  batch: 71 |  batch loss:   0.121\n",
      "Epoch 1/1 |  batch: 72 |  batch loss:   0.166\n",
      "Epoch 1/1 |  batch: 73 |  batch loss:   0.086\n",
      "Epoch 1/1 |  batch: 74 |  batch loss:   0.146\n",
      "Epoch 1/1 |  batch: 75 |  batch loss:   0.160\n",
      "Epoch 1/1 |  batch: 76 |  batch loss:   0.263\n",
      "Epoch 1/1 |  batch: 77 |  batch loss:   0.094\n",
      "Epoch 1/1 |  batch: 78 |  batch loss:   0.133\n",
      "Epoch 1/1 |  batch: 79 |  batch loss:   0.076\n",
      "Epoch 1/1 |  batch: 80 |  batch loss:   0.217\n",
      "Epoch 1/1 |  batch: 81 |  batch loss:   0.099\n",
      "Epoch 1/1 |  batch: 82 |  batch loss:   0.166\n",
      "Epoch 1/1 |  batch: 83 |  batch loss:   0.080\n",
      "Epoch 1/1 |  batch: 84 |  batch loss:   0.099\n",
      "Epoch 1/1 |  batch: 85 |  batch loss:   0.110\n",
      "Epoch 1/1 |  batch: 86 |  batch loss:   0.126\n",
      "Epoch 1/1 |  batch: 87 |  batch loss:   0.097\n",
      "Epoch 1/1 |  batch: 88 |  batch loss:   0.243\n",
      "Epoch 1/1 |  batch: 89 |  batch loss:   0.205\n",
      "Epoch 1/1 |  batch: 90 |  batch loss:   0.132\n",
      "Epoch 1/1 |  batch: 91 |  batch loss:   0.119\n",
      "Epoch 1/1 |  batch: 92 |  batch loss:   0.163\n",
      "Epoch 1/1 |  batch: 93 |  batch loss:   0.135\n",
      "Epoch 1/1 |  batch: 94 |  batch loss:   0.062\n",
      "Epoch 1/1 |  batch: 95 |  batch loss:   0.144\n",
      "Epoch 1/1 |  batch: 96 |  batch loss:   0.165\n",
      "Epoch 1/1 |  batch: 97 |  batch loss:   0.072\n",
      "Epoch 1/1 |  batch: 98 |  batch loss:   0.126\n",
      "Epoch 1/1 |  batch: 99 |  batch loss:   0.089\n",
      "Epoch 1/1 |  batch: 100 |  batch loss:   0.081\n",
      "Epoch 1/1 |  batch: 101 |  batch loss:   0.244\n",
      "Epoch 1/1 |  batch: 102 |  batch loss:   0.125\n",
      "Epoch 1/1 |  batch: 103 |  batch loss:   0.096\n",
      "Epoch 1/1 |  batch: 104 |  batch loss:   0.175\n",
      "Epoch 1/1 |  batch: 105 |  batch loss:   0.092\n",
      "Epoch 1/1 |  batch: 106 |  batch loss:   0.155\n",
      "Epoch 1/1 |  batch: 107 |  batch loss:   0.194\n",
      "Epoch 1/1 |  batch: 108 |  batch loss:   0.250\n",
      "Epoch 1/1 |  batch: 109 |  batch loss:   0.083\n",
      "Epoch 1/1 |  batch: 110 |  batch loss:   0.135\n",
      "Epoch 1/1 |  batch: 111 |  batch loss:   0.626\n",
      "Epoch 1/1 |  batch: 112 |  batch loss:   0.191\n",
      "Epoch 1/1 |  batch: 113 |  batch loss:   0.224\n",
      "Epoch 1/1 |  batch: 114 |  batch loss:   0.074\n",
      "Epoch 1/1 |  batch: 115 |  batch loss:   0.105\n",
      "Epoch 1/1 |  batch: 116 |  batch loss:   0.120\n",
      "Epoch 1/1 |  batch: 117 |  batch loss:   0.185\n",
      "Epoch 1/1 |  batch: 118 |  batch loss:   0.417\n",
      "Epoch 1/1 |  batch: 119 |  batch loss:   0.291\n",
      "Epoch 1/1 |  batch: 120 |  batch loss:   0.088\n",
      "Epoch 1/1 |  batch: 121 |  batch loss:   0.111\n",
      "Epoch 1/1 |  batch: 122 |  batch loss:   0.070\n",
      "Epoch 1/1 |  batch: 123 |  batch loss:   0.177\n",
      "Epoch 1/1 |  batch: 124 |  batch loss:   0.153\n",
      "Epoch 1/1 |  batch: 125 |  batch loss:   0.252\n",
      "Epoch 1/1 |  batch: 126 |  batch loss:   0.106\n",
      "Epoch 1/1 |  batch: 127 |  batch loss:   0.155\n",
      "Epoch 1/1 |  batch: 128 |  batch loss:   0.157\n",
      "Epoch 1/1 |  batch: 129 |  batch loss:   0.059\n",
      "Epoch 1/1 |  batch: 130 |  batch loss:   0.087\n",
      "Epoch 1/1 |  batch: 131 |  batch loss:   0.185\n",
      "Epoch 1/1 |  batch: 132 |  batch loss:   0.088\n",
      "Epoch 1/1 |  batch: 133 |  batch loss:   0.145\n",
      "Epoch 1/1 |  batch: 134 |  batch loss:   0.355\n",
      "Epoch 1/1 |  batch: 135 |  batch loss:   0.129\n",
      "Epoch 1/1 |  batch: 136 |  batch loss:   0.103\n",
      "Epoch 1/1 |  batch: 137 |  batch loss:   0.122\n",
      "Epoch 1/1 |  batch: 138 |  batch loss:   0.081\n",
      "Epoch 1/1 |  batch: 139 |  batch loss:   0.185\n",
      "Epoch 1/1 |  batch: 140 |  batch loss:   0.028\n",
      "Epoch 1/1 |  batch: 141 |  batch loss:   0.100\n",
      "Epoch 1/1 |  batch: 142 |  batch loss:   0.154\n",
      "Epoch 1/1 |  batch: 143 |  batch loss:   0.077\n",
      "Epoch 1/1 |  batch: 144 |  batch loss:   0.124\n",
      "Epoch 1/1 |  batch: 145 |  batch loss:   0.070\n",
      "Epoch 1/1 |  batch: 146 |  batch loss:   0.098\n",
      "Epoch 1/1 |  batch: 147 |  batch loss:   0.113\n",
      "Epoch 1/1 |  batch: 148 |  batch loss:   0.096\n",
      "Epoch 1/1 |  batch: 149 |  batch loss:   0.078\n",
      "Epoch 1/1 |  batch: 150 |  batch loss:   0.090\n",
      "Epoch 1/1 |  batch: 151 |  batch loss:   0.109\n",
      "Epoch 1/1 |  batch: 152 |  batch loss:   0.121\n",
      "Epoch 1/1 |  batch: 153 |  batch loss:   0.082\n",
      "Epoch 1/1 |  batch: 154 |  batch loss:   0.086\n",
      "Epoch 1/1 |  batch: 155 |  batch loss:   0.150\n",
      "Epoch 1/1 |  batch: 156 |  batch loss:   0.157\n",
      "Epoch 1/1 |  batch: 157 |  batch loss:   0.168\n",
      "Epoch 1/1 |  batch: 158 |  batch loss:   0.126\n",
      "Epoch 1/1 |  batch: 159 |  batch loss:   0.106\n",
      "Epoch 1/1 |  batch: 160 |  batch loss:   0.241\n",
      "Epoch 1/1 |  batch: 161 |  batch loss:   0.233\n",
      "Epoch 1/1 |  batch: 162 |  batch loss:   0.047\n",
      "Epoch 1/1 |  batch: 163 |  batch loss:   0.139\n",
      "Epoch 1/1 |  batch: 164 |  batch loss:   0.139\n",
      "Epoch 1/1 |  batch: 165 |  batch loss:   0.065\n",
      "Epoch 1/1 |  batch: 166 |  batch loss:   0.190\n",
      "Epoch 1/1 |  batch: 167 |  batch loss:   0.055\n",
      "Epoch 1/1 |  batch: 168 |  batch loss:   0.046\n",
      "Epoch 1/1 |  batch: 169 |  batch loss:   0.071\n",
      "Epoch 1/1 |  batch: 170 |  batch loss:   0.108\n",
      "Epoch 1/1 |  batch: 171 |  batch loss:   0.125\n",
      "Epoch 1/1 |  batch: 172 |  batch loss:   0.104\n",
      "Epoch 1/1 |  batch: 173 |  batch loss:   0.087\n",
      "Epoch 1/1 |  batch: 174 |  batch loss:   0.060\n",
      "Epoch 1/1 |  batch: 175 |  batch loss:   0.111\n",
      "Epoch 1/1 |  batch: 176 |  batch loss:   0.107\n",
      "Epoch 1/1 |  batch: 177 |  batch loss:   0.084\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/1 |  batch: 178 |  batch loss:   0.152\n",
      "Epoch 1/1 |  batch: 179 |  batch loss:   0.044\n",
      "Epoch 1/1 |  batch: 180 |  batch loss:   0.098\n",
      "Epoch 1/1 |  batch: 181 |  batch loss:   0.102\n",
      "Epoch 1/1 |  batch: 182 |  batch loss:   0.131\n",
      "Epoch 1/1 |  batch: 183 |  batch loss:   0.101\n",
      "Epoch 1/1 |  batch: 184 |  batch loss:   0.106\n",
      "Epoch 1/1 |  batch: 185 |  batch loss:   0.069\n",
      "Epoch 1/1 |  batch: 186 |  batch loss:   0.082\n",
      "Epoch 1/1 |  batch: 187 |  batch loss:   0.041\n",
      "Epoch 1/1 |  batch: 188 |  batch loss:   0.104\n",
      "Epoch 1/1 |  batch: 189 |  batch loss:   0.165\n",
      "Epoch 1/1 |  batch: 190 |  batch loss:   0.092\n",
      "Epoch 1/1 |  batch: 191 |  batch loss:   0.159\n",
      "Epoch 1/1 |  batch: 192 |  batch loss:   0.186\n",
      "Epoch 1/1 |  batch: 193 |  batch loss:   0.141\n",
      "Epoch 1/1 |  batch: 194 |  batch loss:   0.153\n",
      "Epoch 1/1 |  batch: 195 |  batch loss:   0.154\n",
      "Epoch 1/1 |  batch: 196 |  batch loss:   0.067\n",
      "Epoch 1/1 |  batch: 197 |  batch loss:   0.209\n",
      "Epoch 1/1 |  batch: 198 |  batch loss:   0.245\n",
      "Epoch 1/1 |  batch: 199 |  batch loss:   0.103\n",
      "Epoch 1/1 |  batch: 200 |  batch loss:   0.146\n",
      "Epoch 1/1 |  batch: 201 |  batch loss:   0.212\n",
      "Epoch 1/1 |  batch: 202 |  batch loss:   0.109\n",
      "Epoch 1/1 |  batch: 203 |  batch loss:   0.188\n",
      "Epoch 1/1 |  batch: 204 |  batch loss:   0.274\n",
      "Epoch 1/1 |  batch: 205 |  batch loss:   0.060\n",
      "Epoch 1/1 |  batch: 206 |  batch loss:   0.170\n",
      "Epoch 1/1 |  batch: 207 |  batch loss:   0.306\n",
      "Epoch 1/1 |  batch: 208 |  batch loss:   0.109\n",
      "Epoch 1/1 |  batch: 209 |  batch loss:   0.254\n",
      "Epoch 1/1 |  batch: 210 |  batch loss:   0.171\n",
      "Epoch 1/1 |  batch: 211 |  batch loss:   0.087\n",
      "Epoch 1/1 |  batch: 212 |  batch loss:   0.287\n",
      "Epoch 1/1 |  batch: 213 |  batch loss:   0.156\n",
      "Epoch 1/1 |  batch: 214 |  batch loss:   0.046\n",
      "Epoch 1/1 |  batch: 215 |  batch loss:   0.092\n",
      "Epoch 1/1 |  batch: 216 |  batch loss:   0.082\n",
      "Epoch 1/1 |  batch: 217 |  batch loss:   0.196\n",
      "Epoch 1/1 |  batch: 218 |  batch loss:   0.103\n",
      "Epoch 1/1 |  batch: 219 |  batch loss:   0.075\n",
      "Epoch 1/1 |  batch: 220 |  batch loss:   0.122\n",
      "Epoch 1/1 |  batch: 221 |  batch loss:   0.288\n",
      "Epoch 1/1 |  batch: 222 |  batch loss:   0.076\n",
      "Epoch 1/1 |  batch: 223 |  batch loss:   0.140\n",
      "Epoch 1/1 |  batch: 224 |  batch loss:   0.034\n",
      "Epoch 1/1 |  batch: 225 |  batch loss:   0.053\n",
      "Epoch 1/1 |  batch: 226 |  batch loss:   0.181\n",
      "Epoch 1/1 |  batch: 227 |  batch loss:   0.163\n",
      "Epoch 1/1 |  batch: 228 |  batch loss:   0.122\n",
      "Epoch 1/1 |  batch: 229 |  batch loss:   0.124\n",
      "Epoch 1/1 |  batch: 230 |  batch loss:   0.055\n",
      "Epoch 1/1 |  batch: 231 |  batch loss:   0.036\n",
      "Epoch 1/1 |  batch: 232 |  batch loss:   0.116\n",
      "Epoch 1/1 |  batch: 233 |  batch loss:   0.113\n",
      "Epoch 1/1 |  batch: 234 |  batch loss:   0.237\n",
      "Epoch 1/1 |  batch: 235 |  batch loss:   0.237\n",
      "Epoch 1/1 |  batch: 236 |  batch loss:   0.190\n",
      "Epoch 1/1 |  batch: 237 |  batch loss:   0.104\n",
      "Epoch 1/1 |  batch: 238 |  batch loss:   0.084\n",
      "Epoch 1/1 |  batch: 239 |  batch loss:   0.043\n",
      "Epoch 1/1 |  batch: 240 |  batch loss:   0.059\n",
      "Epoch 1/1 |  batch: 241 |  batch loss:   0.129\n",
      "Epoch 1/1 |  batch: 242 |  batch loss:   0.079\n",
      "Epoch 1/1 |  batch: 243 |  batch loss:   0.077\n",
      "Epoch 1/1 |  batch: 244 |  batch loss:   0.033\n",
      "Epoch 1/1 |  batch: 245 |  batch loss:   0.055\n",
      "Epoch 1/1 |  batch: 246 |  batch loss:   0.102\n",
      "Epoch 1/1 |  batch: 247 |  batch loss:   0.104\n",
      "Epoch 1/1 |  batch: 248 |  batch loss:   0.107\n",
      "Epoch 1/1 |  batch: 249 |  batch loss:   0.178\n",
      "Epoch 1/1 |  train loss:     0.170 |  train acc:     93.525%\n",
      "Epoch 1/1 |  valid loss:     0.072 |  valid acc:     97.850%\n"
     ]
    }
   ],
   "source": [
    "N_EPOCHS = 1\n",
    "\n",
    "for epoch in range(N_EPOCHS):\n",
    "    \n",
    "    # Train\n",
    "    model.train()  # IMPORTANT\n",
    "    \n",
    "    total_loss, n_correct, n_samples = 0.0, 0, 0\n",
    "    for batch_i, (X, y) in enumerate(train_dl):\n",
    "        X, y = X.to(DEVICE), y.to(DEVICE)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        y_ = model(X)\n",
    "        loss = criterion(y_, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        # Statistics\n",
    "        print(\n",
    "            f\"Epoch {epoch+1}/{N_EPOCHS} |\"\n",
    "            f\"  batch: {batch_i} |\"\n",
    "            f\"  batch loss:   {loss.item():0.3f}\"\n",
    "        )\n",
    "        _, y_label_ = torch.max(y_, 1)\n",
    "        n_correct += (y_label_ == y).sum().item()\n",
    "        total_loss += loss.item() * X.shape[0]\n",
    "        n_samples += X.shape[0]\n",
    "    \n",
    "    print(\n",
    "        f\"Epoch {epoch+1}/{N_EPOCHS} |\"\n",
    "        f\"  train loss: {total_loss / n_samples:9.3f} |\"\n",
    "        f\"  train acc:  {n_correct / n_samples * 100:9.3f}%\"\n",
    "    )\n",
    "    \n",
    "    \n",
    "    # Eval\n",
    "    model.eval()  # IMPORTANT\n",
    "    \n",
    "    total_loss, n_correct, n_samples = 0.0, 0, 0\n",
    "    with torch.no_grad():  # IMPORTANT\n",
    "        for X, y in val_dl:\n",
    "            X, y = X.to(DEVICE), y.to(DEVICE)\n",
    "                    \n",
    "            y_ = model(X)\n",
    "        \n",
    "            # Statistics\n",
    "            _, y_label_ = torch.max(y_, 1)\n",
    "            n_correct += (y_label_ == y).sum().item()\n",
    "            loss = criterion(y_, y)\n",
    "            total_loss += loss.item() * X.shape[0]\n",
    "            n_samples += X.shape[0]\n",
    "\n",
    "    \n",
    "    print(\n",
    "        f\"Epoch {epoch+1}/{N_EPOCHS} |\"\n",
    "        f\"  valid loss: {total_loss / n_samples:9.3f} |\"\n",
    "        f\"  valid acc:  {n_correct / n_samples * 100:9.3f}%\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercise\n",
    "- Create your own module which takes any of the existing pre-trained model as backbone and adds a problem specific head."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self, backbone: nn.Module, n_classes: int):\n",
    "        super().__init__()\n",
    "        # self.backbone\n",
    "        # self.head = init_head(n_classes)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # TODO\n",
    "        return x"
   ]
  }
 ],
 "metadata": {
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}