assignments/04_finetuning_augmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Свёрточные сети для классификации"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Type\n",
    "\n",
    "import torch\n",
    "from torch import Tensor, nn\n",
    "from torch.nn import functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Задание 1. Skip-connections (2 балла)\n",
    "\n",
    "Постройте архитектуру свёрточной сети, аналогичную архитектуре в примере ниже, но добавьте в неё skip-connections, то есть дополнительные рёбра в вычислительном графе, позволяющие пропускать градиент в более ранние слои напрямую, минуя очередной блок Conv2D + BatchNorm + ReLU:\n",
    "\n",
    "```python\n",
    "def forward(self, x: Tensor) -> Tensor:\n",
    "    x = x + self.block1(x)\n",
    "    x = self.maxpool(x)\n",
    "    x = x + self.block2(x)\n",
    "    x = self.maxpool(x)\n",
    "    ...\n",
    "    x = x.adaptive_maxpool(x).flatten(1)\n",
    "    logits = self.fc(x)\n",
    "    return logits\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Наша верхнеуровневая архитектура будет выглядеть так:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyResNet(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        block: Type[nn.Module],\n",
    "        n_classes: int,\n",
    "        hidden_channels: list[int] = [32, 64],\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "        # входной слой, принимающий изображение с 3-мя каналами\n",
    "        self.in_conv = nn.Conv2d(3, hidden_channels[0], kernel_size=3, stride=1)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "\n",
    "        # собираем свёрточные блоки, каждый задаётся кол-вом входных и выходных каналов\n",
    "        blocks = []\n",
    "        for c_in, c_out in zip(hidden_channels[:-1], hidden_channels[1:]):\n",
    "            # добавляем очередной блок\n",
    "            blocks.append(block(c_in, c_out))\n",
    "            # добавляем Max pooling для уменьшения размерности\n",
    "            blocks.append(nn.MaxPool2d(2, 2))\n",
    "\n",
    "        # собираем блоки в единый Sequential модуль для удобства\n",
    "        self.features = nn.Sequential(*blocks)\n",
    "        self.maxpool = nn.AdaptiveMaxPool2d(1)\n",
    "\n",
    "        # линейный слой для классификации\n",
    "        self.fc = nn.Linear(hidden_channels[-1], n_classes)\n",
    "\n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        h = self.features(self.relu(self.in_conv(x)))\n",
    "        logits = self.fc(self.maxpool(h).flatten(1))\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Базовый блок, без residual connections, состоит из двух свёрток и нормализаций:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BasicBlock(nn.Module):\n",
    "    def __init__(self, inplanes: int, planes: int) -> None:\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(\n",
    "            inplanes, planes, kernel_size=3, stride=1, padding=1, bias=False\n",
    "        )\n",
    "        self.bn1 = nn.BatchNorm2d(planes)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.conv2 = nn.Conv2d(\n",
    "            planes, planes, kernel_size=3, stride=1, padding=1, bias=False\n",
    "        )\n",
    "        self.bn2 = nn.BatchNorm2d(planes)\n",
    "\n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        # first conv + bn + nonlinearity\n",
    "        out = self.relu(self.bn1(self.conv1(x)))\n",
    "        # second conv + bn\n",
    "        out = self.bn2(self.conv2(out))\n",
    "        # final nonlinearity\n",
    "        out = self.relu(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Посмотрим на результат его применения к тензору:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 6, 32, 32])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BasicBlock(4, 6).forward(torch.randn(3, 4, 32, 32)).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Теперь нужно изменить этот блок, добавив в него skip-connection. Теперь в методе `forward` входной тензор `x` пойдёт по двум веткам:\n",
    "1. как в базовом блоке, через наши всёртки и нормализации, до последней нелинейности\n",
    "2. в обход свёрток и нормализаций\n",
    "\n",
    "В конце эти ветки нужно объединить через сумму. Тут есть проблема: в исходном тензоре `x` и обработанном нашим блоком `h(x)` отличается количество каналов (остальные размерности совпадают). То есть нам нужно сравнять количество каналов исходного тензора `inplanes` с количеством выходных каналов `outplanes`.\n",
    "\n",
    "Интуитивно, если рассматривать каждый пиксел входного тензора как вектор размера `inplanes`, в вектор размера `planes` его можно превратить домножением на матрицу размера `inplanes x planes`. Это можно сделать, создав свёрточный слой с размером кернела 1 - он и будет переводить наши пикселы в другую размерность.\n",
    "\n",
    "Не забудьте к сумме каналов применить нелинейность."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, inplanes: int, planes: int) -> None:\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(\n",
    "            inplanes, planes, kernel_size=3, stride=1, padding=1, bias=False\n",
    "        )\n",
    "        self.bn1 = nn.BatchNorm2d(planes)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.conv2 = nn.Conv2d(\n",
    "            planes, planes, kernel_size=3, stride=1, padding=1, bias=False\n",
    "        )\n",
    "        self.bn2 = nn.BatchNorm2d(planes)\n",
    "\n",
    "        # добавьте свёртку 1x1 для изменения кол-ва каналов входного тензора\n",
    "        ...\n",
    "\n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        # сохраним входной тензор на будущее\n",
    "        identity = x\n",
    "\n",
    "        # ВАШ ХОД\n",
    "        ...\n",
    "        out = self.relu(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Проверим размеры:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert ResidualBlock(4, 6).forward(torch.randn(3, 4, 32, 32)).shape == torch.Size(\n",
    "    [3, 6, 32, 32]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Проверим, что модель выдаёт тензор ожидаемого размера:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 7])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "MyResNet(ResidualBlock, 7, hidden_channels=[16, 32, 64, 128]).forward(\n",
    "    torch.randn(3, 3, 32, 32)\n",
    ").shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Теперь мы можем создавать модели разного размера, в том числе достаточно большие и глубокие, чтобы хорошо классифицировать изображения из датасета CIFAR-10."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "151047"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum(\n",
    "    p.numel()\n",
    "    for p in MyResNet(ResidualBlock, 7, hidden_channels=[16, 32, 64, 64]).parameters()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Задание 2. Обучение `MyResNet` с использованием Lightning (5 баллов)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ваша задача: добиться 80% точности на валидационной выборке с вашей реализацией `MyResNet`.\n",
    "\n",
    "После окончания обучения используйте метод `Trainer.validate` для вывода ваших метрик с удачного чекпоинта модели.\n",
    "\n",
    "NB: вызывайте `Trainer.validate` везде, где в задании требуется достичь какой-то точности\n",
    "\n",
    "\n",
    "Советы:\n",
    "- По умолчанию Lightning сохраняет только последний чекпоинт, так что вам может потребоваться `lightning.callbacks.ModelCheckpoint`, чтобы сохранять лучший чекпоинт в процессе обучения.\n",
    "\n",
    "- Используйте tensorboard, чтобы следить за динамикой обучения. Если заметите переобучение - подключайте регуляризацию. Большая модель с регуляризацией обычно лучше маленькой модели без неё.\n",
    "\n",
    "- Чтобы добиться нужной точности, ваша модель должна быть достаточно глубокой, ориентируйтесь на 4-5 блоков. Если необходимо, подключайте регуляризацию"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Задание 3. Добавление аугментаций (1 балл + 2 балла за точность на валидации более 85%)\n",
    "\n",
    "Добавьте к обучающему датасету аугментации - случайные трансформации входных данных. Для этого можно использовать `torchvision.transforms` и `albumentations`.\n",
    "\n",
    "С `torchvision.transforms` совсем просто: вам нужно будет при создании `Datamodule` из практики по `lightning` указать вместо\n",
    "\n",
    "```python\n",
    "transform = transforms.ToTensor()\n",
    "```\n",
    "композицию трансформаций:\n",
    "\n",
    "```python\n",
    "transform = transforms.Compose([\n",
    "    transforms.RandomHorizontalFlip(),  # случайное зеркальное отражение\n",
    "    ...\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В пакете `albumentations` аугментаций значительно больше:\n",
    "\n",
    "![albumentations](https://albumentations.ai/assets/img/custom/top_image.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Задание 4. Использование предобученной модели (4 балла)\n",
    "\n",
    "Теперь мы научимся использовать модели, обученные на других задачах\n",
    "\n",
    "Ваша задача: добиться 90% точности на тестовой выборке CIFAR-10. Постарайтесь уложиться модель с ~5 млн параметров"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В `torchvision.models` есть много реализованных архитектур, размером которых можно удобно управлять. Например, ниже можно создать крошечную версию модели `MobileNetV2`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "46322"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torchvision.models import MobileNetV2\n",
    "\n",
    "mobilenet = MobileNetV2(\n",
    "    num_classes=10,\n",
    "    width_mult=0.4,\n",
    "    inverted_residual_setting=[\n",
    "        # t, c, n, s\n",
    "        [1, 16, 1, 1],\n",
    "        [3, 24, 2, 2],\n",
    "        [3, 32, 3, 2],\n",
    "    ],\n",
    "    dropout=0.2,\n",
    ")\n",
    "\n",
    "sum([param.numel() for param in mobilenet.parameters()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Но кроме архитектуры модели, мы также можем скачать веса, полученные при обучении на каком-то датасете. Например, для нашей задачи можно использовать предобучение на самом известном датасете для классификации изображений - ImageNet:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5288548"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torchvision.models.efficientnet import EfficientNet_B0_Weights, efficientnet_b0\n",
    "\n",
    "# создаём EfficientNet с весами, полученными на ImageNet\n",
    "weights = EfficientNet_B0_Weights.IMAGENET1K_V1\n",
    "efficientnet = efficientnet_b0(weights=weights)\n",
    "sum([param.numel() for param in efficientnet.parameters()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Указание 1.** С использованием модели в исходном виде есть проблема: в ImageNet 1000 классов, а у нас только 10. Поэтому в предобученной модели нужно будет полностью заменить последний линейный слой, который даёт распределение вероятностей классов. Это можно сделать уже в готовом объекте модели, переназначив атрибут.\n",
    "\n",
    "Подсказка: в `efficientnet_b0` линейный слой находится в атрибуте `classifier` \n",
    "\n",
    "\n",
    "**Указание 2.** Все слои, кроме нескольких последних (может быть, только последнего) мы можем заморозить, то есть сделать значения параметров в них неизменными. Это позволит и сохранить способность модели выделять полезные низкоуровневые признаки (она научилась этому на ImageNet), и существенно ускорить дообучение.\n",
    "\n",
    "\n",
    "Чтобы заморозить параметры, нужно всего лишь отключить для них расчёт градиентов. Вернитесь к первой практике, чтобы вспомнить, как это можно сделать. Нам подойдёт самый простой способ с `.requires_grad`.\n",
    "\n",
    "Подсказка: в `efficientnet_b0` свёрточные слои находятся в атрибуте `features` "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Указание 3.** Предобученные модели на ImageNet ожидают специальным образом трансформированные изображения:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ImageClassification(\n",
       "    crop_size=[224]\n",
       "    resize_size=[256]\n",
       "    mean=[0.485, 0.456, 0.406]\n",
       "    std=[0.229, 0.224, 0.225]\n",
       "    interpolation=InterpolationMode.BICUBIC\n",
       ")"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weights.transforms()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Поэтому эти трансформации нужно будет передать в датамодуль (как мы делали с аугментациями)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ВАШ ХОД: Обучите модель и выведите результат метода validate на удачном чекпоинте"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dl-mcs",
   "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}