Train the model

notebook/code.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -17,6 +17,587 @@
     "import numpy as np"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# import resnet from torch\n",
+    "import torch.library\n",
+    "from torchvision.models import squeezenet1_1\n",
+    "from torchvision.models import resnet50\n",
+    "from torchvision.models import resnet18\n",
+    "from torchvision.models import mobilenet_v2\n",
+    "from torchvision import transforms\n",
+    "from torchvision.datasets import ImageFolder\n",
+    "from torch.utils.data import DataLoader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class_num = 5\n",
+    "classes = ['Ak', 'Ala_Idris', 'Buzgulu', 'Dimnit', 'Nazli']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "d:\\Softwares\\Anaconda3\\lib\\site-packages\\torchvision\\models\\_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
+      "  warnings.warn(\n",
+      "d:\\Softwares\\Anaconda3\\lib\\site-packages\\torchvision\\models\\_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=MobileNet_V2_Weights.IMAGENET1K_V1`. You can also use `weights=MobileNet_V2_Weights.DEFAULT` to get the most up-to-date weights.\n",
+      "  warnings.warn(msg)\n"
+     ]
+    }
+   ],
+   "source": [
+    "model = mobilenet_v2(pretrained=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "MobileNetV2(\n",
+      "  (features): Sequential(\n",
+      "    (0): Conv2dNormActivation(\n",
+      "      (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
+      "      (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      (2): ReLU6(inplace=True)\n",
+      "    )\n",
+      "    (1): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
+      "          (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (2): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(16, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(96, 96, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=96, bias=False)\n",
+      "          (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(96, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (3): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(24, 144, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(144, 144, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=144, bias=False)\n",
+      "          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(144, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (4): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(24, 144, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(144, 144, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=144, bias=False)\n",
+      "          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(144, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (5): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(32, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=192, bias=False)\n",
+      "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(192, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (6): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(32, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=192, bias=False)\n",
+      "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(192, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (7): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(32, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(192, 192, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=192, bias=False)\n",
+      "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(192, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (8): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384, bias=False)\n",
+      "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(384, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (9): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384, bias=False)\n",
+      "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(384, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (10): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384, bias=False)\n",
+      "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(384, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (11): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384, bias=False)\n",
+      "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(384, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (12): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(96, 576, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(576, 576, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=576, bias=False)\n",
+      "          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(576, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (13): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(96, 576, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(576, 576, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=576, bias=False)\n",
+      "          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(576, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (14): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(96, 576, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(576, 576, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=576, bias=False)\n",
+      "          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(576, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (15): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
+      "          (1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (16): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
+      "          (1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (17): InvertedResidual(\n",
+      "      (conv): Sequential(\n",
+      "        (0): Conv2dNormActivation(\n",
+      "          (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "          (1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (1): Conv2dNormActivation(\n",
+      "          (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
+      "          (1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "          (2): ReLU6(inplace=True)\n",
+      "        )\n",
+      "        (2): Conv2d(960, 320, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "        (3): BatchNorm2d(320, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      )\n",
+      "    )\n",
+      "    (18): Conv2dNormActivation(\n",
+      "      (0): Conv2d(320, 1280, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+      "      (1): BatchNorm2d(1280, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+      "      (2): ReLU6(inplace=True)\n",
+      "    )\n",
+      "  )\n",
+      "  (classifier): Sequential(\n",
+      "    (0): Dropout(p=0.2, inplace=False)\n",
+      "    (1): Linear(in_features=1280, out_features=1000, bias=True)\n",
+      "  )\n",
+      ")\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "transform = transforms.Compose([\n",
+    "  transforms.RandomResizedCrop(224),\n",
+    "  transforms.RandomHorizontalFlip(),\n",
+    "  transforms.ToTensor(),\n",
+    "  transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "training_set = ImageFolder('../data/train', transform=transform)\n",
+    "test_set = ImageFolder('../data/test', transform=transform)\n",
+    "val_set = ImageFolder('../data/val', transform=transform)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 47,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch_size = 8\n",
+    "epochs = 5\n",
+    "lr = 1e-5\n",
+    "loss_fn = nn.CrossEntropyLoss()\n",
+    "optimizer = optim.Adam(model.parameters(), lr=lr)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 48,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_loader = DataLoader(training_set, batch_size=batch_size, shuffle=True)\n",
+    "test_loader = DataLoader(test_set, batch_size=batch_size)\n",
+    "val_loader = DataLoader(val_set, batch_size=batch_size)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 49,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.classifier[1] = nn.Linear(in_features=1280, out_features=class_num)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 52,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "epochs = 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 55,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Out:  [2, 3, 1, 1, 0, 4, 0, 2]\n",
+      "Target:  tensor([2, 3, 1, 1, 0, 2, 1, 2])\n",
+      "Epoch: 0 [0/400 (0%)]\tLoss: 0.918071\n",
+      "Out:  [3, 3, 4, 2, 2, 0, 1, 4]\n",
+      "Target:  tensor([3, 3, 4, 0, 0, 0, 1, 4])\n",
+      "Out:  [1, 4, 4, 1, 4, 4, 2, 3]\n",
+      "Target:  tensor([1, 4, 4, 1, 4, 4, 2, 3])\n",
+      "Out:  [4, 2, 4, 0, 1, 3, 3, 1]\n",
+      "Target:  tensor([4, 2, 4, 0, 2, 3, 3, 2])\n",
+      "Out:  [0, 2, 3, 4, 1, 4, 3, 1]\n",
+      "Target:  tensor([1, 1, 3, 1, 1, 4, 3, 1])\n",
+      "Out:  [4, 2, 0, 0, 4, 4, 1, 3]\n",
+      "Target:  tensor([0, 2, 0, 0, 4, 2, 0, 3])\n",
+      "Out:  [3, 1, 1, 4, 4, 3, 0, 1]\n",
+      "Target:  tensor([2, 2, 1, 2, 4, 3, 0, 1])\n",
+      "Out:  [2, 4, 3, 1, 1, 0, 4, 3]\n",
+      "Target:  tensor([2, 1, 3, 1, 1, 0, 4, 3])\n",
+      "Out:  [4, 4, 1, 1, 4, 2, 0, 3]\n",
+      "Target:  tensor([4, 3, 1, 1, 4, 2, 0, 3])\n",
+      "Out:  [1, 0, 0, 4, 3, 2, 3, 2]\n",
+      "Target:  tensor([2, 0, 2, 4, 3, 2, 3, 2])\n",
+      "Out:  [2, 2, 0, 3, 1, 4, 0, 4]\n",
+      "Target:  tensor([2, 2, 0, 0, 1, 4, 0, 3])\n",
+      "Epoch: 0 [80/400 (20%)]\tLoss: 0.715335\n",
+      "Out:  [0, 4, 4, 3, 1, 2, 0, 1]\n",
+      "Target:  tensor([1, 4, 2, 0, 1, 2, 0, 1])\n",
+      "Out:  [3, 2, 4, 4, 0, 2, 0, 3]\n",
+      "Target:  tensor([3, 2, 4, 4, 3, 2, 1, 3])\n",
+      "Out:  [0, 0, 3, 4, 1, 4, 2, 2]\n",
+      "Target:  tensor([0, 0, 3, 4, 1, 4, 2, 2])\n",
+      "Out:  [0, 4, 0, 2, 4, 1, 2, 4]\n",
+      "Target:  tensor([0, 3, 0, 2, 0, 1, 3, 3])\n",
+      "Out:  [3, 2, 4, 4, 1, 3, 4, 3]\n",
+      "Target:  tensor([2, 2, 2, 4, 1, 3, 2, 3])\n",
+      "Out:  [4, 1, 3, 2, 3, 3, 4, 0]\n",
+      "Target:  tensor([4, 1, 2, 1, 3, 1, 4, 1])\n",
+      "Out:  [4, 2, 4, 4, 0, 0, 3, 1]\n",
+      "Target:  tensor([4, 2, 4, 4, 1, 0, 3, 1])\n",
+      "Out:  [4, 0, 3, 1, 1, 3, 4, 2]\n",
+      "Target:  tensor([4, 1, 0, 1, 1, 3, 0, 2])\n",
+      "Out:  [1, 0, 0, 3, 2, 4, 2, 4]\n",
+      "Target:  tensor([1, 3, 0, 3, 2, 2, 2, 2])\n",
+      "Out:  [4, 4, 3, 1, 4, 0, 2, 3]\n",
+      "Target:  tensor([4, 4, 3, 1, 4, 0, 2, 4])\n",
+      "Epoch: 0 [160/400 (40%)]\tLoss: 0.672349\n",
+      "Out:  [3, 2, 2, 0, 3, 1, 4, 4]\n",
+      "Target:  tensor([3, 2, 2, 0, 3, 0, 3, 4])\n",
+      "Out:  [1, 1, 3, 0, 2, 4, 0, 3]\n",
+      "Target:  tensor([1, 1, 4, 0, 0, 4, 0, 3])\n",
+      "Out:  [4, 2, 4, 1, 1, 1, 4, 3]\n",
+      "Target:  tensor([4, 1, 4, 1, 1, 1, 4, 3])\n",
+      "Out:  [1, 4, 0, 3, 3, 4, 2, 2]\n",
+      "Target:  tensor([1, 4, 3, 3, 3, 3, 3, 2])\n",
+      "Out:  [3, 2, 0, 2, 1, 4, 0, 4]\n",
+      "Target:  tensor([3, 2, 0, 1, 1, 2, 0, 0])\n",
+      "Out:  [3, 4, 4, 1, 1, 1, 3, 2]\n",
+      "Target:  tensor([3, 4, 4, 1, 1, 1, 2, 1])\n",
+      "Out:  [4, 1, 2, 1, 4, 3, 4, 2]\n",
+      "Target:  tensor([4, 1, 2, 1, 4, 3, 4, 1])\n",
+      "Out:  [0, 3, 4, 2, 0, 3, 3, 3]\n",
+      "Target:  tensor([0, 0, 4, 2, 0, 3, 3, 3])\n",
+      "Out:  [3, 2, 2, 4, 0, 4, 1, 2]\n",
+      "Target:  tensor([3, 2, 2, 4, 0, 4, 0, 2])\n",
+      "Out:  [1, 2, 4, 2, 1, 4, 3, 0]\n",
+      "Target:  tensor([1, 1, 4, 1, 1, 0, 0, 0])\n",
+      "Epoch: 0 [240/400 (60%)]\tLoss: 1.097419\n",
+      "Out:  [4, 4, 1, 3, 4, 1, 0, 3]\n",
+      "Target:  tensor([4, 4, 1, 3, 4, 1, 0, 3])\n",
+      "Out:  [0, 3, 3, 3, 4, 1, 2, 2]\n",
+      "Target:  tensor([0, 3, 0, 0, 4, 2, 2, 2])\n",
+      "Out:  [3, 4, 3, 2, 0, 4, 2, 0]\n",
+      "Target:  tensor([3, 4, 0, 0, 0, 4, 2, 0])\n",
+      "Out:  [3, 1, 3, 0, 1, 3, 4, 1]\n",
+      "Target:  tensor([1, 0, 3, 0, 1, 3, 0, 1])\n",
+      "Out:  [4, 3, 2, 0, 4, 2, 1, 2]\n",
+      "Target:  tensor([4, 3, 2, 0, 4, 2, 2, 2])\n",
+      "Out:  [3, 4, 3, 2, 4, 1, 3, 0]\n",
+      "Target:  tensor([3, 4, 3, 2, 4, 1, 3, 0])\n",
+      "Out:  [0, 0, 3, 2, 1, 0, 4, 4]\n",
+      "Target:  tensor([0, 0, 3, 2, 2, 0, 4, 4])\n",
+      "Out:  [4, 0, 4, 0, 3, 2, 3, 3]\n",
+      "Target:  tensor([4, 3, 4, 0, 3, 2, 3, 0])\n",
+      "Out:  [1, 0, 3, 4, 0, 1, 2, 2]\n",
+      "Target:  tensor([1, 0, 3, 4, 0, 0, 2, 2])\n",
+      "Out:  [1, 2, 4, 3, 2, 3, 4, 0]\n",
+      "Target:  tensor([1, 2, 3, 3, 2, 3, 4, 0])\n",
+      "Epoch: 0 [320/400 (80%)]\tLoss: 0.581446\n",
+      "Out:  [4, 3, 1, 2, 1, 4, 1, 4]\n",
+      "Target:  tensor([4, 1, 1, 2, 1, 3, 2, 2])\n",
+      "Out:  [0, 3, 3, 0, 1, 4, 2, 1]\n",
+      "Target:  tensor([3, 3, 3, 0, 1, 4, 2, 1])\n",
+      "Out:  [1, 3, 4, 3, 2, 2, 0, 4]\n",
+      "Target:  tensor([1, 3, 0, 3, 1, 2, 0, 4])\n",
+      "Out:  [2, 0, 3, 4, 3, 0, 4, 2]\n",
+      "Target:  tensor([2, 0, 0, 4, 4, 0, 4, 2])\n",
+      "Out:  [3, 4, 0, 4, 3, 4, 2, 0]\n",
+      "Target:  tensor([3, 4, 0, 4, 3, 4, 0, 0])\n",
+      "Out:  [1, 0, 3, 2, 2, 0, 1, 4]\n",
+      "Target:  tensor([1, 0, 3, 2, 2, 0, 1, 4])\n",
+      "Out:  [0, 1, 2, 1, 4, 3, 1, 2]\n",
+      "Target:  tensor([0, 1, 2, 1, 4, 3, 1, 2])\n",
+      "Out:  [3, 4, 1, 1, 0, 4, 2, 2]\n",
+      "Target:  tensor([3, 4, 1, 1, 0, 3, 2, 4])\n",
+      "Out:  [4, 4, 2, 1, 2, 3, 0, 4]\n",
+      "Target:  tensor([4, 4, 2, 1, 2, 3, 0, 4])\n",
+      "\n",
+      "Test set: Avg. loss: 0.0947, Accuracy: 43/50 (86%)\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "# train the model\n",
+    "for epoch in range(epochs):\n",
+    "  model.train()\n",
+    "  for batch_idx, (data, target) in enumerate(train_loader):\n",
+    "    optimizer.zero_grad()\n",
+    "    output = model(data)\n",
+    "    print(\"Out: \", [a.argmax().item() for a in output])\n",
+    "    print(\"Target: \", target)\n",
+    "    loss = loss_fn(output, target)\n",
+    "    loss.backward()\n",
+    "    optimizer.step()\n",
+    "    if batch_idx % 10 == 0:\n",
+    "      print('Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
+    "        epoch, batch_idx * len(data), len(train_loader.dataset),\n",
+    "        100. * batch_idx / len(train_loader), loss.item()\n",
+    "      ))\n",
+    "\n",
+    "  # test the model\n",
+    "  model.eval()\n",
+    "  test_loss = 0\n",
+    "  correct = 0\n",
+    "  with torch.no_grad():\n",
+    "    for data, target in test_loader:\n",
+    "      output = model(data)\n",
+    "      test_loss += loss_fn(output, target).item()\n",
+    "      pred = output.argmax(dim=1, keepdim=True)\n",
+    "      correct += pred.eq(target.view_as(pred)).sum().item()\n",
+    "\n",
+    "  test_loss /= len(test_loader.dataset)\n",
+    "  print('\\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
+    "    test_loss, correct, len(test_loader.dataset),\n",
+    "    100. * correct / len(test_loader.dataset)\n",
+    "  ))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 51,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_scripted = torch.jit.script(model)\n",
+    "model_scripted.save('../models/mobilenet.pt')"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,