{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Лабораторная работа. Программирование тензорных ядер\n",
    "\n",
    "Тензорные ядра умеют выполнять только одну операцию - умножать матрицы. Их основное (если не единственное) применение - нейронные сети. В этой лабораторной работе нам предстоит запрограммировать простую нейронную сеть для выполнения на тензорных ядрах. Конечно, эта задача уже решена - есть множество фрейворков удобного построения нейронных сетей (например, pytorch) и уже написаны реализации нейронных слоев, использующие тензорные ядра (библиотека cuDNN). Наш опыт будет практически полезен только в том случае, если нам когда-либо прийдется переносить на видеокарту новые типы нейронных слоев или реализовывать старые на новых архитектурах.\n",
    "\n",
    "## 1. Низкоуровневое программирование нейронной сети\n",
    "\n",
    "Рассмотрим простейшую двухслойную сеть (в квадратных скобках указан выходная размерность, FC - fully connected - полносвязанный слой): \n",
    "```\n",
    "x[128] → FC1[32] → sigmoid → FC2[1]\n",
    "```\n",
    "Если мы будем использовать квадратичную функцию потерь $L = (y^{pred}-y^{true})^2$, то тренировка нашей нейронной сети фактически является оптимизацией функции $L$, зависящей от весов полносвязанных слоев FC1 и FC2. Обозначим их $w1_{ij}$ и $w2_i$ и запишем формулы, полностью определяющие нашу оптимизируемую функцию:\n",
    "$$\n",
    "L(w1,w2) = (FC2-y^{true})^2\n",
    "$$\n",
    "$$\n",
    "FC2 = \\sum\\limits_{i=0}^{31} \\sigma(FC1_i) w2_i\n",
    "$$\n",
    "$$\n",
    "FC1_i = \\sum\\limits_{j=0}^{127} x_j w1_{ji} \n",
    "$$\n",
    "### Батчи\n",
    "\n",
    "На практике объекты x выбирают не по одному, а группами (батчами), чтобы эффективней использовать видеокарту. Пусть размер батча у нас будет равен 8, тогда формулы немного изменятся (индекс координаты батча по традиции всегда указывают на первом месте):\n",
    "$$\n",
    "L(w1,w2) = \\sum\\limits_{i=0}^7 (FC2_i-y^{true}_i)^2\n",
    "$$\n",
    "$$\n",
    "FC2_i = \\sum\\limits_{j=0}^{31} \\sigma(FC1_{ij}) w2_j\n",
    "$$\n",
    "$$\n",
    "FC1_{ij} = \\sum\\limits_{k=0}^{127} x_{ik} w1_{kj}\n",
    "$$\n",
    "### Правило Эйнштейна\n",
    "\n",
    "Обычно те, кто часто работает с такими формулами, пользуются [соглашением Эйнштена](https://ru.wikipedia.org/wiki/%D0%A1%D0%BE%D0%B3%D0%BB%D0%B0%D1%88%D0%B5%D0%BD%D0%B8%D0%B5_%D0%AD%D0%B9%D0%BD%D1%88%D1%82%D0%B5%D0%B9%D0%BD%D0%B0) и опускают символы суммирования, предполагая, что в каждом произведении суммирование происходит по повторяющемуся индексу. Тогда можно эти три формулы объединить в одну:\n",
    "$$\n",
    "L(w1,w2) = \\sum\\limits_{i=0}^7 \\left( \\sigma\\left( x_{ik} w1_{kj} \\right)w2_j - y^{true}_i  \\right)^2\n",
    "$$\n",
    "### Метод стохастического градиента\n",
    "\n",
    "Тренировка нейронной сети методом стохастического градиента состоит в итерационном изменении весов $w1$ и $w2$ на величину, пропорциональную производной (градиенту) оптимизируемой функции $L$ по соответствующей переменной:\n",
    "$$\n",
    "w1_{mn} := w1_{mn} - learning\\_rate * \\frac{\\partial L}{\\partial w1_{mn}} \n",
    "$$\n",
    "$$\n",
    "w2_m := w2_m - learning\\_rate * \\frac{\\partial L}{\\partial w2_m}\n",
    "$$\n",
    "\n",
    "Вычислив производные $L$ (повторите это на листике!!!), мы придем к следующим формулам тренировки нашей нейронной сети:\n",
    "$$\n",
    "w1_{mn} := w1_{mn} - learning\\_rate * 2\\sum\\limits_{i=0}^7 \\left( \\sigma\\left( x_{ik} w1_{kj} \\right)w2_j - y^{true}_i  \\right) \\left(  \\sigma'\\left( x_{ik} w1_{kn} \\right) w2_n \\right) x_{im}\n",
    "$$\n",
    "$$\n",
    "w2_m := w2_m - learning\\_rate * 2\\sum\\limits_{i=0}^7 \\left( \\sigma\\left( x_{ik} w1_{kj} \\right)w2_j - y^{true}_i  \\right)\\sigma\\left( x_{ik} w1_{km} \\right)\n",
    "$$\n",
    "Помните: в каждом произведении суммирование происходит по повторяющимся индексам (кроме $m$ и $n$)!\n",
    "\n",
    "Форумула для производной сигмоиды: $\\sigma'(\\tau) = \\sigma(\\tau)(1-\\sigma(\\tau))$, а также переиспользование выражений, вычисленных в прямом ходе, позволит ускорить вычисления градиентов в обратном.\n",
    "\n",
    "## Реализация средствами numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from time import time\n",
    "\n",
    "batch_sz = 8\n",
    "dim_x = 128\n",
    "dim_fc1 = 32\n",
    "\n",
    "class NumpyModel:\n",
    "    def __init__(self, dim_x):\n",
    "        # конструируем веса со случайными значениями\n",
    "        self.w1 = np.float32(np.random.uniform(low=-1, high=+1, size=(dim_x, dim_fc1)))\n",
    "        self.w2 = np.float32(np.random.uniform(low=-1, high=+1, size=(dim_fc1, 1)))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        t0 = time()\n",
    "        # запоминаем промежуточные вычисления для использования их в обратном ходе\n",
    "        self.x = x\n",
    "        # np.dot - это обычное матричное умножение (для многомерных матриц A*B\n",
    "        # сворачиваются последнее в A и предпоследнее в B измерения)\n",
    "        fc1 = np.dot(x, self.w1)  # shape = 8 x 32\n",
    "        self.sigmoid = 1/(1+np.exp(-fc1))  # shape = 8 x 32\n",
    "        self.fc2 = np.dot(self.sigmoid, self.w2)  # shape = 8\n",
    "        print('Forward time =', (time()-t0)*1000, 'ms')\n",
    "        return self.fc2\n",
    "    \n",
    "    def backward(self, y_true):\n",
    "        t0 = time()\n",
    "        diff = self.fc2-y_true  # shape = 8 x 1\n",
    "        dsigma = self.sigmoid*(1 - self.sigmoid)  # shape = 8 x 32\n",
    "        dsigma_mul_w2 = dsigma * self.w2.flatten()  # shape = 8 x 32\n",
    "        dsigma_mul_w2_mul_x = dsigma_mul_w2[:,np.newaxis,:] * self.x[:,:,np.newaxis]  # shape = 8 x 128 x 32\n",
    "        grad_loss_by_w1 = 2*np.sum(dsigma_mul_w2_mul_x * diff[:,:,np.newaxis], axis=0)  # shape = 128 x 32\n",
    "        assert np.all(grad_loss_by_w1.shape == self.w1.shape), str(grad_loss_by_w1.shape)\n",
    "        \n",
    "        grad_loss_by_w2 = 2*np.dot(self.sigmoid.T, diff)\n",
    "        assert np.all(grad_loss_by_w2.shape == self.w2.shape), str(grad_loss_by_w2.shape)\n",
    "        print('Backward time =', (time()-t0)*1000, 'ms')\n",
    "        return grad_loss_by_w1, grad_loss_by_w2\n",
    "    \n",
    "np.random.seed(0)\n",
    "# \"неизвестная\" точная целевая функция\n",
    "true_model =  NumpyModel(dim_x)\n",
    "\n",
    "# датасет\n",
    "x = np.float32(np.random.uniform(low=-0.5, high=0.5, size=(10000, dim_x)))\n",
    "\n",
    "# значения целевой функции для объектов нашего датасета\n",
    "y0 = true_model.forward(x)\n",
    "\n",
    "#===========================================================================================\n",
    "# Тренировка нейронной сети\n",
    "#===========================================================================================\n",
    "\n",
    "# генерируем случайное начальное приближение к неизвестным весам\n",
    "np.random.seed(1) # мы используем другой seed !!!\n",
    "approx_model = NumpyModel(dim_x)\n",
    "\n",
    "learning_rate = 0.1\n",
    "mean_loss = 0\n",
    "# цикл по батчам\n",
    "for cycle in range(1):\n",
    "    # формируем батч - выбираем случайные объекты из тренировочной выборки\n",
    "    ind = np.random.choice(len(x), size=batch_sz)\n",
    "    x_batch = x[ind, :]\n",
    "    y_true = y0[ind]\n",
    "    \n",
    "    # Прямой ход: вычисляем предсказания y_pred для x_batch\n",
    "    y_pred = approx_model.forward(x_batch)\n",
    "    print('y_pred =', y_pred)\n",
    "    \n",
    "    # ошибка на текущем батче\n",
    "    L = np.sum((y_pred - y_true)**2)\n",
    "    \n",
    "    # Печатаем ошибку, осредненную методом экспоненциально взвешенного скользящего среднего\n",
    "    # https://en.wikipedia.org/wiki/Moving_average#Exponential_moving_average\n",
    "    lamb = 0.01\n",
    "    mean_loss = lamb*L + (1-lamb)*mean_loss\n",
    "    print(f'cycle = {cycle}, mean_loss = {mean_loss}')\n",
    "    \n",
    "    # вычисляем градиенты\n",
    "    grad_loss_by_w1, grad_loss_by_w2 = approx_model.backward(y_true)\n",
    "    print(grad_loss_by_w1)\n",
    "    print(grad_loss_by_w2)\n",
    "    \n",
    "    # корректируем веса в сторону, обратную градиенту ошибки\n",
    "    learning_rate = 0.1\n",
    "    approx_model.w1 -= learning_rate*grad_loss_by_w1\n",
    "    approx_model.w2 -= learning_rate*grad_loss_by_w2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Реализация средствами библиотеки pytorch\n",
    "\n",
    "Чтобы проверить наши формулы вычисления градиентов и сравнить скорость нашей реализации с профессиональной, необходимо реализовать нашу нейронную сеть на pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch as t\n",
    "\n",
    "dev = t.device(\"cpu\")  # используйте cuda:0 или cpu\n",
    "\n",
    "class TorchModel(t.nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        # класс Parameter нужен, что автоматически вычислялись градиенты\n",
    "        self.w1 = t.nn.Parameter(t.Tensor(np.random.uniform(low=-1, high=+1, size=(dim_x, dim_fc1))))\n",
    "        self.w2 = t.nn.Parameter(t.Tensor(np.random.uniform(low=-1, high=+1, size=(dim_fc1, 1))))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        t0 = time()\n",
    "        self.x = x\n",
    "        self.fc1 = x.mm(self.w1)\n",
    "        self.sigmoid = t.sigmoid(self.fc1)\n",
    "        self.fc2 = self.sigmoid.mm(self.w2)\n",
    "        print('Forward time =', (time()-t0)*1000, 'ms')\n",
    "        return self.fc2\n",
    "\n",
    "np.random.seed(1) # если мы хотим проверять градиенты, то мы должны использовать тот же seed!!!\n",
    "torch_model = TorchModel()\n",
    "torch_model.to(dev)\n",
    "\n",
    "mean_loss = 0\n",
    "for cycle in range(1):\n",
    "    \n",
    "    ind = np.random.choice(len(x), size=batch_sz)\n",
    "    x_batch = t.tensor(x[ind], device=dev)\n",
    "    y_true = t.tensor(y0[ind], device=dev)\n",
    "    \n",
    "    # прямой ход: вычисляем приближенный y_pred для x_batch\n",
    "    y_pred = torch_model(x_batch)\n",
    "    print('y_pred =', y_pred.data)\n",
    "    \n",
    "    # считаем величину ошибки\n",
    "    loss = t.sum((y_pred - y_true)**2)\n",
    "    \n",
    "    lamb = 0.01\n",
    "    mean_loss = lamb*loss.data + (1-lamb)*mean_loss # .data - открепляет градиент, иначе будут проблемы!!!\n",
    "    print(f'cycle = {cycle}, loss = {mean_loss}')\n",
    "    \n",
    "    # вычисляем градиенты\n",
    "    t0 = time()\n",
    "    loss.backward()\n",
    "    print('Backward time =', (time()-t0)*1000, 'ms')\n",
    "    \n",
    "    torch_model.w1.data -= learning_rate * torch_model.w1.grad.data\n",
    "    torch_model.w2.data -= learning_rate * torch_model.w2.grad.data\n",
    "    \n",
    "    print('∇w1 =', torch_model.w1.grad)\n",
    "    print('∇w2 =', torch_model.w2.grad)\n",
    "\n",
    "    # после обновления весов нужно обнулить градиенты т.к. при их вычислении используется сумма с накоплением\n",
    "    torch_model.w1.grad.data.zero_()\n",
    "    torch_model.w2.grad.data.zero_()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# если мы хотим проверять градиенты и на C++, то мы должны использовать тот же seed и сохранить все в файл\n",
    "np.random.seed(1) \n",
    "m = NumpyModel(dim_x)\n",
    "np.savetxt('w1.txt', m.w1, delimiter=' ')\n",
    "np.savetxt('w2.txt', m.w2, delimiter=' ')\n",
    "np.savetxt('x.txt', x, delimiter=' ')\n",
    "np.savetxt('y0.txt', y0, delimiter=' ')\n",
    "ind = np.random.choice(len(x), size=batch_sz)\n",
    "np.savetxt('ind.txt', ind, delimiter=' ', fmt='%d')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}