{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import pytorch and machine learning stuff\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# import other stuff\n",
    "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,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}