workshop_10_06_2022_cnn_aws_tsu.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Organize imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import OrderedDict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import PIL\n",
    "from PIL.Image import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch import nn, Tensor, optim\n",
    "from torch.optim import lr_scheduler\n",
    "from torch.utils.data import DataLoader\n",
    "import torchvision\n",
    "from torchvision import (transforms, datasets)\n",
    "from torchvision import models "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "! conda install pytorch torchvision torchaudio -c pytorch\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "! pip install pytorch-lightning\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_lightning import LightningModule, Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchmetrics import Accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workers = multiprocessing.cpu_count()\n",
    "workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = Path('data')\n",
    "path.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! ls {path}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train MNIST classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 16\n",
    "lr = 0.1\n",
    "bs = 2 * 64\n",
    "gamma = 0.7\n",
    "save_model = False\n",
    "log_interval = 10\n",
    "dry_run = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ?? nn.Module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, 3, 1)\n",
    "        self.conv2 = nn.Conv2d(32, 64, 3, 1)\n",
    "        self.dropout1 = nn.Dropout(0.25)\n",
    "        self.dropout2 = nn.Dropout(0.5)\n",
    "        self.fc1 = nn.Linear(9216, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x = F.relu(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = self.dropout1(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.dropout2(x)\n",
    "        x = self.fc2(x)\n",
    "        output = F.log_softmax(x, dim=1)\n",
    "        \n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, device, train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.nll_loss(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if batch_idx % log_interval == 0:\n",
    "            print('Train 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",
    "            if dry_run:\n",
    "                break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss\n",
    "            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability\n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "\n",
    "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
    "        test_loss, correct, len(test_loader.dataset),\n",
    "        100. * correct / len(test_loader.dataset)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_cuda = torch.cuda.is_available()\n",
    "\n",
    "torch.manual_seed(2022)\n",
    "\n",
    "device = torch.device('cuda' if use_cuda else 'cpu')\n",
    "\n",
    "train_kwargs = {'batch_size': 2 * 64}\n",
    "test_kwargs = {'batch_size': 1000}\n",
    "if use_cuda:\n",
    "    cuda_kwargs = {'num_workers': workers,\n",
    "                   'pin_memory': True,\n",
    "                   'shuffle': True}\n",
    "    train_kwargs.update(cuda_kwargs)\n",
    "    test_kwargs.update(cuda_kwargs)\n",
    "\n",
    "transform=transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.1307,), (0.3081,))\n",
    "    ])\n",
    "dataset1 = datasets.MNIST(path / 'mnist', train=True, download=True,\n",
    "                   transform=transform)\n",
    "dataset2 = datasets.MNIST(path / 'mnist', train=False,\n",
    "                   transform=transform)\n",
    "train_loader = DataLoader(dataset1, **train_kwargs)\n",
    "test_loader = DataLoader(dataset2, **test_kwargs)\n",
    "\n",
    "model = Net().to(device)\n",
    "optimizer = optim.Adadelta(model.parameters(), lr=lr)\n",
    "\n",
    "scheduler = lr_scheduler.StepLR(optimizer, step_size=1, gamma=gamma)\n",
    "\n",
    "for epoch in range(1, epochs + 1):\n",
    "    train(model, device, train_loader, optimizer, epoch)\n",
    "    test(model, device, test_loader)\n",
    "    scheduler.step()\n",
    "\n",
    "if save_model:\n",
    "    torch.save(model.state_dict(), path / 'mnist' / 'mnist_cnn.pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train letters classificartion with convolutional neural networkm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prepare data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfms = transforms.Compose([transforms.ToTensor(),\n",
    "                           transforms.Normalize((0.1307,), (0.3081,))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_path = path / 'geomnist_dataset'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geoletters = path / 'geoletters'\n",
    "geoletters.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! wget https://github.com/MaxinAI/school-of-ai/raw/master/data/geoletters/geomnist_dataset.zip -P {geoletters}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import zipfile\n",
    "# with zipfile.ZipFile(path / 'geoletters' / 'geomnist_dataset.zip' , 'r') as zip_ref:\n",
    "#     zip_ref.extractall(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare data loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def img_loader(img_path:Path):\n",
    "    with open(img_path, mode='rb') as fl:\n",
    "        with PIL.Image.open(fl) as img:\n",
    "            return img.convert('L')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = datasets.ImageFolder(image_path / 'train_geo', loader=img_loader, transform=tfms)\n",
    "valid_dataset = datasets.ImageFolder(image_path / 'val_geo', loader=img_loader, transform=tfms)\n",
    "test_dataset = datasets.ImageFolder(image_path / 'test_geo', loader=img_loader, transform=tfms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_dataset, valid_dataset, train_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ?DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, \n",
    "                          num_workers=workers, drop_last=True)\n",
    "valid_loader = DataLoader(valid_dataset, batch_size=64, shuffle=False, \n",
    "                          num_workers=workers)\n",
    "test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False,\n",
    "                        num_workers=workers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset[0][0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def imshow(inp, title=None):\n",
    "    \"\"\"Imshow for Tensor.\"\"\"\n",
    "    inp = inp.numpy().transpose((1, 2, 0))\n",
    "    mean = np.array([0.1307])\n",
    "    std = np.array([0.3081])\n",
    "    inp = std * inp + mean\n",
    "    inp = np.clip(inp, 0, 1)\n",
    "    plt.figure(figsize=(64, 64))\n",
    "    plt.imshow(inp)\n",
    "    if title is not None:\n",
    "        plt.title(title)\n",
    "    plt.pause(0.001)  # pause a bit so that plots are updated\n",
    "\n",
    "\n",
    "# Get a batch of training data\n",
    "inputs, classes = next(iter(train_loader))\n",
    "\n",
    "# Make a grid from batch\n",
    "out = torchvision.utils.make_grid(inputs)\n",
    "\n",
    "imshow(out, title=[classes[x] for x in classes])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_func = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Model initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ?? nn.Flatten"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_channels = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lower_body = nn.Sequential(OrderedDict([('conv1', nn.Conv2d(input_channels, 32, 3)),\n",
    "                                        ('bn1', nn.BatchNorm2d(32)),\n",
    "                                        ('relu1', nn.ReLU(inplace=True)),\n",
    "                                        ('mxpl1', nn.MaxPool2d(2, 2)),\n",
    "                                        ('conv2', nn.Conv2d(32, 64, kernel_size=3)),\n",
    "                                        ('bn2', nn.BatchNorm2d(64)),\n",
    "                                        ('relu2', nn.ReLU(inplace=True)),\n",
    "                                        ('mxpl2', nn.MaxPool2d(2, 2)),\n",
    "                                        ('drop1', nn.Dropout2d(p=0.25))]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_body = nn.Sequential(OrderedDict([('conv3', nn.Conv2d(64, 128, kernel_size=3)),\n",
    "                                       ('bn3', nn.BatchNorm2d(128)),\n",
    "                                       ('relu3', nn.ReLU(inplace=True)),\n",
    "                                       ('mxpl3', nn.MaxPool2d(2, 2)),\n",
    "                                       ('drop2', nn.Dropout2d(p=0.25))]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "linear_body = nn.Sequential(OrderedDict([('flatten', nn.Flatten()),\n",
    "                                         ('ln1', nn.Linear(2 * 2 * 128, 1024, bias=True)),\n",
    "                                         ('bn2', nn.BatchNorm1d(1024)),\n",
    "                                         ('relu3', nn.ReLU(inplace=True)),\n",
    "                                         ('drop2', nn.Dropout(p=0.25))]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "linear_body"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nn.Sequential(lower_body, conv_body, linear_body, nn.Linear(1024, 33))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LetterNet(nn.Module):\n",
    "    \"\"\"Full double letters network implementation\"\"\"\n",
    "\n",
    "    def __init__(self, input_channels=1, num_classes=33):\n",
    "        super(LetterNet, self).__init__()\n",
    "        self.conv_part = lower_body\n",
    "        self.dub_part = conv_body\n",
    "        self.fc_part = linear_body\n",
    "        self.fc = nn.Linear(1024, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv_part(x)\n",
    "        x = self.dub_part(x)\n",
    "        x = self.fc_part(x)\n",
    "        logits = self.fc(x)\n",
    "\n",
    "        return logits\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_out(w:int, k: int, p: int, s: int):\n",
    "    return int(((w - k + 2*p)/s)+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_out(32, 3, 0, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_out(conv_out(conv_out(conv_out(32, 3, 0, 1), 2, 0, 1), 3, 0, 1), 2, 0, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LetterNet2(nn.Module):\n",
    "    \"\"\"Full double letters network implementation\"\"\"\n",
    "\n",
    "    def __init__(self, input_channels=1, num_classes=33):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(input_channels, 32, 3)\n",
    "        self.pool1 = nn.MaxPool2d(2, 2)\n",
    "        self.conv2 = nn.Conv2d(32, 64, 3)\n",
    "        self.pool2 = nn.MaxPool2d(2, 2)\n",
    "        self.fc = nn.Linear(6 * 6 * 64, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.pool1(x)\n",
    "        x = self.conv2(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.pool2(x)\n",
    "        x = torch.flatten(x, start_dim=1)\n",
    "        logits = self.fc(x)\n",
    "\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sz = 32 # 34\n",
    "x_test = torch.randn(4 , 1, sz, sz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = LetterNet2()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "o_test = model(x_test)\n",
    "o_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = LetterNet()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create lightning module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClassifierNetLt(LightningModule):\n",
    "    def __init__(\n",
    "        self, model: nn.Module, loss_fn=loss_func, metrics=Accuracy()):\n",
    "        super().__init__()\n",
    "        self.model = model\n",
    "        self.accuracy = metrics\n",
    "        self.lr=0.01\n",
    "        self.loss_fn = loss_fn\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "    def training_step(self, batch, batch_nb):\n",
    "        x, y = batch\n",
    "        z = self(x)\n",
    "        loss = self.loss_fn(z, y)\n",
    "        \n",
    "        return loss\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        z = self(x)\n",
    "        loss = self.loss_fn(z, y)\n",
    "        preds = torch.argmax(z, dim=1)\n",
    "        self.accuracy(preds, y)\n",
    "\n",
    "        # Calling self.log will surface up scalars for you in TensorBoard\n",
    "        self.log(\"val_loss\", loss, prog_bar=True)\n",
    "        self.log(\"val_acc\", self.accuracy, prog_bar=True)\n",
    "        \n",
    "        return loss\n",
    "    \n",
    "    def test_step(self, batch, batch_idx):\n",
    "        # Here we just reuse the validation_step for testing\n",
    "        return self.validation_step(batch, batch_idx)\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = optim.Adam(self.model.parameters(), lr=self.lr)\n",
    "        lr_scheduler_cos = lr_scheduler.CosineAnnealingWarmRestarts(\n",
    "            optimizer, \n",
    "            T_0=10, T_mult=2, \n",
    "            eta_min=0.0001, \n",
    "            last_epoch=-1,\n",
    "            verbose=True)\n",
    "        \n",
    "        return [optimizer], [lr_scheduler_cos]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "AVAIL_GPUS = min(1, torch.cuda.device_count())\n",
    "AVAIL_GPUS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_lt = ClassifierNetLt(model)\n",
    "trainer = Trainer(\n",
    "    gpus=AVAIL_GPUS,\n",
    "    max_epochs=16,\n",
    "    progress_bar_refresh_rate=20,\n",
    "    auto_lr_find=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! ls {geoletters}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "curr_date = datetime.now().strftime('%d_%m_%Y').lower()\n",
    "curr_date"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geletters_ckpts = geoletters / f'geoletters_{curr_date}'\n",
    "geletters_ckpts.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.fit(\n",
    "    model_lt, train_loader, valid_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Find learning rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_lt_lr = ClassifierNetLt(LetterNet())\n",
    "trainer_lr = Trainer(\n",
    "    gpus=AVAIL_GPUS,\n",
    "    max_epochs=3,\n",
    "    progress_bar_refresh_rate=20,\n",
    "    auto_lr_find=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr_finder = trainer_lr.tuner.lr_find(\n",
    "    model_lt_lr, train_loader, valid_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Results can be found in\n",
    "lr_finder.results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot with\n",
    "fig = lr_finder.plot(suggest=True)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pick point based on plot, or get suggestion\n",
    "new_lr = lr_finder.suggestion()\n",
    "\n",
    "# update hparams of the model\n",
    "model_lt_lr.lr = new_lr\n",
    "new_lr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer_lr.tune(\n",
    "    model_lt_lr, train_loader, valid_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Size dependency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sz = 32 # 34\n",
    "x_test = torch.randn(4 , 1, sz, sz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model_lt.freeze()\n",
    "model_lt(x_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Bigger Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet = models.resnet18(pretrained=False, progress=True)\n",
    "resnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)\n",
    "resnet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Re-initialize data loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 2 * 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(train_dataset, batch_size=bs, shuffle=True, \n",
    "                          num_workers=workers, drop_last=True)\n",
    "valid_loader = DataLoader(valid_dataset, batch_size=bs, shuffle=False, \n",
    "                          num_workers=workers)\n",
    "test_loader = DataLoader(test_dataset, batch_size=bs, shuffle=False,\n",
    "                        num_workers=workers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet_lt = ClassifierNetLt(resnet)\n",
    "trainer = Trainer(\n",
    "    gpus=AVAIL_GPUS,\n",
    "    benchmark=True,\n",
    "    max_epochs=16,\n",
    "    progress_bar_refresh_rate=20,\n",
    "    auto_lr_find=True,\n",
    "#     stochastic_weight_avg=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "trainer.fit(\n",
    "    resnet_lt, train_loader, valid_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save the model state dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = resnet_lt.model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net.state_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(resnet_lt.model.state_dict(), str(geletters_ckpts / 'resnet_18_state_dict.pth'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet_ld = models.resnet18(pretrained=False, progress=False)\n",
    "resnet_ld.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_dict = torch.load(str(geletters_ckpts / 'resnet_18_state_dict.pth'), map_location='cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet_ld.load_state_dict(state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geletters_ckpts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "! ls {geletters_ckpts}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save entire model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(resnet_lt.model, geletters_ckpts / 'resnet_18_entire.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet_ld = torch.load(geletters_ckpts / 'resnet_18_entire.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# resnet_ld.state_dict(), resnet_lt.model.state_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "resnet_ld"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's try to inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run inference service in different session"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "! pip install -U flask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "geletters_ckpts / 'resnet_18_entire.pth'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! ls data/geoletters/geoletters_labels.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# geletters_ckpts.absolute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Union"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "! pip install opencv-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import PIL\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import no_grad\n",
    "from torch.jit import ScriptModule\n",
    "from torchvision.models import (resnet34, resnet50, wide_resnet50_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "size = 256\n",
    "imsz = 224\n",
    "IMG_SUFF = {'.jpg', '.jpeg', '.png'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ToPILImage(object):\n",
    "    \"\"\"Convert inout image to PIL image\"\"\"\n",
    "\n",
    "    def __init__(self, mode=None):\n",
    "        super().__init__()\n",
    "        self.to_pil = transforms.ToPILImage(mode=mode)\n",
    "\n",
    "    def convert(self, img: Union[np.ndarray, PIL.Image.Image]):\n",
    "        \"\"\"\n",
    "        Converts image to the PIL format\n",
    "        Args:\n",
    "            img: inout image\n",
    "\n",
    "        Returns:\n",
    "            converted image\n",
    "        \"\"\"\n",
    "        return img if isinstance(img, PIL.Image.Image) else self.to_pil(img)\n",
    "\n",
    "    def __call__(self, *args, **kwargs):\n",
    "        return self.convert(*args, **kwargs)\n",
    "\n",
    "    def __repr__(self):\n",
    "        format_string = self.__class__.__name__ + '('\n",
    "        if self.to_pil.mode is not None:\n",
    "            format_string += f'mode={self.to_pil.mode}'\n",
    "        format_string += ')'\n",
    "        return format_string\n",
    "\n",
    "\n",
    "class Img2Vec(object):\n",
    "    \"\"\"Model wrapper for image embedding\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self, backbone: Union[nn.Module, ScriptModule], trfm: transforms, device: str = 'cpu'):\n",
    "        super().__init__()\n",
    "        self.device = torch.device(device)\n",
    "        self.backbone = (backbone.eval() if hasattr(backbone, 'eval') else backbone).to(device)\n",
    "        self.trfm = trfm\n",
    "\n",
    "    def preprocess(self, *xs: Union[np.ndarray, PIL.Image.Image]) -> Tensor:\n",
    "        \"\"\"\n",
    "        Transform data before model\n",
    "        Args:\n",
    "            *xs: input data\n",
    "\n",
    "        Returns:\n",
    "            processed data for model\n",
    "        \"\"\"\n",
    "        return torch.stack([self.trfm(x) for x in xs]).to(self.device)\n",
    "\n",
    "    @no_grad()\n",
    "    def forward(self, *xs: np.ndarray) -> np.ndarray:\n",
    "        tns = self.preprocess(*xs)\n",
    "        rts = self.backbone(tns)\n",
    "        y = rts.cpu().data.numpy()\n",
    "\n",
    "        return y\n",
    "\n",
    "    def __call__(self, *args, **kwargs):\n",
    "        return self.forward(*args, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vec_trsfm = transforms.Compose([ToPILImage(mode='RGB'),\n",
    "                                transforms.Resize(size),\n",
    "                                transforms.CenterCrop(imsz),\n",
    "                                transforms.ToTensor(),\n",
    "                                transforms.Normalize(\n",
    "                                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prepare data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "search_path = path / 'search'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_paths = [dp for dp in search_path.iterdir() if dp.is_dir()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_pts = [im_pt for dp in dir_paths for im_pt in dp.iterdir() if im_pt.suffix in IMG_SUFF]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_img(im_pt):\n",
    "    img = cv2.imread(str(im_pt), cv2.IMREAD_ANYCOLOR)\n",
    "    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)\n",
    "    \n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imgs = [read_img(ip) for ip in img_pts]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialize features extractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cut = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "body = wide_resnet50_2(pretrained=True, progress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pool = nn.AdaptiveAvgPool2d(output_size=(1, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flatten = nn.Flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp1 = torch.rand(1, 2048, 7, 7)\n",
    "inp2 = torch.rand(1, 2048, 14, 14)\n",
    "inp3 = torch.rand(1, 2048, 2, 2)\n",
    "inp4 = torch.rand(1, 2048, 28, 28)\n",
    "inp1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "oupt1 = pool(inp1)\n",
    "oupt2 = pool(inp2)\n",
    "oupt3 = pool(inp3)\n",
    "oupt4 = pool(inp4)\n",
    "oupt1.shape, oupt2.shape, oupt3.shape, oupt4.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flatten(oupt1).shape, flatten(oupt2).shape, flatten(oupt2).shape, flatten(oupt2).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "body"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "backbone = nn.Sequential(*list(body.children())[:-cut])\n",
    "net = nn.Sequential(backbone, nn.Flatten())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_vec = Img2Vec(net, vec_trsfm, device='cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tqdm(imgs, desc='Vectorizing images') as p_imgs:\n",
    "    vecs = [img_vec(im)[0] for im in p_imgs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vecs[0].shape, len(vecs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_vecs = list(zip(imgs, vecs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# img_vecs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# vecs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compare vectors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "! pip install scipy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.spatial.distance import cosine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def top_vecs(qi, vec_func=img_vec, db=img_vecs, top_k=5):\n",
    "    qv = vec_func(qi)[0]\n",
    "    resul_pts = [(cosine(qv, vc), pt) for pt, vc in db]\n",
    "    resul_pts = sorted(resul_pts, key=lambda x: x[0], reverse=False)\n",
    "    resul_pts = resul_pts[:top_k]\n",
    "    \n",
    "    return resul_pts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Query images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_path = path / 'queries'\n",
    "query_path.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! ls {query_path}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qim = read_img(query_path / 'rv_2.jpeg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_vec(imgs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = top_vecs(qim, vec_func=img_vec, db=img_vecs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(qim)\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for dist, res_img in res:\n",
    "    plt.imshow(res_img)\n",
    "    plt.show()\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Image similarity with CLIP model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "! pip install git+https://github.com/openai/CLIP.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import clip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model, preprocess = clip.load(\"ViT-B/32\", device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_vec_clip = Img2Vec(model, preprocess, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_vec_clip.backbone = img_vec_clip.backbone.encode_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tqdm(imgs, desc='Reading images') as p_paths:\n",
    "    pil_imgs = [Image.fromarray(im) for im in imgs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pil_imgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ??preprocess.transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# preprocess(pil_imgs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tqdm(pil_imgs, desc='Vectorizing images') as p_imgs:\n",
    "    clip_vecs = [img_vec_clip(im)[0] for im in p_imgs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clip_img_vecs = list(zip(imgs, clip_vecs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clip_vecs[0].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Query images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! ls {query_path}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_path = path / 'queries'\n",
    "query_path.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qim = read_img(query_path / 'rv_2.jpeg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = top_vecs(Image.fromarray(qim), vec_func=img_vec_clip, db=clip_img_vecs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(qim)\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for dist, res_img in res:\n",
    "    plt.imshow(res_img)\n",
    "    plt.show()\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine-tuning model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Dlownload and store data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hymenoptera = path / 'hymenoptera'\n",
    "hymenoptera.mkdir(exist_ok=True)\n",
    "hymenoptera"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! wget https://download.pytorch.org/tutorial/hymenoptera_data.zip -P {hymenoptera}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! ls {hymenoptera}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import zipfile\n",
    "with zipfile.ZipFile(hymenoptera / 'hymenoptera_data.zip' , 'r') as zip_ref:\n",
    "    zip_ref.extractall(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! ls {path / 'hymenoptera_data'}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialize data loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data augmentation and normalization for training\n",
    "# Just normalization for validation\n",
    "train_transforms = 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",
    "    ])\n",
    "val_transforms = transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "    ])\n",
    "\n",
    "data_dir = path / 'hymenoptera_data'\n",
    "\n",
    "train_dataset = datasets.ImageFolder(data_dir / 'train', train_transforms)\n",
    "val_dataset = datasets.ImageFolder(data_dir / 'val', val_transforms)\n",
    "\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True, \n",
    "                          num_workers=workers, drop_last=True)\n",
    "val_datloader = DataLoader(val_dataset, batch_size=64, shuffle=False, \n",
    "                          num_workers=workers)\n",
    "test_datloader = DataLoader(val_dataset, batch_size=4, shuffle=True, \n",
    "                          num_workers=workers)\n",
    "\n",
    "dataset_sizes = len(train_dataloader) + len(val_datloader)\n",
    "class_names = train_dataset.classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def imshow(inp, title=None):\n",
    "    \"\"\"Imshow for Tensor.\"\"\"\n",
    "    inp = inp.numpy().transpose((1, 2, 0))\n",
    "    mean = np.array([0.485, 0.456, 0.406])\n",
    "    std = np.array([0.229, 0.224, 0.225])\n",
    "    inp = std * inp + mean\n",
    "    inp = np.clip(inp, 0, 1)\n",
    "    plt.figure(figsize=(32, 32))\n",
    "    plt.imshow(inp)\n",
    "    if title is not None:\n",
    "        plt.title(title)\n",
    "    plt.pause(0.001)  # pause a bit so that plots are updated\n",
    "\n",
    "\n",
    "# Get a batch of training data\n",
    "inputs, classes = next(iter(train_dataloader))\n",
    "\n",
    "# Make a grid from batch\n",
    "out = torchvision.utils.make_grid(inputs)\n",
    "\n",
    "imshow(out, title=[class_names[x] for x in classes])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_model(model, num_images=6):\n",
    "    model.eval()\n",
    "    images_so_far = 0\n",
    "    fig = plt.figure()\n",
    "    model.freeze()\n",
    "    with torch.no_grad():\n",
    "        for i, (inputs, labels) in enumerate(test_datloader):\n",
    "            outputs = model(inputs)\n",
    "            _, preds = torch.max(outputs, 1)\n",
    "\n",
    "            for j in range(inputs.size()[0]):\n",
    "                images_so_far += 1\n",
    "                ax = plt.subplot(num_images//2, 2, images_so_far)\n",
    "                ax.axis('off')\n",
    "                ax.set_title(f'predicted: {class_names[preds[j]]} ground true: {class_names[labels[j]]}')\n",
    "                imshow(inputs.cpu().data[j])\n",
    "\n",
    "                if images_so_far == num_images:\n",
    "                    return\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, (inputs, labels) in enumerate(test_datloader):\n",
    "    print(i, labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialize the model for fine-tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ft = models.resnet18(pretrained=True, progress=True)\n",
    "num_ftrs = model_ft.fc.in_features\n",
    "# Here the size of each output sample is set to 2.\n",
    "# Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).\n",
    "model_ft.fc = nn.Linear(num_ftrs, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClassifierHymLt(ClassifierNetLt):\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        optimizer_ft = optim.SGD(\n",
    "            self.model.parameters(), lr=0.001, momentum=0.9)\n",
    "        exp_lr_scheduler = lr_scheduler.StepLR(\n",
    "            optimizer_ft, step_size=7, gamma=0.1)\n",
    "        \n",
    "        return [optimizer_ft], [exp_lr_scheduler]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ft_lt = ClassifierHymLt(model_ft)\n",
    "trainer = Trainer(\n",
    "    gpus=AVAIL_GPUS,\n",
    "    benchmark=True,\n",
    "    max_epochs=32,\n",
    "    progress_bar_refresh_rate=20,\n",
    "    auto_lr_find=True,\n",
    "#     stochastic_weight_avg=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.fit(model_ft_lt, train_dataloader, val_datloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ft_lt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "visualize_model(model_ft_lt, num_images=12)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Transfer learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_conv = models.resnet18(pretrained=True)\n",
    "for param in model_conv.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "# Parameters of newly constructed modules have requires_grad=True by default\n",
    "num_ftrs = model_conv.fc.in_features\n",
    "model_conv.fc = nn.Linear(num_ftrs, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClassifierHymFtLt(ClassifierNetLt):\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        optimizer_ft = optim.SGD(\n",
    "            self.model.fc.parameters(), lr=0.001, momentum=0.9)\n",
    "        exp_lr_scheduler = lr_scheduler.StepLR(\n",
    "            optimizer_ft, step_size=7, gamma=0.1)\n",
    "        \n",
    "        return [optimizer_ft], [exp_lr_scheduler]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ft_lt = ClassifierHymFtLt(model_conv)\n",
    "trainer = Trainer(\n",
    "    gpus=AVAIL_GPUS,\n",
    "    benchmark=True,\n",
    "    max_epochs=32,\n",
    "    progress_bar_refresh_rate=20,\n",
    "    auto_lr_find=True,\n",
    "#     stochastic_weight_avg=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.fit(model_ft_lt, train_dataloader, val_datloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "visualize_model(model_ft_lt, num_images=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.4"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}