03_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "descending-sleeping",
   "metadata": {},
   "source": [
    "# Deep Learning with PyTorch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "unexpected-guest",
   "metadata": {},
   "source": [
    "## Tensors\n",
    "*Tensors* represent\n",
    "* *multilinear maps* between vector spaces (mathematics)\n",
    "* generic *n-dimensional arrays* (computer science)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fuzzy-millennium",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.3556e-19, 3.0097e+29],\n",
       "        [7.1853e+22, 4.5145e+27],\n",
       "        [1.8040e+28, 1.5769e-19]])"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "# Create an uninitialized 3x2 tensor of 32-bit floats\n",
    "a = torch.FloatTensor(3, 2)\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "exposed-anatomy",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0., 0.],\n",
       "        [0., 0.],\n",
       "        [0., 0.]])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize the tensor (in-place) with zeros\n",
    "a.zero_()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "painful-advocacy",
   "metadata": {},
   "source": [
    "There are two types of methods in the PyTorch API:\n",
    "* Functional ones that return transformed copies have standard names like `some_function()`\n",
    "* In-place mutating operations will have a trailing underscore in their name, e.g. `some_function_()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "mounted-bulgarian",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 2., 3.],\n",
       "        [3., 2., 1.]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create a tensor from a standard collection\n",
    "torch.FloatTensor([[1, 2, 3], [3, 2, 1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "verified-cause",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0., 0.],\n",
       "       [0., 0.],\n",
       "       [0., 0.]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n = np.zeros(shape=(3, 2))\n",
    "n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "satisfactory-boxing",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0., 0.],\n",
       "        [0., 0.],\n",
       "        [0., 0.]], dtype=torch.float64)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create a tensor from a numpy ndarray\n",
    "b = torch.tensor(n)\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "composed-velvet",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0., 0.],\n",
       "        [0., 0.],\n",
       "        [0., 0.]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Change the numpy array to 64-bit float\n",
    "#  - This translates to the tensor\n",
    "n = np.zeros(shape=(3, 2), dtype=np.float32)\n",
    "torch.tensor(n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "loose-basketball",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0., 0.],\n",
       "        [0., 0.],\n",
       "        [0., 0.]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Alternatively specify a PyTorch dtype\n",
    "n = np.zeros(shape=(3, 2))\n",
    "torch.tensor(n, dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "separated-click",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1, 2, 3])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = torch.tensor([1, 2, 3])\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "pointed-flesh",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(6)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Scalar tensors can be results of some aggregations\n",
    "s = a.sum()\n",
    "s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "infrared-header",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# There's a convenient method to access the value of a scalar tensor\n",
    "s.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "documented-wellington",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(42)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.tensor(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "complimentary-crack",
   "metadata": {},
   "source": [
    "### Tensor Operations\n",
    "Each tensor has associated `device` where the computation takes place. The options are\n",
    "* `cpu` - computation takes place on the CPU\n",
    "* `cuda` or `cuda:<index>` - computation takes place on the GPU (with a device id `<index>`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "hairy-rendering",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([2., 3.])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Determine computation device based on CUDA availability\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "a = torch.FloatTensor([2, 3])\n",
    "\n",
    "# Move the tensor to GPU (if there's CUDA available)\n",
    "a = a.to(device)\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "neural-popularity",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a.device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "animal-failing",
   "metadata": {},
   "source": [
    "### Tensors and Gradients\n",
    "Each tensor has following info related to automatic gradient computation:\n",
    "* `grad` is a property holding computed gradients (tensor of the same shape)\n",
    "* `is_leaf` is true if the tensor was constructed by a user and false if it's a result of a computation\n",
    "* `requires_grad` is true if the tensor requires gradients to be computed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "destroyed-paris",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(12., grad_fn=<SumBackward0>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Define some tensors\n",
    "#  - The first one requires gradients to be computed\n",
    "v1 = torch.tensor([1.0, 1.0], requires_grad=True)\n",
    "v2 = torch.tensor([2.0, 2.0])\n",
    "\n",
    "# Define a computational graph on these tensors\n",
    "#  - Notice: Result contains a function coputing the gradient.\n",
    "v_sum = v1 + v2\n",
    "v_res = (v_sum * 2).sum()\n",
    "v_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "hairy-belief",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(True, True)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v1.is_leaf, v2.is_leaf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "increased-flour",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(False, False)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v_sum.is_leaf, v_res.is_leaf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "fitted-tactics",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(True, False)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v1.requires_grad, v2.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "unavailable-patrick",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(True, True)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v_sum.requires_grad, v_res.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "immune-australia",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([2., 2.])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Calculate the gradients of our graph\n",
    "v_res.backward()\n",
    "\n",
    "# Show backpropagated gradients in v1\n",
    "v1.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "crazy-monroe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# v2 does not require any gradients so there's nothing\n",
    "v2.grad"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ambient-natural",
   "metadata": {},
   "source": [
    "## Neural Network Building Blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "chief-england",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.2939, -1.2834,  0.0977,  1.6714,  0.0031], grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch.nn as nn  # noqa\n",
    "\n",
    "# Construct a 2-to-5 dense layer with an implicit bias\n",
    "#  - Note: Weights of this layer are randomly initialized.\n",
    "dense = nn.Linear(2, 5)\n",
    "\n",
    "# Each PyTorch NN module acts as a callable\n",
    "inputs = torch.FloatTensor([1, 2])\n",
    "dense(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "extended-compression",
   "metadata": {},
   "source": [
    "Some important methods from PyTorch API:\n",
    "* `parameters()` returns an iterable of all trainable variables (those that require gradients)\n",
    "* `zero_grad()` initializes all gradients to zero\n",
    "* `to(device)` moves the computation to a device\n",
    "* `state_dict()` exports all weights to a dictionary for model serialization\n",
    "* `load_state_dict()` oppisite of the previous which imports weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "athletic-layout",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): Linear(in_features=2, out_features=5, bias=True)\n",
       "  (1): ReLU()\n",
       "  (2): Linear(in_features=5, out_features=20, bias=True)\n",
       "  (3): ReLU()\n",
       "  (4): Linear(in_features=20, out_features=10, bias=True)\n",
       "  (5): ReLU()\n",
       "  (6): Dropout(p=0.3, inplace=False)\n",
       "  (7): Softmax(dim=1)\n",
       ")"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Build a sequential model with\n",
    "#  - Three dense layers with ReLU activations\n",
    "#  - A dropout layer\n",
    "#  - And a softmax output over the feature dimension\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(2, 5),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(5, 20),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(20, 10),\n",
    "    nn.ReLU(),\n",
    "    nn.Dropout(p=0.3),\n",
    "    nn.Softmax(dim=1),\n",
    ")\n",
    "\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "compliant-antarctica",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.1137, 0.0886, 0.1400, 0.1144, 0.0886, 0.0886, 0.0886, 0.0886, 0.0886,\n",
       "         0.1005]], grad_fn=<SoftmaxBackward>)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Feed an input tensor through our sequential model\n",
    "#  - There's single instance in the input batch\n",
    "model(torch.FloatTensor([[1, 2]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "revolutionary-alloy",
   "metadata": {},
   "source": [
    "### Custom Layers\n",
    "Creating custom modules (layers) is as easy as inheriting from `nn.Module` class and implementing the `forward()` method. Every other instance of a module assigned to a field is automatically registered under this module.\n",
    "\n",
    "Note that the convention is to use the module as a callable - this is because the `Module` class does some extra work in the `__call__` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "still-birth",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MyModule(\n",
       "  (pipe): Sequential(\n",
       "    (0): Linear(in_features=2, out_features=5, bias=True)\n",
       "    (1): ReLU()\n",
       "    (2): Linear(in_features=5, out_features=20, bias=True)\n",
       "    (3): ReLU()\n",
       "    (4): Linear(in_features=20, out_features=3, bias=True)\n",
       "    (5): Dropout(p=0.3, inplace=False)\n",
       "    (6): Softmax(dim=1)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing import TypeVar  # noqa\n",
    "\n",
    "T = TypeVar(\"T\", bound=torch.Tensor)\n",
    "\n",
    "\n",
    "class MyModule(nn.Module):\n",
    "    \"\"\"Custom PyTorch module\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        n_inputs: int,\n",
    "        n_outputs: int,\n",
    "        dropout_prob: float = 0.3,\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "        # Build a sequential model\n",
    "        #  - Every field that is a Module is auto-discovered\n",
    "        self.pipe = nn.Sequential(\n",
    "            nn.Linear(n_inputs, 5),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(5, 20),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(20, n_outputs),\n",
    "            nn.Dropout(p=dropout_prob),\n",
    "            nn.Softmax(dim=1),\n",
    "        )\n",
    "\n",
    "    def forward(self, x: T) -> T:\n",
    "        # We must treat the sub-module as a callable!\n",
    "        return self.pipe(x)\n",
    "\n",
    "\n",
    "# Build an instance of this model and show its structure\n",
    "net = MyModule(n_inputs=2, n_outputs=3)\n",
    "net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "extreme-hello",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.3075, 0.3657, 0.3268]], grad_fn=<SoftmaxBackward>)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Feed an input batch to the model\n",
    "net(torch.FloatTensor([[2, 3]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "breathing-equation",
   "metadata": {},
   "source": [
    "### Loss Functions and Optimizers\n",
    "\n",
    "PyTorch includes standard set of loss functions and of course allows simple implementation of custom ones. Here's a short list of some loss function classes:\n",
    "* `nn.MSELoss` is the *mean squared error* typically used for regression problems\n",
    "* `nn.BCELoss` and `BCEWithLogits` are *binary cross-entropy* losses for binary classification problems - the former expects single probability value while the latter raw scores (usually preferable)\n",
    "* `nn.CrossEntropyLoss` and `nn.NLLLoss` for multi-class classification problems\n",
    "\n",
    "Similarly there is buch of traditional optimizers such as vanilla `SGD`, `RMSprop`, `Adagrad` or the popular `Adam`. Finally, here's a typical training loop in PyTorch.\n",
    "```python\n",
    "# Define model and loss function\n",
    "model = ...\n",
    "loss_fn = ...\n",
    "\n",
    "# Register all trainable parameters in an optimizer\n",
    "optimizer = optim.Adam(params=model.parameters(), ...)\n",
    "\n",
    "# Iterate over mini-batches of training data\n",
    "for X_train, y_trian in iterate_batches(data, batch_size=32)\n",
    "    \n",
    "    # Wrap examples and labels into tensors\n",
    "    X_train = torch.tensor(X_train)\n",
    "    y_trian = torch.tensor(y_trian)\n",
    "    \n",
    "    # Make predictions using the model\n",
    "    y_pred = model(X_train)\n",
    "    \n",
    "    # Compute model's prediction loss\n",
    "    loss = loss_fn(y_pred, y_train)\n",
    "    \n",
    "    # Compute gradients of the loss function w.r.t. all the weights\n",
    "    #  - The loss is just a computation graph over tensors in the model\n",
    "    #  - And because model's weights \"require gradients\" this calculates dL/dw\n",
    "    loss.backward()\n",
    "    \n",
    "    # Perform one gradient descent step using computed gradients\n",
    "    #  - Note: The optimizer has access to `grad` for registered `params`\n",
    "    optimizer.step()\n",
    "    \n",
    "    # Clear gradients for this step\n",
    "    #  - Alternatively this can be done as the beginning of a step\n",
    "    #  - Note: This is a convenience method for calling it on the model\n",
    "    optimizer.zero_grad()\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "white-silence",
   "metadata": {},
   "source": [
    "## Monitoring with TensorBoard(X)\n",
    "Following example shows how to log arbitrary metrics and investigate them with *TensorBoard*. Because TensorBoard expects data in *TensorFlow* format we use `tensorboardX` for easy integration (also because it's a dependecy of PyTorch Ignite that we'll use later).\n",
    "\n",
    "Example below computes values of few trigonometric functions for varying angles and stores the output in `runs/` director. For later view one can run TensorBoard with\n",
    "```bash\n",
    "tensorboard --log-dir runs\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "packed-fiber",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math  # noqa\n",
    "\n",
    "from tensorboardX import SummaryWriter  # noqa\n",
    "\n",
    "# Define some functions representing our metrics\n",
    "funcs = {\"sin\": math.sin, \"cos\": math.cos, \"tan\": math.tan}\n",
    "\n",
    "# Create tesorboardX writer\n",
    "#  - Note: The default output directory is './runs'\n",
    "#  - Note 2: By default each call creates new \"run\"\n",
    "with SummaryWriter() as writer:\n",
    "\n",
    "    # Register one metric per function\n",
    "    for name, f in funcs.items():\n",
    "\n",
    "        # Evaluate and record f on interval [-360, 360)\n",
    "        for angle in range(-360, 360):\n",
    "            val = f(angle * math.pi / 180)\n",
    "            writer.add_scalar(name, val, angle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "afraid-opportunity",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[01;34mruns\u001b[00m\n",
      "├── \u001b[01;34mMar10_17-31-44_mpc-xps\u001b[00m\n",
      "│   └── events.out.tfevents.1615393904.mpc-xps\n",
      "└── \u001b[01;34mMar10_17-31-47_mpc-xps\u001b[00m\n",
      "    └── events.out.tfevents.1615393907.mpc-xps\n",
      "\n",
      "2 directories, 2 files\n"
     ]
    }
   ],
   "source": [
    "!tree runs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "backed-while",
   "metadata": {},
   "source": [
    "## GAN on Atari Images\n",
    "Let's train a *Generative Adversarial Network (GAN)* on randomly sampled screenshots from three Atari games and collect training metrics for TensorBoard."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "twenty-waters",
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf runs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "hungry-ancient",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random  # noqa\n",
    "from typing import Any, Iterable, Sequence, Tuple  # noqa\n",
    "\n",
    "import cv2  # noqa\n",
    "import gym  # noqa\n",
    "import torch.optim as optim  # noqa\n",
    "import torchvision.utils as vutils  # noqa\n",
    "\n",
    "# setup logging\n",
    "log = gym.logger\n",
    "log.set_level(gym.logger.INFO)\n",
    "\n",
    "# Hyperparameters\n",
    "IMAGE_SIZE = 64\n",
    "LATENT_VECTOR_SIZE = 100\n",
    "DISCR_FILTERS = 64\n",
    "GENER_FILTERS = 64\n",
    "BATCH_SIZE = 16\n",
    "MAX_ITERS = 2_000\n",
    "LEARNING_RATE = 0.0001\n",
    "REPORT_PERIOD = 100\n",
    "SAVE_IMAGE_PERIOD = 1000\n",
    "\n",
    "\n",
    "class InputWrapper(gym.ObservationWrapper):\n",
    "    \"\"\"\n",
    "    Preprocessing of input numpy array:\n",
    "    1. resize image into predefined size\n",
    "    2. move color channel axis to a first place\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, *args: Any) -> None:\n",
    "        super().__init__(*args)\n",
    "        assert isinstance(self.observation_space, gym.spaces.Box)\n",
    "        box = self.observation_space\n",
    "        self.observation_space = gym.spaces.Box(\n",
    "            self.observation(box.low),\n",
    "            self.observation(box.high),\n",
    "            dtype=np.float32,\n",
    "        )\n",
    "\n",
    "    def observation(self, observation: np.ndarray) -> np.ndarray:\n",
    "        # Resize the image\n",
    "        new_obs = cv2.resize(observation, (IMAGE_SIZE, IMAGE_SIZE))\n",
    "        # Move the color dimension to the fromt\n",
    "        #  - PyTorch's conv. layers expect shape [channels, width, height]\n",
    "        new_obs = np.moveaxis(new_obs, 2, 0)\n",
    "        return new_obs.astype(np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "divine-choice",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "    \"\"\"\n",
    "    Converts an image into single number\n",
    "    representing the probability of positive class\n",
    "    (the image being real).\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, input_shape: Tuple[int, ...]) -> None:\n",
    "        super().__init__()\n",
    "        self.conv_pipe = nn.Sequential(\n",
    "            nn.Conv2d(\n",
    "                in_channels=input_shape[0],\n",
    "                out_channels=DISCR_FILTERS,\n",
    "                kernel_size=4,\n",
    "                stride=2,\n",
    "                padding=1,\n",
    "            ),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(\n",
    "                in_channels=DISCR_FILTERS,\n",
    "                out_channels=DISCR_FILTERS * 2,\n",
    "                kernel_size=4,\n",
    "                stride=2,\n",
    "                padding=1,\n",
    "            ),\n",
    "            nn.BatchNorm2d(DISCR_FILTERS * 2),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(\n",
    "                in_channels=DISCR_FILTERS * 2,\n",
    "                out_channels=DISCR_FILTERS * 4,\n",
    "                kernel_size=4,\n",
    "                stride=2,\n",
    "                padding=1,\n",
    "            ),\n",
    "            nn.BatchNorm2d(DISCR_FILTERS * 4),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(\n",
    "                in_channels=DISCR_FILTERS * 4,\n",
    "                out_channels=DISCR_FILTERS * 8,\n",
    "                kernel_size=4,\n",
    "                stride=2,\n",
    "                padding=1,\n",
    "            ),\n",
    "            nn.BatchNorm2d(DISCR_FILTERS * 8),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(\n",
    "                in_channels=DISCR_FILTERS * 8,\n",
    "                out_channels=1,\n",
    "                kernel_size=4,\n",
    "                stride=1,\n",
    "                padding=0,\n",
    "            ),\n",
    "            nn.Sigmoid(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x: T) -> T:\n",
    "        conv_out = self.conv_pipe(x)\n",
    "        return conv_out.view(-1, 1).squeeze(dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "concerned-yesterday",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    \"\"\"Deconvolves input vector into 3x64x64 image\"\"\"\n",
    "\n",
    "    def __init__(self, output_shape: Tuple[int, ...]) -> None:\n",
    "        super().__init__()\n",
    "        self.pipe = nn.Sequential(\n",
    "            nn.ConvTranspose2d(\n",
    "                in_channels=LATENT_VECTOR_SIZE,\n",
    "                out_channels=GENER_FILTERS * 8,\n",
    "                kernel_size=4,\n",
    "                stride=1,\n",
    "                padding=0,\n",
    "            ),\n",
    "            nn.BatchNorm2d(GENER_FILTERS * 8),\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(\n",
    "                in_channels=GENER_FILTERS * 8,\n",
    "                out_channels=GENER_FILTERS * 4,\n",
    "                kernel_size=4,\n",
    "                stride=2,\n",
    "                padding=1,\n",
    "            ),\n",
    "            nn.BatchNorm2d(GENER_FILTERS * 4),\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(\n",
    "                in_channels=GENER_FILTERS * 4,\n",
    "                out_channels=GENER_FILTERS * 2,\n",
    "                kernel_size=4,\n",
    "                stride=2,\n",
    "                padding=1,\n",
    "            ),\n",
    "            nn.BatchNorm2d(GENER_FILTERS * 2),\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(\n",
    "                in_channels=GENER_FILTERS * 2,\n",
    "                out_channels=GENER_FILTERS,\n",
    "                kernel_size=4,\n",
    "                stride=2,\n",
    "                padding=1,\n",
    "            ),\n",
    "            nn.BatchNorm2d(GENER_FILTERS),\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(\n",
    "                in_channels=GENER_FILTERS,\n",
    "                out_channels=output_shape[0],\n",
    "                kernel_size=4,\n",
    "                stride=2,\n",
    "                padding=1,\n",
    "            ),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x: T) -> T:\n",
    "        return self.pipe(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "matched-philadelphia",
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize(batch: np.ndarray) -> np.ndarray:\n",
    "    \"\"\"Normalize to interval [-1, 1]\"\"\"\n",
    "    return np.array(batch, dtype=np.float32) * 2.0 / 255.0 - 1.0\n",
    "\n",
    "\n",
    "def iterate_batches(\n",
    "    envs: Sequence[gym.Env],\n",
    "    batch_size: int = BATCH_SIZE,\n",
    ") -> Iterable[T]:\n",
    "    # Initialize a buffer for current batch\n",
    "    batch = [e.reset() for e in envs]\n",
    "\n",
    "    # Create a sampler over the environments\n",
    "    env_sampler = iter(lambda: random.choice(envs), None)\n",
    "\n",
    "    while True:\n",
    "\n",
    "        # Sample an action from random environment\n",
    "        env = next(env_sampler)\n",
    "        action = env.action_space.sample()\n",
    "\n",
    "        # Get new observation for the action\n",
    "        obs, _, done, _ = env.step(action)\n",
    "\n",
    "        # Just a hack to fix one of the envs\n",
    "        if np.mean(obs) > 0.01:\n",
    "            batch.append(obs)\n",
    "\n",
    "        # Generate new batch\n",
    "        #  - Normalize values to [-1, 1]\n",
    "        #  - Yield new tensor and clean the buffer\n",
    "        if len(batch) == batch_size:\n",
    "            yield torch.tensor(normalize(batch))\n",
    "            batch.clear()\n",
    "\n",
    "        # Restart environment when episode ends\n",
    "        if done:\n",
    "            env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "radio-spokesman",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: Making new env: Breakout-v0\n",
      "INFO: Making new env: AirRaid-v0\n",
      "INFO: Making new env: Pong-v0\n",
      "INFO: Iter: 0000/2000 Gen. Loss: 1.480 Dis. Loss: 1.415\n",
      "INFO: Iter: 0100/2000 Gen. Loss: 5.438 Dis. Loss: 0.041\n",
      "INFO: Iter: 0200/2000 Gen. Loss: 7.007 Dis. Loss: 0.005\n",
      "INFO: Iter: 0300/2000 Gen. Loss: 7.449 Dis. Loss: 0.002\n",
      "INFO: Iter: 0400/2000 Gen. Loss: 8.134 Dis. Loss: 0.084\n",
      "INFO: Iter: 0500/2000 Gen. Loss: 7.594 Dis. Loss: 0.014\n",
      "INFO: Iter: 0600/2000 Gen. Loss: 6.825 Dis. Loss: 0.006\n",
      "INFO: Iter: 0700/2000 Gen. Loss: 8.798 Dis. Loss: 0.005\n",
      "INFO: Iter: 0800/2000 Gen. Loss: 9.548 Dis. Loss: 0.120\n",
      "INFO: Iter: 0900/2000 Gen. Loss: 5.932 Dis. Loss: 0.011\n",
      "INFO: Iter: 1000/2000 Gen. Loss: 6.427 Dis. Loss: 0.011\n",
      "INFO: Iter: 1100/2000 Gen. Loss: 6.506 Dis. Loss: 0.006\n",
      "INFO: Iter: 1200/2000 Gen. Loss: 6.858 Dis. Loss: 0.004\n",
      "INFO: Iter: 1300/2000 Gen. Loss: 7.343 Dis. Loss: 0.002\n",
      "INFO: Iter: 1400/2000 Gen. Loss: 7.197 Dis. Loss: 0.003\n",
      "INFO: Iter: 1500/2000 Gen. Loss: 6.564 Dis. Loss: 0.366\n",
      "INFO: Iter: 1600/2000 Gen. Loss: 6.123 Dis. Loss: 0.074\n",
      "INFO: Iter: 1700/2000 Gen. Loss: 5.758 Dis. Loss: 0.012\n",
      "INFO: Iter: 1800/2000 Gen. Loss: 6.189 Dis. Loss: 0.008\n",
      "INFO: Iter: 1900/2000 Gen. Loss: 6.172 Dis. Loss: 0.006\n"
     ]
    }
   ],
   "source": [
    "# Create new gym environments for 3 Atari games\n",
    "games = (\"Breakout-v0\", \"AirRaid-v0\", \"Pong-v0\")\n",
    "envs = [InputWrapper(gym.make(name)) for name in games]\n",
    "\n",
    "input_shape = envs[0].observation_space.shape\n",
    "\n",
    "# Create GAN components\n",
    "discriminator = Discriminator(input_shape=input_shape).to(device)\n",
    "generator = Generator(output_shape=input_shape).to(device)\n",
    "\n",
    "# The objective is binary cross-entropy (binary classification)\n",
    "loss = nn.BCELoss()\n",
    "\n",
    "# Create two optimizers - one for each component\n",
    "gen_optimizer = optim.Adam(\n",
    "    params=generator.parameters(),\n",
    "    lr=LEARNING_RATE,\n",
    "    betas=(0.5, 0.999),\n",
    ")\n",
    "\n",
    "dis_optimizer = optim.Adam(\n",
    "    params=discriminator.parameters(),\n",
    "    lr=LEARNING_RATE,\n",
    "    betas=(0.5, 0.999),\n",
    ")\n",
    "\n",
    "# Create labels for true and fake images\n",
    "true_labels = torch.ones(BATCH_SIZE, device=device)\n",
    "fake_labels = torch.zeros(BATCH_SIZE, device=device)\n",
    "\n",
    "# Start recording of the training metrics for TensorBoard\n",
    "with SummaryWriter() as writer:\n",
    "\n",
    "    gen_losses = []\n",
    "    dis_losses = []\n",
    "\n",
    "    batch_iter = iter(iterate_batches(envs))\n",
    "\n",
    "    # Main training loop\n",
    "    for i in range(MAX_ITERS):\n",
    "        real_inputs = next(batch_iter)\n",
    "\n",
    "        # Move the input batch to selected device\n",
    "        real_inputs = real_inputs.to(device)\n",
    "\n",
    "        # Generate an equal-sized batch of codings for fake images\n",
    "        #  - Latent vectors are drawn from N(0, 1)\n",
    "        #  - And we move them to selected device too\n",
    "        codings = torch.FloatTensor(BATCH_SIZE, LATENT_VECTOR_SIZE, 1, 1)\n",
    "        codings.normal_(0, 1)\n",
    "        codings = codings.to(device)\n",
    "\n",
    "        # Generate fake images using the generator\n",
    "        fake_inputs = generator(codings)\n",
    "\n",
    "        # Train the discriminator\n",
    "        #  - Note: We must `detach` generator inputs to not alter its weights\n",
    "        dis_optimizer.zero_grad()\n",
    "        pred_true = discriminator(real_inputs)\n",
    "        pred_fake = discriminator(fake_inputs.detach())\n",
    "        dis_loss = loss(pred_true, true_labels) + loss(pred_fake, fake_labels)\n",
    "        dis_loss.backward()\n",
    "        dis_optimizer.step()\n",
    "        dis_losses.append(dis_loss.item())\n",
    "\n",
    "        # Train the generator\n",
    "        gen_optimizer.zero_grad()\n",
    "        dis_pred = discriminator(fake_inputs)\n",
    "        gen_loss = loss(dis_pred, true_labels)\n",
    "        gen_loss.backward()\n",
    "        gen_optimizer.step()\n",
    "        gen_losses.append(gen_loss.item())\n",
    "\n",
    "        # Collect training metrics\n",
    "        if i % REPORT_PERIOD == 0:\n",
    "\n",
    "            # Compute mean losses over buffered window\n",
    "            mean_gen_loss = np.mean(gen_losses)\n",
    "            mean_dis_loss = np.mean(dis_losses)\n",
    "\n",
    "            log.info(\n",
    "                \"Iter: %.4d/%d Gen. Loss: %.3f Dis. Loss: %.3f\",\n",
    "                i,\n",
    "                MAX_ITERS,\n",
    "                mean_gen_loss,\n",
    "                mean_dis_loss,\n",
    "            )\n",
    "\n",
    "            writer.add_scalar(\"loss_gen\", mean_gen_loss, i)\n",
    "            writer.add_scalar(\"loss_dis\", mean_dis_loss, i)\n",
    "\n",
    "            # Reset the loss buffers\n",
    "            gen_losses = []\n",
    "            dis_losses = []\n",
    "\n",
    "        # Collect real and fake images\n",
    "        if i % SAVE_IMAGE_PERIOD == 0:\n",
    "            real_img = vutils.make_grid(real_inputs.data[:64], normalize=True)\n",
    "            fake_img = vutils.make_grid(fake_inputs.data[:64], normalize=True)\n",
    "            writer.add_image(\"img_real\", real_img, i)\n",
    "            writer.add_image(\"img_fake\", fake_img, i)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pacific-christmas",
   "metadata": {},
   "source": [
    "Not great but we've trained the GAN for just few iterations (2k) and it suffices as an example. Now, let's re-implement this using a higher-level library called *PyTorch Ignite*."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "central-minority",
   "metadata": {},
   "source": [
    "## PyTorch Ignite\n",
    "*PyTorch Ignite* is a high-level library that abstracts over the training loop while adding means how to customize it. Customization is achieved via training *event listeners* and serves e.g. for collecting metrics, logging and integration with TensorBoard.\n",
    "\n",
    "Let's draw a basic template for a training script using Ignite.\n",
    "```python\n",
    "from ignite.engine import Engine, Events\n",
    "\n",
    "# Build a model and optimizer\n",
    "model = ...\n",
    "optimizer = ...\n",
    "loss_fn = ...\n",
    "\n",
    "def training(engine, batch):\n",
    "    \n",
    "    # Reset gradients\n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    # Preprocess data and make prediction\n",
    "    X, y = prepare_batch(batch)\n",
    "    y_predict = model(X)\n",
    "    \n",
    "    # Compute loss and backpropagate gradients\n",
    "    loss = loss_fn(y_predict, y)\n",
    "    loss.backward()\n",
    "    \n",
    "    # Make a gradient descent step\n",
    "    optimizer.step()\n",
    "    \n",
    "    # Return metrics for tracking\n",
    "    return loss.item()\n",
    "    \n",
    "# Crate a training engine and run it on training data\n",
    "engine = Engine(training)\n",
    "engine.run(data)\n",
    "```\n",
    "\n",
    "Before we re-create previous GAN example, let's mention that there are other alternatives to Ignite, for instance *Fast.ai* or *PyTorch Lightning*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "correct-middle",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: Making new env: Breakout-v0\n",
      "INFO: Making new env: AirRaid-v0\n",
      "INFO: Making new env: Pong-v0\n",
      "INFO: 0100: gen_loss=5.482, dis_loss=0.201\n",
      "INFO: 0200: gen_loss=7.005, dis_loss=0.030\n",
      "INFO: 0300: gen_loss=7.615, dis_loss=0.005\n",
      "INFO: 0400: gen_loss=8.061, dis_loss=0.003\n",
      "INFO: 0500: gen_loss=8.530, dis_loss=0.002\n",
      "INFO: 0600: gen_loss=8.026, dis_loss=0.387\n",
      "INFO: 0700: gen_loss=6.668, dis_loss=0.066\n",
      "INFO: 0800: gen_loss=6.670, dis_loss=0.014\n",
      "INFO: 0900: gen_loss=7.421, dis_loss=0.007\n",
      "INFO: 1000: gen_loss=7.147, dis_loss=0.098\n",
      "INFO: 1100: gen_loss=7.003, dis_loss=0.065\n",
      "INFO: 1200: gen_loss=6.425, dis_loss=0.228\n",
      "INFO: 1300: gen_loss=5.747, dis_loss=0.236\n",
      "INFO: 1400: gen_loss=5.344, dis_loss=0.282\n",
      "INFO: 1500: gen_loss=4.707, dis_loss=0.251\n",
      "INFO: 1600: gen_loss=4.430, dis_loss=0.320\n",
      "INFO: 1700: gen_loss=3.963, dis_loss=0.307\n",
      "INFO: 1800: gen_loss=4.844, dis_loss=0.199\n",
      "INFO: 1900: gen_loss=4.277, dis_loss=0.254\n",
      "INFO: 2000: gen_loss=4.935, dis_loss=0.117\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "State:\n",
       "\titeration: 2000\n",
       "\tepoch: 1\n",
       "\tepoch_length: 2000\n",
       "\tmax_epochs: 1\n",
       "\toutput: <class 'tuple'>\n",
       "\tbatch: <class 'torch.Tensor'>\n",
       "\tmetrics: <class 'dict'>\n",
       "\tdataloader: <class 'generator'>\n",
       "\tseed: <class 'NoneType'>\n",
       "\ttimes: <class 'dict'>\n",
       "\treal_img_data: <class 'torch.Tensor'>\n",
       "\tfake_img_data: <class 'torch.Tensor'>"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from ignite.contrib.handlers import tensorboard_logger as tb_logger  # noqa\n",
    "from ignite.engine import Engine, Events  # noqa\n",
    "from ignite.metrics import RunningAverage  # noqa\n",
    "\n",
    "# Create new gym environments for 3 Atari games\n",
    "games = (\"Breakout-v0\", \"AirRaid-v0\", \"Pong-v0\")\n",
    "envs = [InputWrapper(gym.make(name)) for name in games]\n",
    "\n",
    "input_shape = envs[0].observation_space.shape\n",
    "\n",
    "# Create GAN components\n",
    "discriminator = Discriminator(input_shape=input_shape).to(device)\n",
    "generator = Generator(output_shape=input_shape).to(device)\n",
    "\n",
    "# The objective is binary cross-entropy (binary classification)\n",
    "loss = nn.BCELoss()\n",
    "\n",
    "# Create two optimizers - one for each component\n",
    "gen_optimizer = optim.Adam(\n",
    "    params=generator.parameters(),\n",
    "    lr=LEARNING_RATE,\n",
    "    betas=(0.5, 0.999),\n",
    ")\n",
    "\n",
    "dis_optimizer = optim.Adam(\n",
    "    params=discriminator.parameters(),\n",
    "    lr=LEARNING_RATE,\n",
    "    betas=(0.5, 0.999),\n",
    ")\n",
    "\n",
    "# Create labels for true and fake images\n",
    "true_labels = torch.ones(BATCH_SIZE, device=device)\n",
    "fake_labels = torch.zeros(BATCH_SIZE, device=device)\n",
    "\n",
    "\n",
    "def training(engine: Engine, batch: T) -> Tuple[float, float]:\n",
    "    # Move the input batch to selected device\n",
    "    real_inputs = batch.to(device)\n",
    "\n",
    "    # Generate an equal-sized batch of codings for fake images\n",
    "    #  - Latent vectors are drawn from N(0, 1)\n",
    "    #  - And we move them to selected device too\n",
    "    codings = torch.FloatTensor(BATCH_SIZE, LATENT_VECTOR_SIZE, 1, 1)\n",
    "    codings.normal_(0, 1)\n",
    "    codings = codings.to(device)\n",
    "\n",
    "    # Generate fake images using the generator\n",
    "    fake_inputs = generator(codings)\n",
    "\n",
    "    # Train the discriminator\n",
    "    #  - Note: We must `detach` generator inputs to not alter its weights\n",
    "    dis_optimizer.zero_grad()\n",
    "    pred_true = discriminator(real_inputs)\n",
    "    pred_fake = discriminator(fake_inputs.detach())\n",
    "    dis_loss = loss(pred_true, true_labels) + loss(pred_fake, fake_labels)\n",
    "    dis_loss.backward()\n",
    "    dis_optimizer.step()\n",
    "\n",
    "    # Train the generator\n",
    "    gen_optimizer.zero_grad()\n",
    "    dis_pred = discriminator(fake_inputs)\n",
    "    gen_loss = loss(dis_pred, true_labels)\n",
    "    gen_loss.backward()\n",
    "    gen_optimizer.step()\n",
    "\n",
    "    # Add generated images to the state\n",
    "    engine.state.real_img_data = real_inputs.data[:64]\n",
    "    engine.state.fake_img_data = fake_inputs.data[:64]\n",
    "\n",
    "    # Return metrics for tracking\n",
    "    return dis_loss.item(), gen_loss.item()\n",
    "\n",
    "\n",
    "# Create an Ignite engine over the training loop\n",
    "#  - and attach a TensorBoard logger\n",
    "engine = Engine(training)\n",
    "engine.tb = tb_logger.TensorboardLogger(log_dir=None)\n",
    "\n",
    "# Register new fields for images in engine's state\n",
    "engine.state_dict_user_keys.append(\"real_img_data\")\n",
    "engine.state_dict_user_keys.append(\"fake_img_data\")\n",
    "\n",
    "# Create running average metrics for bogh losses\n",
    "\n",
    "avg_loss_gen = RunningAverage(output_transform=lambda out: out[1])\n",
    "avg_loss_gen.attach(engine, \"avg_loss_gen\")\n",
    "\n",
    "avg_loss_dis = RunningAverage(output_transform=lambda out: out[0])\n",
    "avg_loss_dis.attach(engine, \"avg_loss_dis\")\n",
    "\n",
    "# Create and attach handler for training lesses\n",
    "\n",
    "log_handler = tb_logger.OutputHandler(\n",
    "    tag=\"train\",\n",
    "    metric_names=[\"avg_loss_gen\", \"avg_loss_dis\"],\n",
    ")\n",
    "\n",
    "engine.tb.attach(\n",
    "    engine,\n",
    "    log_handler=log_handler,\n",
    "    event_name=Events.ITERATION_COMPLETED,\n",
    ")\n",
    "\n",
    "\n",
    "# Register a training progress logger\n",
    "#  - This callback will trigger every `REPORT_PERIOD`\n",
    "@engine.on(Events.ITERATION_COMPLETED(every=REPORT_PERIOD))\n",
    "def log_losses(engine: Engine) -> None:\n",
    "    log.info(\n",
    "        \"%.4d: gen_loss=%.3f, dis_loss=%.3f\",\n",
    "        engine.state.iteration,\n",
    "        engine.state.metrics[\"avg_loss_gen\"],\n",
    "        engine.state.metrics[\"avg_loss_dis\"],\n",
    "    )\n",
    "\n",
    "\n",
    "# Register an image writer\n",
    "#  - This callback will trigger every `SAVE_IMAGE_PERIOD`\n",
    "@engine.on(Events.ITERATION_COMPLETED(every=SAVE_IMAGE_PERIOD))\n",
    "def save_images(engine: Engine) -> None:\n",
    "    real_img = vutils.make_grid(engine.state.real_img_data, normalize=True)\n",
    "    fake_img = vutils.make_grid(engine.state.fake_img_data, normalize=True)\n",
    "    engine.tb.writer.add_image(\"img_real\", real_img, engine.state.iteration)\n",
    "    engine.tb.writer.add_image(\"img_fake\", fake_img, engine.state.iteration)\n",
    "    engine.tb.writer.flush()\n",
    "\n",
    "\n",
    "# Run the training loop\n",
    "engine.run(data=iterate_batches(envs), epoch_length=MAX_ITERS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "israeli-collection",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}