seminars/201/seminar_04.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "1bbc68bf",
      "metadata": {
        "id": "1bbc68bf"
      },
      "source": [
        "## 1. Скачиваем данные\n",
        "\n",
        "[Ссылка на данные](https://www.kaggle.com/lantian773030/pokemonclassification)\n",
        "\n",
        "![](https://storage.googleapis.com/kagglesdsdata/datasets/410745/786434/PokemonData/Psyduck/8eb3ab4da8a44cda9ba15d48ce893a06.jpg?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F20221017%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20221017T070101Z&X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=698923dd2ebe39354e28db4339e5e66fa163815d658877fa929c22d072d3eb8d35fbdb34c7e938a5f4709095d9cff698bec1699716c669d4e27bc3437f1b4e783b8b1fe04a6032a4ec8ff5731da86c89f0d054041bcc400b5efd5f52e8407cc985e026c4a23b4f28e95e2751cfcd926d129914943bd6d77edf03c96e57d58baacfad8f8ea6839ad002fb5efae522258a40be3c06352a230e89c7c7eb8d56bf05a7c5f4e96435fc6360f545aa504d64e175f8d33020ccd1d5d7d59f9528ca40869d8117b083cd89df48e07c03ebb65356d28aa6d6927b8bde4d8407366f342659d2de99b36720e8565847cf1be37f3697a988467f8455e5e60fb0b3c01fee70b0)\n",
        "\n",
        "Скачивание данных в следующей ячейке требует библиотеку `kaggle` и наличия файла с API ключем `~/.kaggle/kaggle.json`.\n",
        "Инструкция вот здесь https://www.kaggle.com/general/156610\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install kaggle"
      ],
      "metadata": {
        "id": "6ftFGQPFh-Uo",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "47255c2c-e148-42fa-b757-bb4632038823"
      },
      "id": "6ftFGQPFh-Uo",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: kaggle in /usr/local/lib/python3.7/dist-packages (1.5.12)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from kaggle) (2.23.0)\n",
            "Requirement already satisfied: certifi in /usr/local/lib/python3.7/dist-packages (from kaggle) (2022.9.24)\n",
            "Requirement already satisfied: python-dateutil in /usr/local/lib/python3.7/dist-packages (from kaggle) (2.8.2)\n",
            "Requirement already satisfied: urllib3 in /usr/local/lib/python3.7/dist-packages (from kaggle) (1.24.3)\n",
            "Requirement already satisfied: six>=1.10 in /usr/local/lib/python3.7/dist-packages (from kaggle) (1.15.0)\n",
            "Requirement already satisfied: python-slugify in /usr/local/lib/python3.7/dist-packages (from kaggle) (6.1.2)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from kaggle) (4.64.1)\n",
            "Requirement already satisfied: text-unidecode>=1.3 in /usr/local/lib/python3.7/dist-packages (from python-slugify->kaggle) (1.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->kaggle) (2.10)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->kaggle) (3.0.4)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir -p ~/.kaggle\n",
        "!cp kaggle.json ~/.kaggle/"
      ],
      "metadata": {
        "id": "uJHXmkMYkbBw"
      },
      "id": "uJHXmkMYkbBw",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!ls -a ~"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "guBSHus6L7zT",
        "outputId": "ff88374e-0a54-4d04-8194-00ef5d62c890"
      },
      "id": "guBSHus6L7zT",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            ".   .bashrc  .config   .jupyter  .keras  .npm\t   .tmux.conf\n",
            "..  .cache   .ipython  .kaggle\t .local  .profile  .wget-hsts\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pwd"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "T1zFyAjVMC7A",
        "outputId": "64257231-f447-49c4-952c-8a4364d8fb43"
      },
      "id": "T1zFyAjVMC7A",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6c36e60c",
      "metadata": {
        "id": "6c36e60c",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "abfc67ed-00d2-4eeb-b3e4-a3218e916999"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Warning: Your Kaggle API key is readable by other users on this system! To fix this, you can run 'chmod 600 /root/.kaggle/kaggle.json'\n",
            "Downloading pokemonclassification.zip to /content\n",
            " 99% 411M/417M [00:02<00:00, 170MB/s]\n",
            "100% 417M/417M [00:02<00:00, 173MB/s]\n"
          ]
        }
      ],
      "source": [
        "!kaggle datasets download -d lantian773030/pokemonclassification"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!unzip -qq pokemonclassification.zip"
      ],
      "metadata": {
        "id": "ZnZ32IpXkPKk"
      },
      "id": "ZnZ32IpXkPKk",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3ec5937a",
      "metadata": {
        "id": "3ec5937a",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "39969983-c856-4be0-f9a9-a48231be850a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "total 608\n",
            "drwxr-xr-x 152 root root 4096 Oct 18 08:37 .\n",
            "drwxr-xr-x   1 root root 4096 Oct 18 08:37 ..\n",
            "drwxr-xr-x   2 root root 4096 Oct 18 08:37 Abra\n",
            "drwxr-xr-x   2 root root 4096 Oct 18 08:37 Aerodactyl\n",
            "drwxr-xr-x   2 root root 4096 Oct 18 08:37 Alakazam\n",
            "drwxr-xr-x   2 root root 4096 Oct 18 08:37 Alolan Sandslash\n",
            "drwxr-xr-x   2 root root 4096 Oct 18 08:37 Arbok\n",
            "drwxr-xr-x   2 root root 4096 Oct 18 08:37 Arcanine\n",
            "drwxr-xr-x   2 root root 4096 Oct 18 08:37 Articuno\n"
          ]
        }
      ],
      "source": [
        "!ls -la PokemonData/ | head"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5da4a06d",
      "metadata": {
        "id": "5da4a06d"
      },
      "source": [
        "Выпилим векторные картинки из датасета"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "85efc6c8",
      "metadata": {
        "id": "85efc6c8"
      },
      "outputs": [],
      "source": [
        "import glob\n",
        "import os\n",
        "\n",
        "bad_images = glob.glob('PokemonData/*/*.svg') \n",
        "#Return a possibly empty list of path names that match pathname 'PokemonData/*/*.svg'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0fc81faa",
      "metadata": {
        "id": "0fc81faa"
      },
      "outputs": [],
      "source": [
        "for bad_image in bad_images:\n",
        "    os.remove(bad_image)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3000da3d",
      "metadata": {
        "id": "3000da3d"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torchvision.transforms as T\n",
        "from torch.utils.data import Dataset, DataLoader"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9faec118",
      "metadata": {
        "id": "9faec118"
      },
      "source": [
        "Напишем свой класс датасета\n",
        "\n",
        "![](https://miro.medium.com/max/500/1*OWXgKpWZuD9jyclBdFHaZg.jpeg)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e8bf2af8",
      "metadata": {
        "id": "e8bf2af8"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import shutil\n",
        "from PIL import Image\n",
        "from sklearn.model_selection import train_test_split\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "\n",
        "class PokemonDataset(Dataset):\n",
        "    SPLIT_RANDOM_SEED = 42\n",
        "    TEST_SIZE = 0.25\n",
        "\n",
        "    def __init__(self, root, train=True, load_to_ram=True, transform=None):\n",
        "        super().__init__()\n",
        "        self.root = root\n",
        "        self.train = train\n",
        "        self.load_to_ram = load_to_ram\n",
        "        self.transform = transform\n",
        "        self.to_tensor = T.ToTensor()\n",
        "        self.all_files = []\n",
        "        self.all_labels = []\n",
        "        self.images = []\n",
        "\n",
        "        self.classes = sorted(os.listdir(self.root))\n",
        "        for i, class_name in tqdm(enumerate(self.classes), total=len(self.classes)):\n",
        "            files = sorted(os.listdir(os.path.join(self.root, class_name)))\n",
        "            train_files, test_files = train_test_split(files, random_state=self.SPLIT_RANDOM_SEED + i,\n",
        "                                                       test_size=self.TEST_SIZE)\n",
        "            if self.train:\n",
        "                self.all_files += train_files\n",
        "                self.all_labels += [i] * len(train_files)\n",
        "                if self.load_to_ram:\n",
        "                    self.images += self._load_images(train_files, i)\n",
        "\n",
        "            else:\n",
        "                self.all_files += test_files\n",
        "                self.all_labels += [i] * len(test_files)\n",
        "                if self.load_to_ram:\n",
        "                    self.images += self._load_images(test_files, i)\n",
        "\n",
        "    def _load_images(self, image_files, label):\n",
        "        images = []\n",
        "        for filename in image_files:\n",
        "            image = Image.open(os.path.join(self.root, self.classes[label], filename)).convert('RGB')\n",
        "            images += [image]\n",
        "\n",
        "        return images\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.all_files)\n",
        "\n",
        "    def __getitem__(self, item):\n",
        "        label = self.all_labels[item]\n",
        "        if self.load_to_ram:\n",
        "            image = self.images[item]\n",
        "        else:\n",
        "            filename = self.all_files[item]\n",
        "            image = Image.open(os.path.join(self.root, self.classes[label], filename)).convert('RGB')\n",
        "\n",
        "        if self.transform is not None:\n",
        "            image = self.transform(image)\n",
        "\n",
        "        return image, label"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "758f1aa7",
      "metadata": {
        "id": "758f1aa7"
      },
      "outputs": [],
      "source": [
        "train_dataset = PokemonDataset(root='PokemonData', train=True, load_to_ram=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d383ca98",
      "metadata": {
        "id": "d383ca98"
      },
      "outputs": [],
      "source": [
        "image, label = train_dataset[100]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "57c179c0",
      "metadata": {
        "id": "57c179c0"
      },
      "outputs": [],
      "source": [
        "image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0fe353e2",
      "metadata": {
        "id": "0fe353e2"
      },
      "outputs": [],
      "source": [
        "train_dataset.classes[label]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d97e3931",
      "metadata": {
        "id": "d97e3931"
      },
      "source": [
        "В следующей ячейке описаны обучающие (с аугментациями) и тестовые трансформы (без аугментаций) для картинок"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "af120685",
      "metadata": {
        "id": "af120685"
      },
      "outputs": [],
      "source": [
        "normalize = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
        "\n",
        "train_transform = T.Compose([\n",
        "    T.RandomResizedCrop(224, scale=(0.5, 1.0)),\n",
        "    T.RandomHorizontalFlip(),\n",
        "    T.ToTensor(),\n",
        "    normalize,\n",
        "])\n",
        "\n",
        "test_transform = T.Compose([\n",
        "    T.Resize(256),\n",
        "    T.CenterCrop(224),\n",
        "    T.ToTensor(),\n",
        "    normalize,\n",
        "])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a5aca187",
      "metadata": {
        "id": "a5aca187",
        "outputId": "72710eea-1d23-4ef8-89e1-41635a743555",
        "colab": {
          "referenced_widgets": [
            "8ee39b110185466abd7cd47705b442a6",
            "7d8dc1b992f74a02b37fb19de314f18e",
            "0461139a6e9648c98c0bb4d80c702d2e",
            "ae1583c600f34b5094ef6f6522765426",
            "06c096b05c5e4e80a07dc99eebeecd9b",
            "a8909c927380408eb0a81d78192ef329",
            "f42a63d1455e43998fca7fc90bb759e3",
            "4aace35f3b424fd0b46fe1c063e54c57",
            "c972632f93df4961b058cf03f33e599a",
            "aba4467a0c99442dac9ab7c240c0bf76",
            "24d43cf416434dedbd5482145f406252"
          ],
          "base_uri": "https://localhost:8080/",
          "height": 69
        }
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/150 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "8ee39b110185466abd7cd47705b442a6"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "train_dataset = PokemonDataset(root='PokemonData', train=True, load_to_ram=True, transform=test_transform)\n",
        "test_dataset = PokemonDataset(root='PokemonData', train=False, load_to_ram=True, transform=test_transform)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6141dd67",
      "metadata": {
        "id": "6141dd67"
      },
      "outputs": [],
      "source": [
        "print(f'Число классов: {len(train_dataset.classes)}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "10a24bcd",
      "metadata": {
        "id": "10a24bcd"
      },
      "outputs": [],
      "source": [
        "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, pin_memory=True, num_workers=1)\n",
        "test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, pin_memory=True, num_workers=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d317e8ee",
      "metadata": {
        "id": "d317e8ee"
      },
      "source": [
        "## 2. Обучаем модель из torchvision"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8e3c48d7",
      "metadata": {
        "id": "8e3c48d7"
      },
      "outputs": [],
      "source": [
        "import seaborn as sns\n",
        "import matplotlib.pyplot as plt\n",
        "from IPython.display import clear_output\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "\n",
        "sns.set_style('whitegrid')\n",
        "plt.rcParams.update({'font.size': 15})\n",
        "\n",
        "\n",
        "def plot_losses(train_losses, test_losses, train_accuracies, test_accuracies):\n",
        "    clear_output()\n",
        "    fig, axs = plt.subplots(1, 2, figsize=(13, 4))\n",
        "    axs[0].plot(range(1, len(train_losses) + 1), train_losses, label='train')\n",
        "    axs[0].plot(range(1, len(test_losses) + 1), test_losses, label='test')\n",
        "    axs[0].set_ylabel('loss')\n",
        "\n",
        "    axs[1].plot(range(1, len(train_accuracies) + 1), train_accuracies, label='train')\n",
        "    axs[1].plot(range(1, len(test_accuracies) + 1), test_accuracies, label='test')\n",
        "    axs[1].set_ylabel('accuracy')\n",
        "\n",
        "    for ax in axs:\n",
        "        ax.set_xlabel('epoch')\n",
        "        ax.legend()\n",
        "\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "87cef18b",
      "metadata": {
        "id": "87cef18b"
      },
      "outputs": [],
      "source": [
        "def training_epoch(model, optimizer, criterion, train_loader, tqdm_desc):\n",
        "    train_loss, train_accuracy = 0.0, 0.0\n",
        "    model.train()\n",
        "    for images, labels in tqdm(train_loader, desc=tqdm_desc):\n",
        "        images = images.to(device)  # images: batch_size x num_channels x height x width\n",
        "        labels = labels.to(device)  # labels: batch_size\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        logits = model(images)  # logits: batch_size x num_classes\n",
        "        loss = criterion(logits, labels)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        train_loss += loss.item() * images.shape[0]\n",
        "        train_accuracy += (logits.argmax(dim=1) == labels).sum().item()\n",
        "    \n",
        "    train_loss /= len(train_loader.dataset)\n",
        "    train_accuracy /= len(train_loader.dataset)\n",
        "    return train_loss, train_accuracy\n",
        "\n",
        "\n",
        "@torch.no_grad()\n",
        "def validation_epoch(model, criterion, test_loader, tqdm_desc):\n",
        "    test_loss, test_accuracy = 0.0, 0.0\n",
        "    model.eval()\n",
        "    for images, labels in tqdm(test_loader, desc=tqdm_desc):\n",
        "        images = images  # images: batch_size x num_channels x height x width\n",
        "        labels = labels  # labels: batch_size\n",
        "        logits = model(images)  # logits: batch_size x num_classes\n",
        "        loss = criterion(logits, labels)\n",
        "\n",
        "        test_loss += loss.item() * images.shape[0]\n",
        "        test_accuracy += (logits.argmax(dim=1) == labels).sum().item()\n",
        "\n",
        "    test_loss /= len(test_loader.dataset)\n",
        "    test_accuracy /= len(test_loader.dataset)\n",
        "    return test_loss, test_accuracy\n",
        "\n",
        "    \n",
        "def train(model, optimizer, scheduler, criterion, train_loader, test_loader, num_epochs):\n",
        "    train_losses, train_accuracies = [], []\n",
        "    test_losses, test_accuracies = [], []\n",
        "\n",
        "    for epoch in range(1, num_epochs + 1):\n",
        "        train_loss, train_accuracy = training_epoch(\n",
        "            model, optimizer, criterion, train_loader,\n",
        "            tqdm_desc=f'Training {epoch}/{num_epochs}'\n",
        "        )\n",
        "        test_loss, test_accuracy = validation_epoch(\n",
        "            model, criterion, test_loader,\n",
        "            tqdm_desc=f'Validating {epoch}/{num_epochs}'\n",
        "        )\n",
        "\n",
        "        if scheduler is not None: #better to keep it inside training_step\n",
        "            scheduler.step()\n",
        "\n",
        "        train_losses += [train_loss]\n",
        "        train_accuracies += [train_accuracy]\n",
        "        test_losses += [test_loss]\n",
        "        test_accuracies += [test_accuracy]\n",
        "        plot_losses(train_losses, test_losses, train_accuracies, test_accuracies)\n",
        "\n",
        "    return train_losses, test_losses, train_accuracies, test_accuracies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b5f1dedc",
      "metadata": {
        "id": "b5f1dedc"
      },
      "outputs": [],
      "source": [
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "device"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0f386173",
      "metadata": {
        "id": "0f386173"
      },
      "source": [
        "**Attention!** Интерфейс `torchvision.models` поменялся в версии **0.13**."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "device"
      ],
      "metadata": {
        "id": "2Ah01POwrO5T"
      },
      "id": "2Ah01POwrO5T",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0786ce92",
      "metadata": {
        "id": "0786ce92"
      },
      "outputs": [],
      "source": [
        "from torchvision.models import mobilenet_v2, MobileNet_V2_Weights\n",
        "\n",
        "num_epochs = 50\n",
        "model = mobilenet_v2(num_classes=len(train_dataset.classes)).to(device)\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)\n",
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, num_epochs)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3958e369",
      "metadata": {
        "id": "3958e369"
      },
      "source": [
        "Вот так можно посчитать число обучаемых параметров в модели:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "00574d5d",
      "metadata": {
        "id": "00574d5d"
      },
      "outputs": [],
      "source": [
        "sum(param.numel() for param in model.parameters())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "61179914",
      "metadata": {
        "scrolled": false,
        "id": "61179914"
      },
      "outputs": [],
      "source": [
        "train_losses, test_losses, train_accuracies, test_accuracies = train(\n",
        "    model, optimizer, scheduler, criterion, train_loader, test_loader, num_epochs\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def training_epoch(model, optimizer, criterion, train_loader, tqdm_desc):\n",
        "    train_loss, train_accuracy = 0.0, 0.0\n",
        "    model.train()\n",
        "    for images, labels in tqdm(train_loader, desc=tqdm_desc):\n",
        "        images = images.to(device)  # images: batch_size x num_channels x height x width\n",
        "        labels = labels.to(device)  # labels: batch_size\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        logits = model(images)  # logits: batch_size x num_classes\n",
        "        loss = criterion(logits, labels)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        train_loss += loss.item() * images.shape[0]\n",
        "        train_accuracy += (logits.argmax(dim=1) == labels).sum().item()\n",
        "    \n",
        "    train_loss /= len(train_loader.dataset)\n",
        "    train_accuracy /= len(train_loader.dataset)\n",
        "    return train_loss, train_accuracy\n",
        "\n",
        "\n",
        "@torch.no_grad()\n",
        "def validation_epoch(model, criterion, test_loader, tqdm_desc):\n",
        "    test_loss, test_accuracy = 0.0, 0.0\n",
        "    model.eval()\n",
        "    for images, labels in tqdm(test_loader, desc=tqdm_desc):\n",
        "        images = images.to(device)  # images: batch_size x num_channels x height x width\n",
        "        labels = labels.to(device)  # labels: batch_size\n",
        "        logits = model(images)  # logits: batch_size x num_classes\n",
        "        loss = criterion(logits, labels)\n",
        "\n",
        "        test_loss += loss.item() * images.shape[0]\n",
        "        test_accuracy += (logits.argmax(dim=1) == labels).sum().item()\n",
        "\n",
        "    test_loss /= len(test_loader.dataset)\n",
        "    test_accuracy /= len(test_loader.dataset)\n",
        "    return test_loss, test_accuracy\n",
        "\n",
        "    \n",
        "def train(model, optimizer, scheduler, criterion, train_loader, test_loader, num_epochs):\n",
        "    train_losses, train_accuracies = [], []\n",
        "    test_losses, test_accuracies = [], []\n",
        "\n",
        "    for epoch in range(1, num_epochs + 1):\n",
        "        train_loss, train_accuracy = training_epoch(\n",
        "            model, optimizer, criterion, train_loader,\n",
        "            tqdm_desc=f'Training {epoch}/{num_epochs}'\n",
        "        )\n",
        "        test_loss, test_accuracy = validation_epoch(\n",
        "            model, criterion, test_loader,\n",
        "            tqdm_desc=f'Validating {epoch}/{num_epochs}'\n",
        "        )\n",
        "\n",
        "        if scheduler is not None:\n",
        "            scheduler.step()\n",
        "\n",
        "        train_losses += [train_loss]\n",
        "        train_accuracies += [train_accuracy]\n",
        "        test_losses += [test_loss]\n",
        "        test_accuracies += [test_accuracy]\n",
        "        plot_losses(train_losses, test_losses, train_accuracies, test_accuracies)\n",
        "\n",
        "    return train_losses, test_losses, train_accuracies, test_accuracies"
      ],
      "metadata": {
        "id": "7tWIzIIWqOwj"
      },
      "id": "7tWIzIIWqOwj",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "id": "129b7f83",
      "metadata": {
        "id": "129b7f83"
      },
      "source": [
        "## 3. Добавляем аугментации"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bba8f370",
      "metadata": {
        "id": "bba8f370"
      },
      "outputs": [],
      "source": [
        "train_dataset.transform = train_transform\n",
        "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, pin_memory=True, num_workers=4)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ca8a21bf",
      "metadata": {
        "id": "ca8a21bf"
      },
      "outputs": [],
      "source": [
        "model = mobilenet_v2(num_classes=len(train_dataset.classes)).to(device)\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)\n",
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, num_epochs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b8ee8b08",
      "metadata": {
        "id": "b8ee8b08"
      },
      "outputs": [],
      "source": [
        "train_losses, test_losses, train_accuracies, test_accuracies = train(\n",
        "    model, optimizer, scheduler, criterion, train_loader, test_loader, num_epochs\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bbf1d115",
      "metadata": {
        "id": "bbf1d115"
      },
      "source": [
        "## 4. FIne-tuning предобученной модели"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "03b180d3",
      "metadata": {
        "id": "03b180d3"
      },
      "outputs": [],
      "source": [
        "model = mobilenet_v2(weights=MobileNet_V2_Weights.IMAGENET1K_V1)\n",
        "# Заменяем исходную голову модели на голову с правильным числом классов\n",
        "model.classifier[1] = torch.nn.Linear(1280, len(train_dataset.classes))\n",
        "model = model.to(device)\n",
        "\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)\n",
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, num_epochs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9ffa704b",
      "metadata": {
        "id": "9ffa704b"
      },
      "outputs": [],
      "source": [
        "train_losses, test_losses, train_accuracies, test_accuracies = train(\n",
        "    model, optimizer, scheduler, criterion, train_loader, test_loader, num_epochs\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0c98e2e8",
      "metadata": {
        "id": "0c98e2e8"
      },
      "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.8.13"
    },
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "d317e8ee",
        "129b7f83"
      ]
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "8ee39b110185466abd7cd47705b442a6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_7d8dc1b992f74a02b37fb19de314f18e",
              "IPY_MODEL_0461139a6e9648c98c0bb4d80c702d2e",
              "IPY_MODEL_ae1583c600f34b5094ef6f6522765426"
            ],
            "layout": "IPY_MODEL_06c096b05c5e4e80a07dc99eebeecd9b"
          }
        },
        "7d8dc1b992f74a02b37fb19de314f18e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a8909c927380408eb0a81d78192ef329",
            "placeholder": "​",
            "style": "IPY_MODEL_f42a63d1455e43998fca7fc90bb759e3",
            "value": " 39%"
          }
        },
        "0461139a6e9648c98c0bb4d80c702d2e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_4aace35f3b424fd0b46fe1c063e54c57",
            "max": 150,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_c972632f93df4961b058cf03f33e599a",
            "value": 58
          }
        },
        "ae1583c600f34b5094ef6f6522765426": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_aba4467a0c99442dac9ab7c240c0bf76",
            "placeholder": "​",
            "style": "IPY_MODEL_24d43cf416434dedbd5482145f406252",
            "value": " 58/150 [00:18&lt;00:49,  1.87it/s]"
          }
        },
        "06c096b05c5e4e80a07dc99eebeecd9b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a8909c927380408eb0a81d78192ef329": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f42a63d1455e43998fca7fc90bb759e3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "4aace35f3b424fd0b46fe1c063e54c57": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c972632f93df4961b058cf03f33e599a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "aba4467a0c99442dac9ab7c240c0bf76": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "24d43cf416434dedbd5482145f406252": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}