assignments/03_init_act_norm_optim.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Инициализация и нормализация"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В этом задании вам предстоит реализовать два вида нормализации: по батчам (BatchNorm1d) и по признакам (LayerNorm1d)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Callable, NamedTuple\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch import Tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Реализация BatchNorm1d и LayerNorm1d."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.1. (2 балла) Реализуйте BatchNorm1d\n",
    "\n",
    "Подсказка: чтобы хранить текущие значения среднего и дисперсии, вам потребуется метод `torch.nn.Module.register_buffer`, ознакомьтесь с документацией к нему. Подумайте, какие проблемы возникнут, если вы будете просто сохранять ваши значения в тензор.\n",
    "\n",
    "\n",
    "Важно помнить:\n",
    "- Понятно, что нормализацию мы добавляем после очередного слоя с параметрами, но до применения функции активации или после? Подумайте, есть ли у одного из этих способов преимущества над другим.\n",
    "- Модуль нормализации по батчам работает по-разному при обучении и при валидации, и ему нужно понимать, в каком он состоянии. Эта информация доступна в атрибуте модуля `self.training: bool`, его значение определит ветвление логики в вашей реализации метода `forward`.\n",
    "- Переключение модулей между режимами осуществляется вызовами у объекта модели методов `.train()` (переключение в режим обучения) и `.eval()` (переключение в режим валидации) Почитайте документацию к этим методам. Ваши функции `train_epoch` и `test_epoch` теперь должны переводить модель в нужный режим перед началом обработки данных."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BatchNorm1d(nn.Module):\n",
    "    def __init__(\n",
    "        self, num_features: int, momentum: float = 0.9, eps: float = 1e-5\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "        self.scale = nn.Parameter(...)\n",
    "        self.shift = nn.Parameter(...)\n",
    "        self.register_buffer(\"running_mean\", ...)\n",
    "        self.register_buffer(\"running_var\", ...)\n",
    "        self.momentum = momentum\n",
    "        self.eps = eps\n",
    "\n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.2. (1 балл) Реализуйте LayerNorm1d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Отличия LayerNorm от BatchNorm - в том, что расчёт средних и дисперсий в BatchNorm происходит вдоль размерности батча (см. рисунок слева), а в LayerNorm - вдоль размерности признаков (см. рисунок справа).\n",
    "\n",
    "<img src=\"../attachments/norm.png\" width=\"800\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm1d(nn.Module):\n",
    "    def __init__(self, num_features: int, eps: float = 1e-5) -> None:\n",
    "        super(LayerNorm1d, self).__init__()\n",
    "        self.scale = nn.Parameter(...)\n",
    "        self.shift = nn.Parameter(...)\n",
    "        self.eps = eps\n",
    "\n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Замечание: явных тестов на корректность в этом задании нет, так как конкретные реализации могут давать немного разные результаты. Но вы можете проверить корректность своей реализации в эксперименте сами, нормализация должна немного исправлять проблемы неудачной инициализации. Ну и как минимум вы можете вручную проверить, что ваши активации действительно нормализуются в результате применения вашего модуля."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Эксперименты"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В этом задании ваша задача - проверить, какие из приёмов хорошо справляются с нездоровыми активациями в промежуточных слоях. Вам будет дана базовая модель, у которой есть проблемы с инициализацией параметров, попробуйте несколько приёмов для устранения проблем обучения:\n",
    "1. Хорошая инициализация параметров\n",
    "2. Ненасыщаемая функция активации (например, `F.leaky_relu`)\n",
    "3. Нормализация по батчам или по признакам (можно использовать встроенные `nn.BatchNorm1d` и `nn.LayerNorm`)\n",
    "4. Более продвинутый оптимизатор (`torch.optim.RMSprop`)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.0. Подготовка: датасет, функции для обучения"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Проверять наши гипотезы будем на датасете MNIST, для отладки добавим в функции для обучения возможность использовать только несколько батчей данных"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "train_dataset = datasets.MNIST(\n",
    "    \"data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=transforms.ToTensor(),\n",
    ")\n",
    "test_dataset = datasets.MNIST(\n",
    "    \"data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=transforms.ToTensor(),\n",
    ")\n",
    "\n",
    "batch_size = 32\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def training_step(\n",
    "    batch: tuple[torch.Tensor, torch.Tensor],\n",
    "    model: nn.Module,\n",
    "    optimizer: torch.optim.Optimizer,\n",
    ") -> torch.Tensor:\n",
    "    # прогоняем батч через модель\n",
    "    x, y = batch\n",
    "    logits = model(x)\n",
    "    # оцениваем значение ошибки\n",
    "    loss = F.cross_entropy(logits, y)\n",
    "    # обновляем параметры\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()\n",
    "    # возвращаем значение функции ошибки для логирования\n",
    "    return loss\n",
    "\n",
    "\n",
    "def train_epoch(\n",
    "    dataloader: DataLoader,\n",
    "    model: nn.Module,\n",
    "    optimizer: torch.optim.Optimizer,\n",
    "    max_batches: int = 100,\n",
    ") -> Tensor:\n",
    "    loss_values: list[float] = []\n",
    "    for i, batch in enumerate(dataloader):\n",
    "        loss = training_step(batch, model, optimizer)\n",
    "        loss_values.append(loss.item())\n",
    "        if i == max_batches:\n",
    "            break\n",
    "    return torch.tensor(loss_values).mean()\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def test_epoch(\n",
    "    dataloader: DataLoader, model: nn.Module, max_batches: int = 100\n",
    ") -> Tensor:\n",
    "    loss_values: list[float] = []\n",
    "    for i, batch in enumerate(dataloader):\n",
    "        x, y = batch\n",
    "        logits = model(x)\n",
    "        # оцениваем значение ошибки\n",
    "        loss = F.cross_entropy(logits, y)\n",
    "        loss_values.append(loss.item())\n",
    "        if i == max_batches:\n",
    "            break\n",
    "    return torch.tensor(loss_values).mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.1. Определение класса модели (2 балла)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Для удобства проведения экспериментов мы немного усложним создание модели, чтобы можно было задать разные способы инициализации параметров и нормализации промежуточных активаций, не меняя определение класса.\n",
    "\n",
    "Добавьте в метод `__init__`:\n",
    "- аргумент, который позволит использовать разные функции активации для промежуточных слоёв\n",
    "- аргумент, который позволит задавать разные способы нормализации: `None` (без нормализации), `nn.BatchNorm` и `nn.LayerNorm`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_std_normal(model: nn.Module) -> None:\n",
    "    \"\"\"Функция для инициализации параметров модели стандартным нормальным распределением.\"\"\"\n",
    "    for param in model.parameters():\n",
    "        torch.nn.init.normal_(param.data, mean=0, std=1)\n",
    "\n",
    "\n",
    "from typing import Type\n",
    "\n",
    "\n",
    "class MLP(nn.Module):\n",
    "    \"\"\"Базовая модель для экспериментов\n",
    "\n",
    "    Args:\n",
    "        input_dim (int): размерность входных признаков\n",
    "        hidden_dim (int): размерност скрытого слоя\n",
    "        output_dim (int): кол-во классов\n",
    "        act_fn (Callable[[Tensor], Tensor], optional): Функция активации. Defaults to F.tanh.\n",
    "        init_fn (Callable[[nn.Module], None], optional): Функция для инициализации. Defaults to init_std_normal.\n",
    "        norm (Type[nn.BatchNorm1d  |  nn.LayerNorm] | None, optional): Способ нормализации промежуточных активаций.\n",
    "            Defaults to None.\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_dim: int,\n",
    "        hidden_dim: int,\n",
    "        output_dim: int,\n",
    "        act_fn: Callable[[Tensor], Tensor] = F.tanh,\n",
    "        init_fn: Callable[[nn.Module], None] = init_std_normal,\n",
    "        norm: Type[nn.BatchNorm1d | nn.LayerNorm] | None = None,\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "        # теперь линейные слои будем задавать\n",
    "        self.fc1 = nn.Linear(input_dim, hidden_dim)\n",
    "        self.fc2 = nn.Linear(hidden_dim, output_dim)\n",
    "        self.act_fn = ...\n",
    "        self.norm = ...\n",
    "\n",
    "        # reinitialize parameters\n",
    "        init_fn(self)\n",
    "\n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        h = self.fc1.forward(x.flatten(1))\n",
    "        # here you can do normalization\n",
    "        if self.norm:\n",
    "            h = self.norm(h)\n",
    "        return self.fc2.forward(self.act_fn(h))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.2. Эксперименты (7 баллов)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Проведите по 3 эксперимента с каждой из модификаций с разными значениями `seed`, соберите статистику значений тестовой ошибки после 10 эпох обучения, сделайте выводы о том, что работает лучше\n",
    "\n",
    "Проверяем:\n",
    "1. Метод инициализации весов модели: $\\mathcal{N}(0, 1)$ / Kaiming normal\n",
    "2. Функция активации: tanh /  (или любая другая без насыщения)\n",
    "3. Слой нормализации: None / BatchNorm / LayerNorm\n",
    "4. Выбранный оптимизатор: SGD / RMSprop / Adam\n",
    "\n",
    "Итого у нас 7 экспериментов:\n",
    "- исходный (1)\n",
    "- смена инициализации (1)\n",
    "- смена нелинейности (1)\n",
    "- смена нормализации (2)\n",
    "- смена оптимизатора (2)\n",
    "\n",
    "Каждый эксперимент нужно повторить 3 раза с разными значениями random seed, посчитать среднее и вывести результаты в pandas.DataFrame.\n",
    "Можно дополнительно потестировать разные сочетания опций, например инициализация + нормализация\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Чтобы автоматизировать проведение экспериментов, можно использовать функцию, которая будет принимать все необходимые настройки эксперимента, запускать его и сохранять нужные метрики:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_experiment(\n",
    "    model_gen: Callable[[], nn.Module],\n",
    "    optim_gen: Callable[[nn.Module], torch.optim.Optimizer],\n",
    "    seed: int,\n",
    "    n_epochs: int = 10,\n",
    "    max_batches: int | None = None,\n",
    "    verbose: bool = False,\n",
    ") -> float:\n",
    "    \"\"\"Функция для запуска экспериментов.\n",
    "\n",
    "    Args:\n",
    "        model_gen (Callable[[], nn.Module]): Функция для создания модели\n",
    "        optim_gen (Callable[[nn.Module], torch.optim.Optimizer]): Функция для создания оптимизатора для модели\n",
    "        seed (int): random seed\n",
    "        n_epochs (int, optional): Число эпох обучения. Defaults to 10.\n",
    "        max_batches (int | None, optional): Если указано, только `max_batches` минибатчей\n",
    "            будет использоваться при обучении и тестировании. Defaults to None.\n",
    "        verbose (bool, optional): Выводить ли информацию для отладки. Defaults to False.\n",
    "\n",
    "    Returns:\n",
    "        float: Значение ошибки на тестовой выборке в конце обучения\n",
    "    \"\"\"\n",
    "    torch.manual_seed(seed)\n",
    "    # создадим модель и выведем значение ошибки после инициализации\n",
    "    model = model_gen()\n",
    "    optim = optim_gen(model)\n",
    "    epoch_losses: list[float] = []\n",
    "    for i in range(n_epochs):\n",
    "        train_loss = train_epoch(train_loader, model, optim, max_batches=max_batches)\n",
    "        test_loss = test_epoch(test_loader, model, max_batches=max_batches)\n",
    "        if verbose:\n",
    "            print(f\"Epoch {i} train loss = {train_loss:.4f}\")\n",
    "            print(f\"Epoch {i} test loss = {test_loss:.4f}\")\n",
    "\n",
    "        epoch_losses.append(test_loss.item())\n",
    "\n",
    "    last_epoch_loss = epoch_losses[-1]\n",
    "    return last_epoch_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Пример использования:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0 train loss = 12.6168\n",
      "Epoch 0 test loss = 9.9327\n",
      "Epoch 1 train loss = 9.0954\n",
      "Epoch 1 test loss = 7.5498\n",
      "Epoch 2 train loss = 6.9607\n",
      "Epoch 2 test loss = 6.2342\n",
      "Epoch 3 train loss = 5.8992\n",
      "Epoch 3 test loss = 5.3655\n",
      "Epoch 4 train loss = 4.9951\n",
      "Epoch 4 test loss = 4.7433\n",
      "Epoch 5 train loss = 4.4778\n",
      "Epoch 5 test loss = 4.3001\n",
      "Epoch 6 train loss = 3.9693\n",
      "Epoch 6 test loss = 3.9605\n",
      "Epoch 7 train loss = 3.7261\n",
      "Epoch 7 test loss = 3.6844\n",
      "Epoch 8 train loss = 3.4223\n",
      "Epoch 8 test loss = 3.4538\n",
      "Epoch 9 train loss = 2.9975\n",
      "Epoch 9 test loss = 3.2638\n"
     ]
    }
   ],
   "source": [
    "losses = run_experiment(\n",
    "    model_gen=lambda: MLP(784, 128, 10, init_fn=init_std_normal, norm=None),\n",
    "    optim_gen=lambda x: torch.optim.SGD(x.parameters(), lr=0.01),\n",
    "    seed=42,\n",
    "    n_epochs=10,\n",
    "    max_batches=100,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Для удобства задания настроек эксперимента можно определять их с помощью класса `Experiment`, в котором можно также реализовать логику для строкового представления:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_dim = 784\n",
    "hidden_dim = 128\n",
    "output_dim = len(train_dataset.classes)\n",
    "\n",
    "\n",
    "class Experiment(NamedTuple):\n",
    "    init_fn: Callable[[nn.Module], None]\n",
    "    act_fn: Callable[[Tensor], Tensor]\n",
    "    norm: Type[nn.BatchNorm1d | nn.LayerNorm] | None\n",
    "    optim_cls: Type[torch.optim.Optimizer]\n",
    "\n",
    "    @property\n",
    "    def model_gen(self) -> Callable[[], nn.Module]:\n",
    "        return lambda: MLP(\n",
    "            input_dim, hidden_dim, output_dim, init_fn=self.init_fn, norm=self.norm\n",
    "        )\n",
    "\n",
    "    @property\n",
    "    def optim_gen(self) -> Callable[[nn.Module], torch.optim.Optimizer]:\n",
    "        return lambda x: self.optim_cls(x.parameters(), lr=0.01)\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        # TODO: попробуйте сделать представление эксперимента более читаемым\n",
    "        return str(self)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Описываем все эксперименты:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[init: init_std_normal, act: tanh, norm: null, optim: SGD,\n",
       " init: init_std_normal, act: silu, norm: LayerNorm, optim: SGD,\n",
       " init: init_std_normal, act: relu, norm: BatchNorm1d, optim: RMSprop]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "options = [\n",
    "    Experiment(\n",
    "        init_fn=init_std_normal,\n",
    "        act_fn=F.tanh,\n",
    "        norm=None,\n",
    "        optim_cls=torch.optim.SGD,\n",
    "    ),\n",
    "    Experiment(\n",
    "        init_fn=init_std_normal,\n",
    "        act_fn=F.silu,\n",
    "        norm=nn.LayerNorm,\n",
    "        optim_cls=torch.optim.SGD,\n",
    "    ),\n",
    "    Experiment(\n",
    "        init_fn=init_std_normal,\n",
    "        act_fn=F.relu,\n",
    "        norm=nn.BatchNorm1d,\n",
    "        optim_cls=torch.optim.RMSprop,\n",
    "    ),\n",
    "]\n",
    "\n",
    "options"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Запускаем расчёты:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "init: init_std_normal, act: tanh, norm: null, optim: SGD\n",
      "init: init_std_normal, act: silu, norm: LayerNorm, optim: SGD\n",
      "init: init_std_normal, act: relu, norm: BatchNorm1d, optim: RMSprop\n"
     ]
    }
   ],
   "source": [
    "seeds = [42]  # здесь вам нужно 3 разных значения\n",
    "results = []\n",
    "\n",
    "for option in options:\n",
    "    print(option)\n",
    "    for seed in seeds:\n",
    "        loss = run_experiment(\n",
    "            model_gen=...,\n",
    "            optim_gen=...,\n",
    "            seed=seed,\n",
    "            n_epochs=10,\n",
    "            max_batches=None,\n",
    "            verbose=False,\n",
    "        )\n",
    "        results.append([str(option), seed, loss])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Выводим результаты:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>loss</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>desc</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>init: init_std_normal, act: relu, norm: BatchNorm1d, optim: RMSprop</th>\n",
       "      <td>0.443120</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>init: init_std_normal, act: silu, norm: LayerNorm, optim: SGD</th>\n",
       "      <td>1.821582</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>init: init_std_normal, act: tanh, norm: null, optim: SGD</th>\n",
       "      <td>3.263790</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                        loss\n",
       "desc                                                        \n",
       "init: init_std_normal, act: relu, norm: BatchNo...  0.443120\n",
       "init: init_std_normal, act: silu, norm: LayerNo...  1.821582\n",
       "init: init_std_normal, act: tanh, norm: null, o...  3.263790"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "pd.DataFrame(...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ВЫВОДЫ:"
   ]
  }
 ],
 "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
}