{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.utils.flop_counter import FlopCounterMode"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "print(torch.__version__)",
   "id": "5172f205b1827513",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Как и в прошлой лабораторной работе, проверьте, что у torch есть доступ к графическому ускорителю и драйверу cuda.",
   "id": "73e641775babe64"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Датасет Tiny ImageNet и свёрточные сети",
   "id": "dc87e5e630d86c22"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "Рассмотрим небольшой срез датасета `https://huggingface.co/datasets/zh-plus/tiny-imagenet`. В нем выбраны только 11 классов достаточно близких друг к другу: кошки и собаки. Описание каждого из классов можете найти в файле `data_info.json`. Датасет содержится в двух срезах `train.pt` и `test.pt`. Каждый из файлов имеет структуру\n",
    "\n",
    "```\n",
    "{\n",
    "    \"images\": images[],\n",
    "    \"labels\": labels[],\n",
    "}\n",
    "```\n",
    "\n",
    "Изображения хранятся в формате `NCHW` с типом `uint8` (для экономии места), а метки классов в формате `N`."
   ],
   "id": "d17e5d67be5e12e4"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Разберите написанный класс-наследник `Dataset`. Зачем здесь перегружены методы `__len__` и `__getitem__`? Где происходит переход от типа `uint8` к `float`? Зачем?",
   "id": "b81478eb0fbb0166"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "class ImageDataset(Dataset):\n",
    "    def __init__(self, path, transform=None):\n",
    "        data = torch.load(path, weights_only=True)\n",
    "\n",
    "        self.images = data[\"images\"]\n",
    "        self.labels = data[\"labels\"]\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.labels)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        image = self.images[idx]\n",
    "        label = self.labels[idx]\n",
    "        image = image.float() / 255.0\n",
    "\n",
    "        if self.transform is not None:\n",
    "            image = self.transform(image)\n",
    "\n",
    "        return image, label"
   ],
   "id": "4ab5dc6ac3f3497d",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Теперь определим непосредственно наши датасеты с `batch_size=64` (можете изменять по своему усмотрению). Чем отличается определение тренировочного лоадера от тестового? Почему так? Какая трансформация, присутствовавшая в прошлой лабораторной работе пропущена? Добавьте ее.",
   "id": "d2ae077e2c5fdeaa"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "train_dataset = ImageDataset(\"train.pt\")\n",
    "test_dataset = ImageDataset(\"test.pt\")\n",
    "\n",
    "train_loader = DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=64,\n",
    "    shuffle=True,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    ")\n",
    "\n",
    "test_loader = DataLoader(\n",
    "    test_dataset,\n",
    "    batch_size=64,\n",
    "    shuffle=False,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    ")"
   ],
   "id": "11d1ab4c83b548d4",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Проверьте, что датасет загружен правильно. С помощью лоадера получите экземпляр батча и посмотрите его размерность. Проверьте, что элементы матрицы действительно `float` и находятся в диапазоне от 0 до 1.",
   "id": "313154a8949add54"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "8407b50c5e0570fa",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Визуализируйте батч изображений, выведите его на экран с помощью `matplotlib` вместе с метками классов.",
   "id": "16877769f5e6710c"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "2cd8ce9d1016b515",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Определите простую свёрточную нейронную сеть (не более 100_000) параметров. За основу берите архитектуру VGG, LeNet или AlexNet",
   "id": "793a74719a06f5ef"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "class SimpleCNN(nn.Module):\n",
    "    ..."
   ],
   "id": "1eb2a08cf94cc0c5",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Определите экземпляр модели и перенесите ее веса на GPU",
   "id": "5c401803fbb84efd"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "num_classes = len(torch.unique(train_dataset.labels))\n",
    "\n",
    "print(\"Classes:\", num_classes)\n",
    "\n",
    "device = torch.device(\n",
    "    \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    ")\n",
    "\n",
    "model = SimpleCNN(num_classes).to(device)\n",
    "\n",
    "print(model)"
   ],
   "id": "2c0ff8e9d08dab38",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Определите функцию потерь (рекомендуется перекрестная энтропия) и оптимизатор (рекомендуется Adam с weight decay)",
   "id": "32e8699983f02fe8"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "f2b68f2ac2b2ead0",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Определите обучающий цикл, ориентируясь на предыдущую лабораторную работу. Также, как и там, сохраняйте статистику",
   "id": "894920a769a30463"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "max_epochs = 50\n",
    "\n",
    "statistics = {\n",
    "    \"train_loss\": [],\n",
    "    \"train_acc\": [],\n",
    "    \"val_loss\": [],\n",
    "    \"val_acc\": []\n",
    "}\n",
    "\n",
    "\n",
    "for epoch in range(1, max_epochs+1):\n",
    "\n",
    "    # TRAIN\n",
    "    model.train()\n",
    "\n",
    "\n",
    "    # EVAL\n",
    "\n",
    "    print(f\"Epoch {epoch:3d}/{max_epochs} │ Train Loss: {statistics['train_loss'][-1]:.4f}  Acc: {statistics['train_acc'][-1]:.4f} │ Val Loss: {statistics['val_loss'][-1]:.4f}  Acc: {statistics['val_acc'][-1]:.4f}\")"
   ],
   "id": "f11d2898a89c5f33",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Сделайте визуализацию значений функции потерь, а также точности на тренировочных и валидационных данных с помощью `matplotlib`. Найдите точку ранней остановки (если она есть)",
   "id": "878e921ece3933bb"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "a282e8d520842c39",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "Добавьте в свою нейросеть одну из технологий, помогших решить задачу ImageNet. Возьмите свой номер в группе и вычислите остаток от деления на 3.\n",
    "\n",
    "0. Добавьте технологию residual block\n",
    "1. Добавьте технологию inception\n",
    "2. Добавьте SE block\n",
    "\n",
    "С модифицированной нейронной сетью проведите тот же эксперимент, что и выше, и сравните результаты (выведите их на один график для визуализации)"
   ],
   "id": "8b4328b4e679edc2"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "58beacab4766c8b9",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Визуализируйте матрицу ошибок модели, сопоставьте с названиями классов. Какие выводы можно сделать?",
   "id": "c7d1e9143c799a77"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "292839084c52a175",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "**Соревнование** (5 баллов). Модифицируйте свою нейросеть, а также трансформации в датасете, чтобы получить наилучший результат при наименьших вычислительных потребностях. Метрика будет вычисляться в зависимости не от количества параметров, а от вычислительной сложности в GFLOP, то есть количестве операций с плавающей точкой, которые нужно произвести для одного изображения. Метрика, по которой будут оцениваться работы в ячейке ниже. Лучшие три результата приносят 5, 4 и 3 балла соответственно.",
   "id": "f0817f29086cf708"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "INPUT_SHAPE = (1, 3, 64, 64)\n",
    "F0 = 1e8\n",
    "\n",
    "\n",
    "def competition_score(model, accuracy: float) -> float:\n",
    "    model.eval()\n",
    "\n",
    "    x = torch.zeros(INPUT_SHAPE, device=device)\n",
    "\n",
    "    with torch.inference_mode():\n",
    "        with FlopCounterMode(display=False) as counter:\n",
    "            logits = model(x)\n",
    "\n",
    "    flops = counter.get_total_flops()\n",
    "    num_classes = logits.shape[1]\n",
    "\n",
    "    adjusted_accuracy = max(\n",
    "        0.0, (accuracy - 1 / num_classes) / (1 - 1 / num_classes)\n",
    "    )\n",
    "    penalty = (F0 / max(flops, F0)) ** 0.05\n",
    "    return adjusted_accuracy * penalty"
   ],
   "id": "4cd181b10f0d8cba",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Пояснение к функции для score. За каждое удвоение сложности, результат понижается примерно на 3.5%",
   "id": "31b6b73fe7a8d093"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "cca2f3908fa5c6cc",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}