{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Обучение нейросетей: NumPy, PyTorch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Данный ноутбук посвещен изучению функционала и абстракций, которые нам представляет PyTorch для разработки нейронных сетей. Основные моменты, которые будут рассматриваться на данном семинаре:\n",
    "* Реализация оптимизации градиентным спуском, используя NumPy\n",
    "* Использование PyTorch Tensors, как замены NumPy ndarrays\n",
    "* Автоматизация подсчета градиентов при помощи `torch.autograd`\n",
    "* Использование `torch.nn` для построения вычислительного графа при помощи уже реализованных слоев\n",
    "* Пакет `torch.optim`, предоставляющий абтракцию оптимизации нейросетей и обновления ее весов"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!curl -O 'https://courses.cv-gml.ru/storage/seminars/nn-training-numpy-pytorch/test.zip'\n",
    "!unzip -o test.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import traceback\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "# Increase these if figures appear small\n",
    "plt.rcParams[\"figure.figsize\"] = fx, fy = (14.08, 6.40)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_random_data(num_samples, num_features, num_targets, seed=42):\n",
    "    rng = np.random.default_rng(seed)\n",
    "    feature_locs = rng.normal(size=(1, num_features))\n",
    "    feature_scales = np.exp(rng.normal(size=(1, num_features)))\n",
    "    features = rng.normal(\n",
    "        loc=feature_locs,\n",
    "        scale=feature_scales,\n",
    "        size=(num_samples, num_features),\n",
    "    )\n",
    "\n",
    "    targets = []\n",
    "    for _ in range(num_targets):\n",
    "        num_deps = 1 + rng.integers(num_features)\n",
    "        dep_inds = rng.choice(num_features, num_deps)\n",
    "        deps = features[:, dep_inds]\n",
    "\n",
    "        offsets = rng.uniform(0, 2 * np.pi, size=num_deps)\n",
    "        scales = rng.normal(size=num_deps)\n",
    "\n",
    "        target = scales[None, :] * np.sin(deps + offsets[None, :])\n",
    "        targets.append(target.sum(axis=-1))\n",
    "\n",
    "    targets = np.stack(targets, axis=-1)\n",
    "    return features, targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_results():\n",
    "    # Show training loss history\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.title(\"Training loss history\")\n",
    "    plt.plot(loss_hist)\n",
    "    plt.yscale(\"log\")\n",
    "    plt.xlabel(\"step\")\n",
    "    plt.ylabel(\"loss\")\n",
    "\n",
    "    # Visualize ground truth vs predictions\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.title(\"Ground truth vs predictions\")\n",
    "    for i in range(D_out):\n",
    "        plt.scatter(y_vis[:, i], p_vis[:, i], s=1)\n",
    "    plt.xlabel(\"y\")\n",
    "    plt.ylabel(\"p\")\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Разминка на NumPy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Перед тем, как переходить к рассмотрению PyTorch, давайте сначала реализуем логику для запуска и оптимизации градиентным спуском простой нейросети с помощью NumPy.\n",
    "\n",
    "NumPy - это framework для научных вычислений; он ничего не знает о вычислительных графах, глубоком обучении или градиентах. Однако мы можем легко использовать NumPy для обучения простой двухслойной нейросети на случайных данных вручную, реализуя прямые и обратные проходы через сеть с помощью NumPy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В данном семинаре мы будем реализовывать \"классический\" метод градиентного спуска, то есть ищем решение для задача оптимизации\n",
    "\n",
    "$$\n",
    "f\\!\\left(P\\right) = \\mathrm{loss}\\!\\left(Y,\\,\\mathrm{nn}\\!\\left(P,\\,X\\right)\\right)\\\\\n",
    "\\min\\limits_{P} f\\!\\left(P\\right)\n",
    "$$\n",
    "\n",
    "где $X, Y$ - **все** имеющиеся эталонные данные, $P$ - обучаемые параметры сети, $\\mathrm{nn}$ - заранее фиксированная математическая функция, задающая архитектуру сети и $\\mathrm{loss}$ - функция потерь (дифференцируемая функция оценивающая ошибку предсказания, которую мы хотим минимизировать).\n",
    "\n",
    "Тогда метод градиентного спуска предлагает на каждом шаге вычислить локальный градиент функции по градиентам и обновить эти параметры по правилу\n",
    "\n",
    "$$\n",
    "P \\leftarrow P - \\alpha \\cdot \\frac{\\partial f}{\\partial P}\n",
    "$$\n",
    "\n",
    "Подробнее про метод градиентного спуска:\n",
    "- [machinelearning.ru](http://www.machinelearning.ru/wiki/index.php?title=Метод_градиентного_спуска)\n",
    "- [Wikipedia](https://en.wikipedia.org/wiki/Gradient_descent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# N is number of samples\n",
    "N = 100\n",
    "# D_in is the input dimensionality\n",
    "D_in = 16\n",
    "# D_out is the target dimensionality\n",
    "D_out = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create random input and output data\n",
    "x, y = create_random_data(N, D_in, D_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's define a Multi-Layer Perceptron\n",
    "# H is the dimensionality of the hidden layer in the MLP\n",
    "H = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Randomly initialize weights\n",
    "scale1 = 1 / np.sqrt(D_in)\n",
    "scale2 = 1 / np.sqrt(H)\n",
    "w1 = np.random.uniform(-scale1, scale1, size=(D_in, H))\n",
    "w2 = np.random.uniform(-scale2, scale2, size=(H, D_out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters for the train loop\n",
    "learning_rate = 1e-6\n",
    "num_steps = 50_000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_hist = []\n",
    "progress = tqdm(range(num_steps))\n",
    "for t in progress:\n",
    "    # Forward pass: compute predicted y\n",
    "    h = x @ w1\n",
    "    h_relu = h.clip(0, None)\n",
    "    p = h_relu @ w2\n",
    "\n",
    "    # Compute the loss\n",
    "    delta = p - y\n",
    "    loss = (delta ** 2).sum()\n",
    "\n",
    "    # Save the loss value to history and display it in the progress bar\n",
    "    loss_hist.append(loss)\n",
    "    progress.desc = f\"loss: {loss:.8f}\"\n",
    "\n",
    "    # Backprop to compute gradients of w1 and w2 with respect to loss\n",
    "    grad_p = 2 * delta\n",
    "    grad_w2 = h_relu.T @ grad_p\n",
    "\n",
    "    grad_h_relu = grad_p @ w2.T\n",
    "    grad_h = (h >= 0) * grad_h_relu\n",
    "    grad_w1 = x.T @ grad_h\n",
    "\n",
    "    # Update weights\n",
    "    w1 -= learning_rate * grad_w1\n",
    "    w2 -= learning_rate * grad_w2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Implement inference/prediction\n",
    "p = (x @ w1).clip(0, None) @ w2\n",
    "\n",
    "p_vis = p\n",
    "y_vis = y\n",
    "visualize_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch Tensors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NumPy - отличный фреймворк, но он не может использовать GPU для ускорения численных вычислений. Для современных нейронных сетей графические процессоры часто обеспечивают ускорение в 50 раз или больше, поэтому использование NumPy не достаточно для реализации современных методов глубинного обучения.\n",
    "\n",
    "Здесь мы представляем самую фундаментальную концепцию PyTorch - тензоры. Тензоры в PyTorch - концептуально идентичны N-мерным массивам NumPy. PyTorch предоставляет множество функций для работы с этими тензорами. Любые вычисления, которые вы, возможно, захотите выполнить с помощью NumPy, также могут быть выполнены с помощью тензоров PyTorch. Многие (**но не все**) функции и методы для работы с тензорами имеют те же имена и интерфейсы как в NumPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = np.arange(6, dtype=np.float32).reshape(2, 3)\n",
    "a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Для конвертации между существующими NumPy массивами и PyTorch тензорами, можно использовать функцию `torch.from_numpy(ndarray)` и метод `torch.Tensor.numpy()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use torch.from_numpy() to convert from NumPy ndarray to PyTorch tensors\n",
    "t = torch.from_numpy(a)\n",
    "t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the torch.Tensor.numpy() method to convert a PyTorch tensor to Numpy\n",
    "t.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: they share the same memory (both are \"views\" of the same buffer)\n",
    "a[0, 0] = 42\n",
    "t"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В отличие от N-мерных массивов NumPy, тензоры PyTorch могут использовать GPU для ускорения численных вычислений. Чтобы создать тензор на GPU, можно передать соответствующий аргумент `device` при создании этого тензора. Уже существующий тензор можно скопировать на нужный девайс, используя методы `tensor.to(device)`, `tensor.cpu()` или `tensor.cuda()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    print(\"Using the GPU 😊\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    print(\"Using the CPU 😞\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = t.to(device)\n",
    "c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c.device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Обратите внимание, что для конвертации PyTorch тензора в NumPy массив необходимо, чтобы этот тензор находился на CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NumPy only supports CPU ndarrays\n",
    "try:\n",
    "    c.numpy()\n",
    "except TypeError:\n",
    "    traceback.print_exc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c.cpu().numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Давайте используем тензоры PyTorch, чтобы обучить нейросеть на случайных данных. Как и в приведенном выше примере на NumPy, давайте вручную реализуем прямой и обратный проходы по сети, используя операции с тензорами PyTorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = create_random_data(N, D_in, D_out)\n",
    "\n",
    "# Load input and output data onto the device\n",
    "x = torch.from_numpy(x).to(dtype=torch.float32, device=device)\n",
    "y = torch.from_numpy(y).to(dtype=torch.float32, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Randomly initialize weights\n",
    "w1 = torch.empty((D_in, H), device=device).uniform_(-scale1, scale1)\n",
    "w2 = torch.empty((H, D_out), device=device).uniform_(-scale2, scale2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy-paste the train loop from previous section\n",
    "# (you might not even need to change anything)\n",
    "\n",
    "# The following code is **exactly** the same as in the\n",
    "# previous train loop, sans comments and empty lines\n",
    "loss_hist = []\n",
    "progress = tqdm(range(num_steps))\n",
    "for t in progress:\n",
    "    h = x @ w1\n",
    "    h_relu = h.clip(0, None)\n",
    "    p = h_relu @ w2\n",
    "    delta = p - y\n",
    "    loss = (delta ** 2).sum()\n",
    "    loss_hist.append(loss)\n",
    "    progress.desc = f\"loss: {loss:.8f}\"\n",
    "    grad_p = 2 * delta\n",
    "    grad_w2 = h_relu.T @ grad_p\n",
    "    grad_h_relu = grad_p @ w2.T\n",
    "    grad_h = (h >= 0) * grad_h_relu\n",
    "    grad_w1 = x.T @ grad_h\n",
    "    w1 -= learning_rate * grad_w1\n",
    "    w2 -= learning_rate * grad_w2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy-paste and edit the inference code from previous section\n",
    "p = (x @ w1).clip(0, None) @ w2\n",
    "\n",
    "# Matplotlib and other CPU-based applications\n",
    "# may need CPU tensors (or NumPy ndarrays).\n",
    "if False:\n",
    "    # Before\n",
    "    p_vis = p\n",
    "    y_vis = y\n",
    "\n",
    "# After\n",
    "y_vis = y.cpu()\n",
    "p_vis = p.cpu()\n",
    "loss_hist = [l.cpu() for l in loss_hist]\n",
    "\n",
    "visualize_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch `autograd`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В примере выше нами был реализован, как прямой проход по сети, так и обратное распространение ошибки. Реализовывать каждый раз обратное распространение ошибки неэффективно и громоздко, ведь как мы знаем современные нейросетевые архитектуры могут состоять из сотен слоев и содержать тысячи операций в самых разных нетривиальных конфигурациях.\n",
    "\n",
    "Для этого в PyTorch реализована поддержка [автоматического дифференциирования в обратном режиме](https://en.wikipedia.org/wiki/Automatic_differentiation) (aka backprop). При использовании `torch.autograd`, выполнение прямого прохода вашей сети будет автоматически сохранять необходимые промежуточные результаты и записывать граф вычислительных зависимостей. Узлы в этом графе будут тензорами, а ребра - функциями, которые производят выходные тензоры из входных тензоров. Обратное распространение ошибки через этот граф затем позволяет легко вычислять градиенты.\n",
    "\n",
    "На практике использование данной технологии весьма простое: если мы хотим вычислить градиенты относительно некоторого тензора, тогда мы устанавливаем атрибут `requires_grad = True` при создании этого тензора или вызвав метод `tensor.requires_grad_()`. Любые дифференциируемые операции PyTorch с этим тензором будут автоматически возвращать тензоры, которые будут \"помнить\" свой вычислительный граф, что позволит нам позже выполнить обратное распространение ошибки через них. Если `x` - тензор с `requires_grad = True`, то после обратного распространения ошибки, `x.grad` будет другим тензором, содержащим градиент `x` относительно некоторого скалярного значения."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_demo = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)\n",
    "x_demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_demo = x_demo ** 2 + 3 * x_demo\n",
    "y_demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# requires_grad is \"viral\" by default:\n",
    "# the result of any differentiable operation involving a\n",
    "# tensor that requires a gradient will also require a gradient\n",
    "y_demo.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = y_demo.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# d(loss)/dx_demo = 2x + 3\n",
    "x_demo.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "2 * x_demo + 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_demo.grad # computed, but not saved"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Если нам наоборот не хочется сторить граф для выполнения обратного распространения ошибки через тензор с установленным атрибутом `requires_grad = True` (например, когда мы обновляем веса модели), то мы можем использовать контекстный менеджер `torch.no_grad()`, чтобы избежать построения вычислительного графа. В ситуации, когда у нас уже есть тензор `x`, являющийся частью графа, то можно получить копию данных этого тензора \"отщепленную\" от графа зависимостей выполнив операцию `x.detach()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    print(2 * x_demo + 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# .detach() returns a tensor with the same data but \"disconnected\" from the graph\n",
    "y_detached = y_demo.detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_demo + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_detached + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_detached.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# FYI: the underlying data is shared\n",
    "# (!) Don't do this (!)\n",
    "y_detached[0] = 100.0\n",
    "y_demo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Давайте реализуем обучние нейронной сети с использованием автоматического подсчета градиентов `autograd`.\n",
    "\n",
    "Больше не нужно вручную выводить формулы градиентов - достаточно описать forward-проход и вызвать `loss.backward()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create random Tensors for weights; setting requires_grad=True means that we\n",
    "# want to compute gradients for these Tensors during the backward pass.\n",
    "w1 = torch.empty((D_in, H), device=device, requires_grad=True)\n",
    "w2 = torch.empty((H, D_out), device=device, requires_grad=True)\n",
    "w1.data.uniform_(-scale1, scale1)\n",
    "w2.data.uniform_(-scale2, scale2)\n",
    "\n",
    "loss_hist = []\n",
    "progress = tqdm(range(num_steps))\n",
    "for t in progress:\n",
    "    # Forward pass: compute predicted y using operations on Tensors. Since w1 and\n",
    "    # w2 have requires_grad=True, operations involving these Tensors will cause\n",
    "    # PyTorch to build a computational graph, allowing automatic computation of\n",
    "    # gradients. Since we are no longer implementing the backward pass by hand we\n",
    "    # don't need to keep references to intermediate values.\n",
    "    p = (x @ w1).clip(0, None) @ w2\n",
    "\n",
    "    # Compute the loss. Loss must be a scalar tensor\n",
    "    loss = ((p - y) ** 2).sum()\n",
    "\n",
    "    # Save the loss value to history and display it in the progress bar\n",
    "    # CAUTION! Don't forget that the loss tensor lives on the device\n",
    "    # and keeps a reference to the whole computation graph!\n",
    "    if False:\n",
    "        # WRONG! This will leak memory!\n",
    "        loss_hist.append(loss)\n",
    "\n",
    "    loss_hist.append(loss.detach())\n",
    "    progress.desc = f\"loss: {loss.item():.8f}\"\n",
    "\n",
    "    # Use autograd to compute the backward pass. This call will compute the\n",
    "    # gradient of loss with respect to all Tensors with requires_grad=True.\n",
    "    # After this call w1.grad and w2.grad will be Tensors holding the gradient\n",
    "    # of the loss with respect to w1 and w2 respectively.\n",
    "    loss.backward()\n",
    "\n",
    "    # Update weights using gradient descent. For this step we just want to mutate\n",
    "    # the values of w1 and w2 in-place; we don't want to build up a computational\n",
    "    # graph for the update steps, so we use the torch.no_grad() context manager\n",
    "    # to prevent PyTorch from building a computational graph for the updates.\n",
    "    with torch.no_grad():\n",
    "        w1 -= learning_rate * w1.grad\n",
    "        w2 -= learning_rate * w2.grad\n",
    "\n",
    "    # CAUTION! Autograd loss.backward() ACCUMULATES the gradients, it effectively\n",
    "    # does `w.grad += ...` for each tensor with requires_grad=True in the graph.\n",
    "    # You need to manually clear the gradients between backward passes!\n",
    "    if False:\n",
    "        # Almost equivalent, but can be suboptimal in terms of performance\n",
    "        w1.grad.zero_()\n",
    "        w2.grad.zero_()\n",
    "\n",
    "    w1.grad = None\n",
    "    w2.grad = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy-paste and edit the inference code from previous section\n",
    "with torch.no_grad():\n",
    "\n",
    "    # Don't forget to wrap the forward computation in torch.no_grad()!\n",
    "    p = (x @ w1).clip(0, None) @ w2\n",
    "\n",
    "y_vis = y.cpu()\n",
    "p_vis = p.cpu()\n",
    "loss_hist = [l.cpu() for l in loss_hist]\n",
    "\n",
    "visualize_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch `nn`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Построение вычислительного графа и `autograd` уже являются очень мощными инструментами для разработки моделей глубокого обучения, но тем не менее для полноценной работы с нейронными сетями являются слишком низкоуровневыми.\n",
    "\n",
    "В PyTorch реализован еще один уровень абстракции в виде пакета `nn`, который позволяет не задумываться про написание часто используемых нейросетевых функций, а использовать уже готовые.\n",
    "\n",
    "Пакет `nn` определяет набор модулей, которые примерно эквивалентны слоям нейронной сети. Модуль получает входные тензоры и вычисляет выходные тензоры, но также может хранить внутреннее состояние, например, тензоры, содержащие обучаемые параметры."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Custom PyTorch Modules\n",
    "\n",
    "`torch.nn.Module` - это базовый класс, от которого должны наследоваться все модули. В данном классе уже за вас реализована логика необходимая для хранения состояния нейросети (например, обучаемых параметров).\n",
    "\n",
    "Реализуя свой подкласс `Module` необходимо определить метод `forward`, который принимает входные тензоры и создает выходные тензоры, используя другие модули или `autograd`-дифференцируемые операции с тензорами.\n",
    "\n",
    "Давайте напишем тот же самый двухслойный перцептрон, что и раньше, но без готовых слоёв - только с явными параметрами:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TwoLayerNet(torch.nn.Module):\n",
    "    def __init__(self, D_in, H, D_out):\n",
    "        \"\"\"\n",
    "        Create the parameters / buffers in the constructor,\n",
    "        wrap them in torch.nn.Parameter / torch.nn.Buffer,\n",
    "        assign them as member variables and initialize them.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        # nn.Parameter is a special object: assigning it to an attribute\n",
    "        # registers it as a trainable parameter of this module\n",
    "        self.w1 = torch.nn.Parameter(torch.empty(D_in, H))\n",
    "        self.w2 = torch.nn.Parameter(torch.empty(H, D_out))\n",
    "        self.w1.data.uniform_(-scale1, scale1)\n",
    "        self.w2.data.uniform_(-scale2, scale2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Copy the forward code from the train loop in previous section\n",
    "        \"\"\"\n",
    "        return (x @ self.w1).clip(0, None) @ self.w2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct our model by instantiating the class defined above.\n",
    "model = TwoLayerNet(D_in, H, D_out).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Built-in `nn` modules\n",
    "\n",
    "Писать каждый слой руками не обязательно: в `nn` уже есть готовые модули.\n",
    "\n",
    "- `nn.Linear(D_in, H)` - линейный слой, который сам создаёт и хранит свои веса (и, если нужно, смещение);\n",
    "- `nn.ReLU()` - модуль без параметров, который применяет нелинейность (то же самое, что `.clip(0, None)`);\n",
    "- `nn.Sequential(...)` - контейнер, который применяет модули друг за другом.\n",
    "\n",
    "Соберём ту же сеть из готовых блоков:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the nn package to define our model as a sequence of layers.\n",
    "# After constructing the model we use the .to() method to move it\n",
    "# (and all of its constituent parameters) to the desired device.\n",
    "#\n",
    "# Note: Disable the bias in Linear layers here for consistency\n",
    "# with previous simplified, manually implemented examples.\n",
    "model = torch.nn.Sequential(\n",
    "    torch.nn.Linear(D_in, H, bias=False),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(H, D_out, bias=False),\n",
    ").to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Лайфхак: в ситуации, когда `forward` метод состоит из последовательного вызова слоев определенных в `__init__`, можно вместо `torch.nn.Module` отнаследоваться от `torch.nn.Sequential` и получить дефолтную реализацию `forward`.\n",
    "\n",
    "Данный подход сочитает в себе плюсы создания `torch.nn.Sequential` из списка/словаря и наследования от `torch.nn.Module`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TwoLayerNet(torch.nn.Sequential):\n",
    "    def __init__(self, D_in, H, D_out):\n",
    "        super().__init__()\n",
    "        self.linear1 = torch.nn.Linear(D_in, H)\n",
    "        self.relu = torch.nn.ReLU()\n",
    "        self.linear2 = torch.nn.Linear(H, D_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct our model by instantiating the class defined above.\n",
    "model = TwoLayerNet(D_in, H, D_out).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Пакет `nn` также определяет набор полезных функций потерь, которые обычно используются при обучении нейронных сетей."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The nn package also contains definitions of popular loss functions;\n",
    "# in this case we will use Mean Squared Error (MSE) as our loss function.\n",
    "#\n",
    "# Note: Pass reduction=\"sum\" in order to compute the *sum* of squared\n",
    "# errors rather than the mean; this is for consistency with the previous\n",
    "# simplified examples. In practice it is more common to use mean squared\n",
    "# error as a loss by using reduction=\"mean\" (which is the default).\n",
    "loss_fn = torch.nn.MSELoss(reduction=\"sum\").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Switch model into training mode\n",
    "model = model.train()\n",
    "\n",
    "# Copy-paste and edit the train loop from previous section\n",
    "loss_hist = []\n",
    "progress = tqdm(range(num_steps))\n",
    "for t in progress:\n",
    "\n",
    "    # Forward pass and Loss computation:\n",
    "    # - Module objects override the __call__\n",
    "    #     operator so you can call them like functions.\n",
    "    # - The torch.nn.*Loss functions are also just modules that\n",
    "    #     take two arguments - predicted values and ground truth.\n",
    "    p = model(x)\n",
    "    loss = loss_fn(p, y)\n",
    "\n",
    "    loss_hist.append(loss.detach())\n",
    "    progress.desc = f\"loss: {loss.item():.8f}\"\n",
    "    loss.backward()\n",
    "\n",
    "    # Iterate over all model parameters and update them.\n",
    "    with torch.no_grad():\n",
    "        for param in model.parameters():\n",
    "            param.data -= learning_rate * param.grad\n",
    "\n",
    "    # Zero the gradients for all parameters in the model.\n",
    "    model.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy-paste and edit the inference code from previous section\n",
    "\n",
    "# Don't forget to switch the model into evaluation mode\n",
    "model = model.eval()\n",
    "with torch.no_grad():\n",
    "    p = model(x)\n",
    "\n",
    "y_vis = y.cpu()\n",
    "p_vis = p.cpu()\n",
    "loss_hist = [l.cpu() for l in loss_hist]\n",
    "\n",
    "visualize_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Control Flow and Weight Sharing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В качестве примера динамических графов и распределения весов мы реализуем очень странную модель: `fully connected` сеть с `ReLU`, которая на каждом прямом проходе выбирает случайное число от 1 до 4 и использует соотвествующее количество скрытых слоев, многократно используя одни и те же веса для скрытых слоев.\n",
    "\n",
    "Мы можем реализовать распределение веса между скрытыми слоями, просто повторно используя один и тот же модуль несколько раз при прямом проходе."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DynamicNet(torch.nn.Module):\n",
    "    def __init__(self, D_in, H, D_out):\n",
    "        \"\"\"\n",
    "        In the constructor we construct three nn.Linear instances\n",
    "        a trainable parameter tensor and a buffer tensor that we\n",
    "        will use in the forward pass.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.input_linear = torch.nn.Linear(D_in, H)\n",
    "        self.middle_linear = torch.nn.Linear(H, H)\n",
    "        self.output_linear = torch.nn.Linear(H, D_out)\n",
    "\n",
    "        # Important notes:\n",
    "        if False:\n",
    "            # Parameters -> part of model, updated by optimizer\n",
    "            self.trainable_values = torch.nn.Parameter(torch.ones(D_in))\n",
    "            # Buffers -> part of model, NOT updated by optimizer\n",
    "            self.fixed_values = torch.nn.Buffer(torch.ones(D_in))\n",
    "            # Regular tensors -> NOT \"registered\" as part of the model at all\n",
    "            self.not_registered = torch.ones(D_in)\n",
    "\n",
    "            # Layers / Parameters / Buffers are automatically \"registered\"\n",
    "            # when they are assigned directly to attributes on `self`.\n",
    "\n",
    "            # So stuff like this doesn't work:\n",
    "            self.layer_list = [torch.nn.Linear(1, 1)]\n",
    "            self.layer_dict = {\"foo\": torch.nn.Linear(1, 1)}\n",
    "            # Instead, use\n",
    "            self.layer_list = torch.nn.ModuleList([...])\n",
    "            self.layer_dict = torch.nn.ModuleDict({...})\n",
    "            # Or register the modules / parameters / buffers manually:\n",
    "            self.register_module(...)\n",
    "            self.register_buffer(...)\n",
    "            self.register_parameter(...)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        We can use Modules defined in the constructor as\n",
    "        well as arbitrary (differentiable) operations on Tensors.\n",
    "\n",
    "        For the forward pass of the model, we randomly choose either 1, 2, 3, or 4\n",
    "        and reuse the middle_linear Module that many times to compute hidden layer\n",
    "        representations.\n",
    "\n",
    "        Since each forward pass builds a dynamic computation graph, we can use normal\n",
    "        Python control-flow operators like loops or conditional statements when\n",
    "        defining the forward pass of the model.\n",
    "\n",
    "        Here we also see that it is perfectly safe to reuse the same Module many\n",
    "        times when defining a computational graph.\n",
    "        \"\"\"\n",
    "        # Note: NEVER manually call mod.forward(...) on Modules!\n",
    "        # Instead, you should ALWAYS call mod(...) directly.\n",
    "        h_relu = self.input_linear(x)\n",
    "        h_relu = h_relu.clip(0, None)\n",
    "\n",
    "        if self.training:\n",
    "            # Choose random depth during training\n",
    "            num_repeats = 1 + np.random.randint(4)\n",
    "        else:\n",
    "            # Use the maximum depth during evaluation\n",
    "            num_repeats = 3\n",
    "\n",
    "        # Repeat the middle_linear layer num_repeats times\n",
    "        for _ in range(num_repeats):\n",
    "            h_relu = self.middle_linear(h_relu)\n",
    "            h_relu = h_relu.clip(0, None)\n",
    "\n",
    "        p = self.output_linear(h_relu)\n",
    "        return p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct our model by instantiating the class defined above.\n",
    "model = DynamicNet(D_in, H, D_out).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy-paste the train loop from previous section\n",
    "model = model.train()\n",
    "loss_hist = []\n",
    "progress = tqdm(range(num_steps))\n",
    "for t in progress:\n",
    "    p = model(x)\n",
    "    loss = loss_fn(p, y)\n",
    "    loss_hist.append(loss.detach())\n",
    "    progress.desc = f\"loss: {loss.item():.8f}\"\n",
    "    loss.backward()\n",
    "    with torch.no_grad():\n",
    "        for param in model.parameters():\n",
    "            param.data -= learning_rate * param.grad\n",
    "    model.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy-paste the inference code from previous section\n",
    "model = model.eval()\n",
    "with torch.no_grad():\n",
    "    p = model(x)\n",
    "\n",
    "y_vis = y.cpu()\n",
    "p_vis = p.cpu()\n",
    "loss_hist = [l.cpu() for l in loss_hist]\n",
    "\n",
    "visualize_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch `optim`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "До этого момента мы обновляли веса наших моделей вручную изменяя тензоры, содержащие обучаемые параметры. Это не проблема для простых алгоритмов оптимизации, таких как стохастический градиентный спуск, но на практике мы часто обучаем нейронные сети, используя более сложные оптимизаторы, такие как `AdaGrad`, `RMSProp`, `Adam` и т. д.\n",
    "\n",
    "Пакет `optim` в PyTorch абстрагирует алгоритмы оптимизации и предоставляет реализации часто используемых из них.\n",
    "\n",
    "В этом примере мы будем использовать пакет `nn` для определения нашей модели, как и раньше, но мы будем оптимизировать модель, используя алгоритм `SGD`, предоставляемый пакетом `optim`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy-paste model and loss function definition from previous section\n",
    "model = TwoLayerNet(D_in, H, D_out).to(device)\n",
    "loss_fn = torch.nn.MSELoss(reduction=\"sum\").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the optim package to define an Optimizer that will update the weights of\n",
    "# the model for us. Here we will use SGD; the optim package contains many other\n",
    "# optimization algorithms. The first argument to the Adam constructor tells the\n",
    "# optimizer which Tensors it should update.\n",
    "#\n",
    "# Note: The best/recommended hyperparameters may vary for different optimizers.\n",
    "learning_rate = 1e-6\n",
    "num_steps = 50_000\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)\n",
    "\n",
    "# Copy-paste and edit the train loop from previous section\n",
    "model = model.train()\n",
    "loss_hist = []\n",
    "progress = tqdm(range(num_steps))\n",
    "for t in progress:\n",
    "    p = model(x)\n",
    "    loss = loss_fn(p, y)\n",
    "    loss_hist.append(loss.detach())\n",
    "    progress.desc = f\"loss: {loss.item():.8f}\"\n",
    "    loss.backward()\n",
    "\n",
    "    # Calling the step function on an Optimizer performs\n",
    "    # an update to all its tracked parameters\n",
    "    optimizer.step()\n",
    "\n",
    "    # Use the optimizer object to zero all of the\n",
    "    # gradients for the Tensors it will update\n",
    "    optimizer.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy-paste the inference code from previous section\n",
    "model = model.eval()\n",
    "with torch.no_grad():\n",
    "    p = model(x)\n",
    "\n",
    "y_vis = y.cpu()\n",
    "p_vis = p.cpu()\n",
    "loss_hist = [l.cpu() for l in loss_hist]\n",
    "\n",
    "visualize_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Что дальше\n",
    "\n",
    "Во всех примерах этого семинара мы обучались на одной и той же маленькой фиксированной выборке - это был градиентный спуск в \"классической\" формулировке. На практике выборки на много порядков больше, и посчитать градиент сразу по всей выборке невозможно.\n",
    "\n",
    "В следующем семинаре мы разберём, как устроена загрузка данных в PyTorch: абстракции `Dataset` и `DataLoader`, стохастический и пакетный (minibatch) градиентный спуск и работу с данными, которые целиком не помещаются в память."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus: Custom PyTorch `autograd` Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Под капотом каждый примитивный оператор `autograd` - это на самом деле две функции, которые работают с тензорами. Функция `forward` вычисляет выходные тензоры из входных тензоров. Функция `backward` получает градиент выходных тензоров относительно некоторого скалярного значения и вычисляет градиент входных тензоров относительно того же скалярного значения.\n",
    "\n",
    "В PyTorch мы можем легко определить наш собственный оператор `autograd`, определив подкласс `torch.autograd.Function` и реализовав на нем статические функции `forward` и `backward`. Затем мы можем использовать наш новый оператор, вызвая функцию `.apply` на этом подкласее, передавая в неё произвольные аргументы, которые будут затем переданы в наш метод `forward`.\n",
    "\n",
    "В этом примере мы определяем нашу собственную `autograd` функцию для выполнения нелинейности `ReLU` и используем ее в реализации нашей сети:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyReLU(torch.autograd.Function):\n",
    "    \"\"\"\n",
    "    We can implement our own custom autograd Functions by subclassing\n",
    "    torch.autograd.Function and implementing the forward and backward passes\n",
    "    which operate on Tensors.\n",
    "    \"\"\"\n",
    "\n",
    "    @staticmethod\n",
    "    def forward(ctx, x):\n",
    "        \"\"\"\n",
    "        In the forward pass we receive a context object and one or more input\n",
    "        Tensors; we must return a single Tensor (or a tuple of Tensors)\n",
    "        containing the output(s). We can use the context object to cache\n",
    "        tensors and arbitrary python objects for use in the backward pass.\n",
    "        \"\"\"\n",
    "        if False:\n",
    "            # You can save `x` here and compute the mask in backward\n",
    "            ctx.save_for_backward(x)\n",
    "\n",
    "        # Or you can precompute the boolean mask ahead of time and save a bit of\n",
    "        # memory (by only storing a boolean mask instead of the full float tensor)\n",
    "        mask = x > 0\n",
    "        ctx.save_for_backward(mask)\n",
    "        output = x.clip(0, None)\n",
    "        return output\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad_output):\n",
    "        \"\"\"\n",
    "        In the backward pass we receive the context object and Tensors containing\n",
    "        the gradient of the loss with respect to each of the outputs produced during\n",
    "        the forward pass. We can retrieve cached data from the context object, and\n",
    "        must compute and return the gradient of the loss with respect to the inputs\n",
    "        originally passed to the forward function (or None, if no gradient wrt that\n",
    "        input exists).\n",
    "        \"\"\"\n",
    "        if False:\n",
    "            [x] = ctx.saved_tensors\n",
    "            mask = x > 0\n",
    "\n",
    "        [mask] = ctx.saved_tensors\n",
    "        jacobian_output_x = mask.float()\n",
    "        grad_x = jacobian_output_x * grad_output\n",
    "        return grad_x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check that our function behaves like the built-in relu\n",
    "t = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0], requires_grad=True)\n",
    "\n",
    "for name, impl in [\n",
    "    (\"MyReLU\", MyReLU.apply),\n",
    "    (\"clip\", lambda t: t.clip(0.0, None)),\n",
    "    (\"torch\", torch.nn.functional.relu),\n",
    "]:\n",
    "    print(\"#\", name)\n",
    "    t.grad = None\n",
    "    out = impl(t)\n",
    "    print(\"output:\", out.detach())\n",
    "    out.mean().backward()\n",
    "    print(\"t.grad:\", t.grad)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Чтобы ближе ознакомиться с `autograd` вам предлагаемся самим написать функцию активации `SoftPlus`: $\\log(1 + \\exp(x))$.\n",
    "\n",
    "\n",
    "\n",
    "![](https://courses.cv-gml.ru/storage/seminars/nn-training-numpy-pytorch/activation-softplus.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def my_softplus(x):\n",
    "    exp = x.exp()\n",
    "    output = exp.log1p()\n",
    "\n",
    "    if False:\n",
    "        # 1) exp(x) may overflow even for relatively small x\n",
    "        return output\n",
    "\n",
    "    if False:\n",
    "        # 2) This is better, but clipping changes the function\n",
    "        exp = exp.clip(None, 1e6)\n",
    "        output = exp.log1p()\n",
    "        return output\n",
    "\n",
    "    # 3) This still produces nan gradients\n",
    "    mask = exp > 1e6\n",
    "    output[mask] = x[mask]\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Подсказка\n",
    "Производная `SoftPlus` имеет вид: $\\frac{\\exp(x)}{1 + \\exp(x)}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MySoftPlus(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, x):\n",
    "        # Same as my_softplus\n",
    "        exp = x.exp()\n",
    "        mask = exp > 1e6\n",
    "        ctx.save_for_backward(exp, mask)\n",
    "\n",
    "        # softplus(x) == x for sufficiently large x\n",
    "        output = exp.log1p()\n",
    "        output[mask] = x[mask]\n",
    "        return output\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad_output):\n",
    "        [exp, mask] = ctx.saved_tensors\n",
    "\n",
    "        if False:\n",
    "            # 4) After exp(x) overflows, inf/(1 + inf) is NaN\n",
    "            grad_x = exp / (1 + exp)\n",
    "            return grad_x\n",
    "\n",
    "        if False:\n",
    "            # 5) We were asked to compute dL/dx, not doutput/dx\n",
    "            grad_x = exp / (1 + exp)\n",
    "            grad_x[mask] = 1\n",
    "            return grad_x\n",
    "\n",
    "        jacobian_output_x = exp / (1 + exp)\n",
    "        jacobian_output_x[mask] = 1\n",
    "        grad_x = jacobian_output_x * grad_output\n",
    "        return grad_x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Протестируйте `SoftPlus`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check that our function behaves like the built-in softplus\n",
    "t = torch.tensor([-100.0, -2.0, 0.0, 2.0, 100.0], requires_grad=True)\n",
    "\n",
    "for name, impl in [\n",
    "    (\"torch\", torch.nn.functional.softplus),\n",
    "    (\"manual\", my_softplus),\n",
    "    (\"MySoftPlus\", MySoftPlus.apply),\n",
    "]:\n",
    "    print(\"#\", name)\n",
    "    t.grad = None\n",
    "    out = impl(t)\n",
    "    print(\"output:\", out.detach())\n",
    "    out.mean().backward()\n",
    "    print(\"t.grad:\", t.grad)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp_test = torch.from_numpy(np.load(\"test/inp.npy\")).to(device)\n",
    "inp_test.requires_grad = True\n",
    "\n",
    "out_pred = MySoftPlus.apply(inp_test)\n",
    "loss = out_pred.sum()\n",
    "loss.backward()\n",
    "\n",
    "out_test = np.load(\"test/out.npy\")\n",
    "inp_grad_test = np.load(\"test/inp_grad.npy\")\n",
    "assert np.allclose(\n",
    "    out_pred.detach().cpu().numpy(), out_test\n",
    "), \"Most likely, the forward pass implementation is incorrect\"\n",
    "assert np.allclose(\n",
    "    inp_test.grad.cpu().numpy(), inp_grad_test\n",
    "), \"Most likely, the backward pass implementation is incorrect\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Попробуйем обучить что-нибудь с ипользованием данной функции"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For convenience, define a MySoftPlusModule(torch.nn.Module)\n",
    "class MySoftPlusModule(torch.nn.Module):\n",
    "    def forward(self, x):\n",
    "        return MySoftPlus.apply(x)\n",
    "\n",
    "\n",
    "# Use MySoftPlusModule in TwoLayerNet\n",
    "class TwoLayerNet(torch.nn.Sequential):\n",
    "    def __init__(self, D_in, H, D_out):\n",
    "        super().__init__()\n",
    "        self.linear1 = torch.nn.Linear(D_in, H)\n",
    "        self.softplus = MySoftPlusModule()\n",
    "        self.linear2 = torch.nn.Linear(H, D_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy-paste model, loss function and optimizer definition from previous section\n",
    "model = TwoLayerNet(D_in, H, D_out).to(device)\n",
    "loss_fn = torch.nn.MSELoss(reduction=\"sum\").to(device)\n",
    "\n",
    "# Let's also try to adjust the optimizer hyperparameters a bit\n",
    "learning_rate = 1e-4  # was 1e-6\n",
    "num_steps = 5_000  # was 50_000\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy-paste the train loop from previous section\n",
    "model = model.train()\n",
    "loss_hist = []\n",
    "progress = tqdm(range(num_steps))\n",
    "for t in progress:\n",
    "    p = model(x)\n",
    "    loss = loss_fn(p, y)\n",
    "    loss_hist.append(loss.detach())\n",
    "    progress.desc = f\"loss: {loss.item():.8f}\"\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy-paste the inference code from previous section\n",
    "model = model.eval()\n",
    "with torch.no_grad():\n",
    "    p = model(x)\n",
    "\n",
    "y_vis = y.cpu()\n",
    "p_vis = p.cpu()\n",
    "loss_hist = [l.cpu() for l in loss_hist]\n",
    "\n",
    "visualize_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 1
}