week11_diffprogramming/homework.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.5"
  },
  "colab": {
   "name": "DL21_fall_shw9.ipynb",
   "provenance": [],
   "collapsed_sections": []
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dr8Bk7tpUCbh"
   },
   "source": [
    "# Дифференцируемое программирование\n",
    "\n",
    "**Разработчик: Александр Шевченко**\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mryab/dl-hse-ami/blob/main/week11_diffprogramming/homework.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Ucdyy1Hw-Vjr"
   },
   "source": [
    "!wget -q \"https://raw.githubusercontent.com/mryab/dl-hse-ami/main/week11_diffprogramming/pairwise_example.pth\"\n",
    "!wget -q \"https://raw.githubusercontent.com/mryab/dl-hse-ami/main/week11_diffprogramming/state_dict.pth\"\n",
    "!wget -q \"https://raw.githubusercontent.com/mryab/dl-hse-ami/main/week11_diffprogramming/Structure-of-linear-chain-conditional-random-field.png\"\n",
    "!wget -q \"https://raw.githubusercontent.com/mryab/dl-hse-ami/main/week11_diffprogramming/unary_example.pth\"\n",
    "!wget -q \"https://raw.githubusercontent.com/mryab/dl-hse-ami/main/week11_diffprogramming/utils.py\""
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aRMYMXWVUCbh"
   },
   "source": [
    "## Дифференцируемое программирование\n",
    "\n",
    "В этом задании будет реализовываться система распознавания рукописных слов, основанная на совмещении алгоритмов предсказания (динамическое программирование) и глубинного обучения. Мы будем использовать датасет Stanford OCR (http://ai.stanford.edu/~btaskar/ocr/), состоящий из слов на английском языке и изображений рукописных букв.\n",
    "\n",
    "Для начала загрузим и подготовим данные. Для распаковки необходим gunzip. Пользователям Windows нужно скачать и распаковать датасет вручную."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "IEcSyPFXUCbh"
   },
   "source": [
    "!rm -rf letter.data\n",
    "!wget http://ai.stanford.edu/~btaskar/ocr/letter.data.gz\n",
    "!gunzip letter.data.gz"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "DPBifp6zUCbi"
   },
   "source": [
    "from utils import prepare_data\n",
    "train_x, train_y, test_x, test_y, val_x, val_y, le = prepare_data()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X19HQNDGUCbi"
   },
   "source": [
    "Каждый элемент датасета содержит данные об одном слове. Списки $*\\_x[i]$ содержат numpy-массивы размера [word_len, 1, 32, 32], содержащие изображения рукописных букв. Списки $*\\_y[i]$ содержат numpy-массивы размера [word_len] с метками для каждого изображения."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "k3D7aHuKUCbi"
   },
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sb\n",
    "import numpy as np\n",
    "%matplotlib inline"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VWj_ClToUCbi"
   },
   "source": [
    "Изображения выглядят следующим образом. Метки классов уже сконвертированы в числа для использования.\n",
    "\n",
    "Обратите внимание, что в нашем датасете первые буквы в каждом слове обрезаны и не используются (это не баг, а сделано специально, потому что первая буква часто бывает заглавной и, соответственно, её вариабельность сильно выше)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "oNGo6y4HUCbi"
   },
   "source": [
    "sb.set()\n",
    "fig, ax = plt.subplots(1, train_x[0].shape[0], figsize=(15,15))\n",
    "ax = np.array(ax)\n",
    "\n",
    "word = ''.join(le.inverse_transform(train_y[0]))\n",
    "for idx in range(train_x[0].shape[0]):\n",
    "    ax[idx].set_title(word[idx])\n",
    "    ax[idx].axis('off')\n",
    "    ax[idx].imshow(train_x[0][idx,0,:,:])\n",
    "plt.tight_layout()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zmZvbkBcUCbi"
   },
   "source": [
    "### Score функция и правдоподобие"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lhzvx7RQUCbi"
   },
   "source": [
    "<img src=\"Structure-of-linear-chain-conditional-random-field.png\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9OkDbH0WUCbi"
   },
   "source": [
    "Мы будем использовать модель вида цепочка (то есть нас будут интересовать только связи между соседними буквами) со следующей score-функцией:\n",
    "$$\n",
    "F(Y| X, \\Theta) = \\sum\\limits_{i=0}^{L-1} U(x_i, y_i) + \\sum\\limits_{i=0}^{L-2} W(y_{i}, y_{i+1})\n",
    "$$\n",
    "\n",
    "$\\Theta$ содержит параметры унарных $U$ и парных $W$ потенциалов. В этом задании для унарных потенциалов мы будем использовать простую нейросеть для классификации изображений (как для MNIST), а парные потенциалы параметризуем при помощи матрицы размера $26 \\times 26$ (обратите внимание, что парные потенциалы зависят только от меток, т.е., не зависят непосредственно от изображений).\n",
    "$U$ унарные потенциалы отвечают за совместимость метки $y_i$ и входного изображения буквы $x_i$. Парные потенциалы показывают, насколько вероятно сочетание букв $(y_i, y_{i+1})$.\n",
    "\n",
    "Используя score-функцию $F$, мы можем задать распределение вероятностей над всеми возможными разметками последовательности $X$ (это распределение связано с графической моделью Conditional Random Field, CRF):\n",
    "$$\n",
    "P(Y| X,\\Theta) = \\frac{1}{Z(\\Theta)} \\exp\\{F(Y| X, \\Theta)\\}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Fdd67DK2UCbi"
   },
   "source": [
    "### Предсказание\n",
    "\n",
    "Для фиксированных значений параметров $\\Theta$ предсказание может быть сделано, например, при помощи максимизации score функции $F$ (соответствует моде распределения над разметками). Для функций, связи между переменными в которых образуют граф вида цепочка (возможно для любого дерева), задача максимизации может быть решена точно за полиномиальное время при помощи метода динамического программирования.\n",
    "\n",
    "Выведем конкретный алгоритм, используя подход динамического программирования для решения задачи максимизации score функции:\n",
    "$$\n",
    "\\max_{Y} F(Y|X,\\Theta) = \\max_{Y} \\left[ \\sum\\limits_{i=0}^{L-1} U(x_i, y_i) + \\sum\\limits_{i=0}^{L-2} W(y_{i}, y_{i+1}) \\right]\n",
    "$$\n",
    "\n",
    "Выполняя алгебраические преобразования, задачу можно переписать следующим образом:\n",
    "$$\n",
    "\\max_{Y} \\left[ \\sum\\limits_{i=0}^{L-1} U(x_i, y_i) + \\sum\\limits_{i=0}^{L-2} W(y_{i}, y_{i+1}) \\right] = \\max_{y_0} \\left[U(x_0, y_0) + \\max_{y_1,...,y_{L-1}}\\left( W(y_0, y_1) + \\sum\\limits_{i=1}^{L-1} U(x_i, y_i) + \\sum\\limits_{i=1}^{L-2} W(y_{i}, y_{i+1}) \\right) \\right]\n",
    "$$\n",
    "\n",
    "В качестве подзадач динамического программирования будем использовать внутренние максимумы. Обозначим через $V_j(y_j)$ такой максимум по переменным с индексами большими чем $j$, т.е.,\n",
    "$$\n",
    "V_j(y_j) = U(x_j, y_j) + \\max_{y_{j+1},...,y_{L-1}}\\left( W(y_j, y_{j+1}) + \\sum\\limits_{i=j+1}^{L-1} U(x_i, y_i) + \\sum\\limits_{i=j+1}^{L-2} W(y_{i}, y_{i+1}) \\right).\n",
    "$$\n",
    "Динамическое программирование основано на итеративном вычислении $V_j(y_j)$ на основе ранее вычисленных значений. Используется следующая формула пересчёта:\n",
    "$$\n",
    "V_j(y_j) = U(x_j, y_j) + \\max_{y_{j+1}}\\left[ W(y_j, y_{j+1}) + V_{j+1}(y_{j+1}) \\right].\n",
    "$$\n",
    "Инициализировать пересчёт можно так: $V_{L-1}(y_{L-1}) = U(x_{L-1}, y_{L-1})$. Значение score на наилучшей конфигурации (решение задачи) можно найти при помощи максимизации $\\max_{y_0} [V_0(y_0)]$.\n",
    "\n",
    "Используя сохраненные индексы максимумов в каждой из задач максимизации, можно сделать проход в обратном направлении и восстановить оптимальную разметку.\n",
    "\n",
    "**Задание 1 (0.3 балла):** напишите недостающий код в ячейке ниже."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "DzBSkb-oUCbi"
   },
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "NUM_LABELS = 26\n",
    "\n",
    "def dynamic_programming(U, W):\n",
    "    \"\"\"\n",
    "    Parameters:\n",
    "        U: unary potentials, torch tensor shape (len_word, NUM_LABELS)\n",
    "        W: pairwise potentials, torch tensor shape (NUM_LABELS, NUM_LABELS)\n",
    "    \n",
    "    Returns:\n",
    "        arg_classes: argmaximum, torch long tensor shape (len_word)\n",
    "    \"\"\"\n",
    "    L = U.size(0)\n",
    "    V, argmax = torch.zeros(L, NUM_LABELS),\\\n",
    "                torch.zeros(L, NUM_LABELS)\n",
    "    \n",
    "    ### code starts here ###\n",
    "    \n",
    "    ### code ends here ###\n",
    "\n",
    "    return arg_classes.long()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZsjJKp2FUCbi"
   },
   "source": [
    "Если все реализовано верно, вы должны получить вывод: \"nconsequential\""
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "eDKNdklsUCbi"
   },
   "source": [
    "U = torch.load('unary_example.pth')\n",
    "W = torch.load('pairwise_example.pth')\n",
    "pred = dynamic_programming(U, W)\n",
    "pred = [int(i) for i in pred]\n",
    "print(''.join(le.inverse_transform(pred)))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I6YB059-UCbi"
   },
   "source": [
    "Обратите внимание, что если делать предсказание только по унарным потенциалам, то алгоритм делает ошибки."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "PVe-q5GeUCbi"
   },
   "source": [
    "_, u_labels = U.max(1)\n",
    "u_labels = [int(i) for i in u_labels]\n",
    "print(''.join(le.inverse_transform(u_labels)))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qi1uPjldUCbi"
   },
   "source": [
    "## Настройка параметров $\\Theta$ при помощи структурного метода опорных векторов\n",
    "Для настройки параметров $\\Theta$ будем использовать структурный метод опорных векторов (Structured SVM, SSVM). Интуитивно оптимизация данного функционала позволит обеспечить высокий score на правильных разметках и низкий score на неправильных. Функция потерь SSVM на одном объекте выборки $X$, $Y$ выглядит так:\n",
    "$$\n",
    "\\max_{Y'} \\left[\\Delta(Y,Y') + F(Y',X,\\Theta)\\right] - F(Y,X,\\Theta).  \n",
    "$$\n",
    "Здесь $\\Delta(Y,Y')$ - это функция, обобщающая отступ (margin) из классического SVM. Мы будет к качестве $\\Delta$ использовать нормированное расстояние Хэмминга между последовательностями $Y$ и $Y'$, т.е. $\\Delta(Y,Y') = \\frac{1}{L}\\sum\\limits_{i=1}^{L} [y_i \\neq y_i']$.\n",
    "\n",
    "Задача максимизации, возникающая в функции потерь, может быть решена при помощи уже реализованного алгоритма динамического программирования (добавление функции $\\Delta$ в данном случае не усложняет задачу, поскольку не меняет структуру графа).\n",
    "Для добавления $\\Delta$ в score-функцию достаточно добавить $\\frac{1}{L}$ ко всем унарным потенциалам, соответствующим неправильным меткам.\n",
    "\n",
    "Процедура обучения (настройки параметров $\\Theta$) состоит в минимизации функции потерь (усредненной по обучающей выборке) по $\\Theta$ при помощи методов стохастической оптимизации. При обработке каждого элемента выборки нужно решать задачу максимизации $F+\\Delta$. После нахождения оптимальной конфигурации (argmax) достаточно подставить полученные $Y'$ и вести оптимизацию по $\\Theta$ опуская слагаемое, отвечающее $\\Delta$. На лекции методы этой группы назывались \"структурным пулингом\"."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6lc0pW2SUCbi"
   },
   "source": [
    "**Задание 2 (0.2 балла):** реализуйте решение задачи максимизации из функции потерь SSVM (loss-augmented inference) при помощи вызова ранее реализованного алгоритма динамического программирования. Для тестирования кода добавьте возможность умножения расстояния Хэмминга на вес weight."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "t3nJtU-YUCbi"
   },
   "source": [
    "def loss_aug_inference(U, W, target, weight=1.0):\n",
    "    \"\"\"\n",
    "    Parameters:\n",
    "        U: unary potentials, torch tensor shape (len_word, NUM_LABELS)\n",
    "        W: pairwise potentials, torch tensor shape (NUM_LABELS, NUM_LABELS)\n",
    "        target: true configuration, torch long tensor shape (len_word)\n",
    "        weight: (for debug) put more weight on the loss term\n",
    "    Returns:\n",
    "        arg_classes: argmaximum, torch long tensor shape (len_word)\n",
    "    \"\"\"\n",
    "    ### code starts here ###\n",
    "    \n",
    "    ### code ends here ###\n",
    "    return arg_classes.long()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "SmdDJ5GAUCbi"
   },
   "source": [
    "U = torch.load('unary_example.pth')\n",
    "W = torch.load('pairwise_example.pth')\n",
    "target = torch.LongTensor([13, 2, 14, 13, 18, 4, 16, 20, 4, 13, 19, 8, 0, 11])\n",
    "pred = loss_aug_inference(U, W, target, weight=60.0)\n",
    "correct = torch.LongTensor([13,2, 14, 13, 18, 5, 14, 21, 4, 13, 19, 8, 13, 2])\n",
    "assert pred.eq(correct).sum() == correct.numel(), \"Check your loss_aug_inference\""
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XGNfw5ZeUCbi"
   },
   "source": [
    "### Унарные потенциалы\n",
    "Для извлечения унарных потенциалов будем использовать нейросеть вида LeNet. Для экономии времени сеть обучена заранее (стандартная схема для MNIST, но с 26 классами).\n",
    "\n",
    "Стоит также заметить, что unary network без использования парных потенциалов достигает качества на валидации 0.92 (точность предсказания всех символов датасета)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "3af9ylwEUCbi"
   },
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "class LeNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(LeNet, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 10, 5)\n",
    "        self.conv2 = nn.Conv2d(10, 20, 5)\n",
    "        self.fc1 = nn.Linear(5 * 5 * 20, 140)\n",
    "        self.fc2 = nn.Linear(140, 26)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.max_pool2d(F.relu(self.conv1(x)), 2)\n",
    "        x = F.max_pool2d(F.relu(self.conv2(x)), 2)\n",
    "        x = x.view(-1, 5 * 5 * 20)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "unary_net = LeNet()\n",
    "unary_net.load_state_dict(torch.load('state_dict.pth'))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-ErlQd_DUCbi"
   },
   "source": [
    "### Обучение\n",
    "\n",
    "**Задание 3 (0.1 балла):** реализуйте подсчёт score-функции для вычисления функции потерь SSVM."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "z80MAlygUCbj"
   },
   "source": [
    "def score_function(Y, U, W):\n",
    "    \"\"\"\n",
    "    Parameters:\n",
    "        U: unary potentials, torch tensor shape (len_word, NUM_LABELS)\n",
    "        W: pairwise potentials, torch tensor shape (NUM_LABELS, NUM_LABELS)\n",
    "        Y: configuration, torch long tensor shape (len_word)\n",
    "    \n",
    "    Returns:\n",
    "        value of score function\n",
    "    \"\"\"\n",
    "    ### code starts here ###\n",
    "    \n",
    "    ### code ends here ###\n",
    "    return score_value"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "u3waP2NjUCbj"
   },
   "source": [
    "U = torch.load('unary_example.pth')\n",
    "W = torch.load('pairwise_example.pth')\n",
    "target = torch.LongTensor([13, 2, 14, 13, 18, 4, 16, 20, 4, 13, 19, 8, 0, 11])\n",
    "s = score_function(target, U, W)\n",
    "assert np.allclose(score_function(target, U, W), 175.58605), 'Check you score function'"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j8NjkyOXUCbj"
   },
   "source": [
    "Теперь нужно реализовать подсчёт функции потерь SSVM и вызов оптимизатора.\n",
    "\n",
    "**Задание 4 (0.1 балла):** напишите недостающий код в ячейке ниже."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "XMCkm1ZYUCbj"
   },
   "source": [
    "from torch.optim import Adam\n",
    "\n",
    "from tqdm import trange\n",
    "from IPython.display import clear_output\n",
    "\n",
    "trace_values = []\n",
    "\n",
    "torch.manual_seed(42)\n",
    "W = torch.randn(NUM_LABELS, NUM_LABELS, requires_grad=True)\n",
    "opt = Adam([W], lr=1e-2)\n",
    "\n",
    "n_epoch = 3\n",
    "for epoch in range(n_epoch):\n",
    "    print('Epoch', epoch)\n",
    "    mean_val = 0.\n",
    "    for i in trange(len(train_x)):\n",
    "        word, target = torch.from_numpy(train_x[i]).float(),\\\n",
    "                       torch.from_numpy(train_y[i]).long()\n",
    "        \n",
    "        U = unary_net(word)\n",
    "        y_ = loss_aug_inference(U, W, target)\n",
    "        ### code starts here ###\n",
    "        \n",
    "        ### code ends here ###\n",
    "        mean_val += loss.item() + 1. - y_.eq(target).float().sum() / U.size(0)\n",
    "        if i % 500 == 0 and i:\n",
    "            trace_values.append(mean_val / 500.)\n",
    "            mean_val = 0.\n",
    "            \n",
    "            clear_output()\n",
    "            plt.title('SSVM loss, current epoch {}'.format(epoch))\n",
    "            plt.plot(np.arange(len(trace_values)), trace_values)\n",
    "            plt.show()\n",
    "\n",
    "        \n",
    "    glob_acc = 0.\n",
    "    letters_num = np.sum([i.shape[0] for i in val_x])\n",
    "    for i in range(len(val_x)):\n",
    "        word, target = torch.from_numpy(val_x[i]).float(),\\\n",
    "                       torch.from_numpy(val_y[i]).long()\n",
    "        U, P = unary_net(word), W\n",
    "        pred = dynamic_programming(U, P)\n",
    "        eq_count = pred.eq(target).sum()\n",
    "        glob_acc += eq_count\n",
    "    glob_acc = glob_acc.float() / letters_num\n",
    "\n",
    "    print('global val accuracy {}'.format(glob_acc))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vgxf8xCmUCbj"
   },
   "source": [
    "Точность на валидации должна получиться близкой к 0.965."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "--oMv6ViUCbj"
   },
   "source": [
    "## Настройка параметров $\\Theta$ при помощи метода максимального правдоподобия\n",
    "\n",
    "Другим подходом к настройке параметров $\\Theta$ является метод максимального правдоподобия.\n",
    "Метод состоит в максимизации лог-правдоподобия на обучающей выборке. Правдоподобие задаётся следущим распределением вероятностей:\n",
    "$$\n",
    "P(Y| X,\\Theta) = \\frac{1}{Z(\\Theta)} \\exp\\{F(Y| X, \\Theta)\\}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CPpVDsDYUCbj"
   },
   "source": [
    "### Вычисление нормировочной константы $Z$\n",
    "Для модели цепочки нормировочная константа может быть посчитана эффективно с использованием sum-product belief propagation (динамическое программирование) \n",
    "\n",
    "$$\n",
    "Z(\\Theta) = \\sum\\limits_{Y'} \\exp\\{F(Y'| X, \\Theta)\\} = \\sum\\limits_{Y'} \\exp\\left\\{\\sum\\limits_{i=0}^{L-1} U(x_i, y_i') + \\sum\\limits_{i=0}^{L-2} W(y_{i}', y_{i+1}'))\\right\\}\n",
    "$$\n",
    "$$\n",
    "= \\sum\\limits_{y_0'} \\exp\\{U(x_0, y_0')\\}\\sum\\limits_{y_1',...,y_{L-1}'}\\exp\\left\\{W(y_{0}', y_{1}') + \\sum\\limits_{i=1}^{L-1} U(x_i, y_i') + \\sum\\limits_{i=2}^{L-1} W(y_{i}', y_{i+1}'))\\right\\}\n",
    "$$\n",
    "Определеним подзадачи динамического программирования (аналогично предсказанию, но сумма заменена на произведение, а максимум на сумму).\n",
    "$$\n",
    "V_j(y_j) = \\exp\\{U(x_j, y_j)\\} \\sum_{y_{j+1},...,y_{L-1}}\\left( \\exp\\{W(y_j, y_{j+1})\\} \\prod\\limits_{i=j+1}^{L-1}  \\exp\\{U(x_i, y_i)\\} \\prod\\limits_{i=j+1}^{L-2}  \\exp\\{W(y_{i}, y_{i+1})\\} \\right).\n",
    "$$\n",
    "Динамическое программирование основано на итеративном вычислении $V_j(y_j)$ на основе ранее вычисленных значений. Используется следующая формула пересчёта:\n",
    "$$\n",
    "V_j(y_j) = \\exp\\{U(x_{j}, y_{j})\\} \\sum_{y_{j+1}}\\left[ \\exp\\{W(y_j, y_{j+1})\\} V_{j+1}(y_{j+1}) \\right].\n",
    "$$\n",
    "Инициализировать пересчёт можно так: $V_{L-1}(y_{L-1}) = \\exp\\{U(x_{L-1}, y_{L-1})\\}$. Окончательное значение нормировочной константы можно найти как $\\sum_{y_0} [V_0(y_0)]$.\n",
    "\n",
    "Для численно устойчивой реализации необходимо использовать функцию log_sum_exp и производить вычисления в логарифмической шкале, т.е. найти $\\log Z(\\Theta)$.\n",
    "\n",
    "HINT: для log_sum_exp используйте max trick:\n",
    "\n",
    "$$\n",
    " \\log \\sum\\limits_{i=1}^N \\exp\\{x_i\\} =  \\log \\sum\\limits_{i=1}^N \\exp\\{x_i - \\max_{j}[x_j]\\} + \\max_{j}[x_j]\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oLZ0F0SFUCbj"
   },
   "source": [
    "**Задание 5 (0.2 балла):** Реализуйте подсчет нормировочной константы."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "WVHDCA-KUCbj"
   },
   "source": [
    "def log_sum_exp(vec, axis=0):\n",
    "    ### code starts here ###\n",
    "    \n",
    "    ### code ends here ###\n",
    "    return result\n",
    "\n",
    "def compute_log_partition(U, W):\n",
    "    \"\"\"\n",
    "    Parameters:\n",
    "        U: unary potentials, torch tensor shape (len_word, NUM_LABELS)\n",
    "        W: pairwise potentials, torch tensor shape (NUM_LABELS, NUM_LABELS)\n",
    "    \n",
    "    Returns:\n",
    "        value of partition function\n",
    "    \"\"\"\n",
    "    ### code starts here ###\n",
    "    \n",
    "    ### code ends here ###\n",
    "    return logZ"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "1CDsmtWCUCbj"
   },
   "source": [
    "U = torch.load('unary_example.pth')\n",
    "W = torch.load('pairwise_example.pth')\n",
    "assert np.allclose(compute_log_partition(U, W), 175.63, rtol=1e-4, atol=1e-6), 'Check you compatability function'"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pflfq0NnUCbj"
   },
   "source": [
    "### Обучение\n",
    "**Задание 6 (0.1 балла)** Реализуйте подсчет negative log-likelihood и шаг оптимизации."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Ptjl4_MgUCbj"
   },
   "source": [
    "from torch.optim import Adam\n",
    "from tqdm import trange\n",
    "\n",
    "torch.manual_seed(42)\n",
    "W = torch.randn(NUM_LABELS, NUM_LABELS, requires_grad=True)\n",
    "opt = Adam([W], lr=1e-2)\n",
    "\n",
    "trace_values = []\n",
    "\n",
    "n_epoch = 3\n",
    "for epoch in range(n_epoch):\n",
    "    print('Epoch', epoch)\n",
    "    mean_val = 0.\n",
    "    for i in trange(len(train_x)):\n",
    "        word, target = torch.from_numpy(train_x[i]).float(),\\\n",
    "                       torch.from_numpy(train_y[i]).long()\n",
    "        \n",
    "        U = unary_net(word)\n",
    "        logZ = compute_log_partition(U, W)\n",
    "\n",
    "        ### code starts here ###\n",
    "        \n",
    "        ### code ends here ###\n",
    "        \n",
    "        mean_val += loss.item()\n",
    "        if i % 500 == 0 and i:\n",
    "            trace_values.append(mean_val / 500.)\n",
    "            mean_val = 0.\n",
    "            \n",
    "            clear_output()\n",
    "            plt.title('Negative loglikelihood loss, current epoch {}'.format(epoch))\n",
    "            plt.plot(np.arange(len(trace_values)), trace_values)\n",
    "            plt.show()\n",
    "        \n",
    "    glob_acc = 0.\n",
    "    letters_num = np.sum([i.shape[0] for i in val_x])\n",
    "    for i in range(len(val_x)):\n",
    "        word, target = torch.from_numpy(val_x[i]).float(),\\\n",
    "                       torch.from_numpy(val_y[i]).long()\n",
    "        U, P = unary_net(word), W\n",
    "        pred = dynamic_programming(U, P)\n",
    "        eq_count = pred.eq(target).sum()\n",
    "        glob_acc += eq_count\n",
    "    glob_acc = glob_acc.float() / letters_num\n",
    "\n",
    "    print('global val accuracy {}'.format(glob_acc))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n8U51SRAUCbj"
   },
   "source": [
    "Точность на валидации должна получиться в районе 0.97."
   ]
  }
 ]
}