examples/notebooks/CycleGAN_with_torch_cuda_amp.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "CycleGAN_with_torch_cuda_amp.ipynb",
   "provenance": [],
   "collapsed_sections": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XwaQH08zFRZW"
   },
   "source": [
    "# CycleGAN with Ignite and `torch.cuda.amp`\n",
    "\n",
    "In this notebook we provide an implementation of [CycleGAN](https://arxiv.org/abs/1703.10593) and its training on \"Horse 2 Zebra\" dataset using Ignite. This notebook is almost similar to another our [notebook on CycleGAN with Nvidia/Apex](https://github.com/pytorch/ignite/blob/master/examples/notebooks/CycleGAN_with_ignite_and_nvdia_apex.ipynb).\n",
    "\n",
    "In contrast, we will use recently added [`torch.cuda.amp`](https://pytorch.org/docs/master/notes/amp_examples.html#working-with-multiple-models-losses-and-optimizers) module to perform automatic mixed precision training instead of using Nvidia/Apex package. This module is available in pytorch (>=1.6.0) release.\n",
    "\n",
    "\n",
    "### CycleGAN in a Nutshell\n",
    "\n",
    "CycleGAN is unpaired image-to-image translation task from $A$ to $B$ and represented by two generative networks $G$ and $F$:\n",
    "$$\n",
    "\\hat{y} = G(x) \\in B,\\text{ for } x \\in A \\\\\n",
    "\\hat{x} = F(y) \\in A,\\text{ for } y \\in B\n",
    "$$\n",
    "\n",
    "and two discriminators $D_A$ and $D_B$. Training of the networks is done by minimizing the loss is a sum of 3 components:\n",
    "$$\n",
    "\\mathcal{L}(G, F, D_A, D_B) = \\mathcal{L}_{GAN}(G, D_B) + \\mathcal{L}_{GAN}(F, D_A) + \\lambda \\mathcal{L}_{cyc}(G, F)\n",
    "$$\n",
    "with GAN loss:\n",
    "$$\n",
    "\\mathcal{L}_{GAN}(G, D_B) = \\text{mean}_{x \\in A}\\left[ (D_B(G(x)) - 1)^2 \\right]+ \\text{mean}_{y \\in B}\\left[ (D_B(y) - 1)^2 \\right] \\\\\n",
    "\\mathcal{L}_{GAN}(F, D_A) = \\text{mean}_{y \\in B}\\left[ (D_A(F(y)) - 1)^2 \\right]+ \\text{mean}_{x \\in A}\\left[ (D_A(x) - 1)^2 \\right]\n",
    "$$\n",
    "and forward and backward cycle consistency term:\n",
    "$$\n",
    "\\mathcal{L}_{cyc}(G, F) = \\text{mean}_{x \\in A}\\left[ |F(G(x)) - x|_1 \\right] + \\text{mean}_{y \\in B}\\left[ |G(F(y)) - y|_1 \\right]\n",
    "$$\n",
    "\n",
    "Optionally, one can add identity loss terms. See [here](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/qa.md#what-is-the-identity-loss-322-373-362)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "eqe1kXPcXj1U"
   },
   "source": [
    "!nvidia-smi"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6Y60e3m2FRZY"
   },
   "source": [
    "## Requirements\n",
    "\n",
    "1) Let's download the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "l3LdmHAuFRZa"
   },
   "source": [
    "!wget https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets/horse2zebra.zip -O/tmp/horse2zebra.zip\n",
    "!7z x /tmp/horse2zebra.zip -o/tmp/"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "quI5HJekFRZm"
   },
   "source": [
    "2) Next we need to install `torchvision` and PyTorch Nightly Release (>=1.6.0)\n",
    "\n",
    "```python\n",
    "# Torchvision:\n",
    "pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "3c6PHUZeFRZu"
   },
   "source": [
    "!pip install --upgrade --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DU14wTLeZBh3"
   },
   "source": [
    "3) We install nightly version of pytorch-ignite"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "dWN63EToZA-G"
   },
   "source": [
    "!pip install --pre pytorch-ignite"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "p8M6GlpmQ5jZ"
   },
   "source": [
    "import torch\n",
    "import ignite\n",
    "torch.__version__, ignite.__version__"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "d3LCqCufFRZ6"
   },
   "source": [
    "import random\n",
    "import torch\n",
    "\n",
    "seed = 17\n",
    "random.seed(seed)\n",
    "_ = torch.manual_seed(seed)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "20G39vEqFRaA"
   },
   "source": [
    "## Dataflow\n",
    "\n",
    "Let's setup the datase"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "nOr3nd4qFRaB"
   },
   "source": [
    "from torch.utils.data import Dataset, DataLoader\n",
    "from PIL import Image\n",
    "\n",
    "class FilesDataset(Dataset):\n",
    "    \n",
    "    def __init__(self, path, extension=\"*.jpg\"):\n",
    "        self.path = Path(path)\n",
    "        assert self.path.exists(), \"Path '{}' is not found\".format(path)\n",
    "        self.images = list(self.path.rglob(extension))\n",
    "        assert len(self.images) > 0, \"No images with extension {} found at '{}'\".format(extension, path)\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.images)\n",
    "    \n",
    "    def __getitem__(self, i):\n",
    "        return Image.open(self.images[i]).convert('RGB')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "byG73rBHFRaG"
   },
   "source": [
    "from pathlib import Path\n",
    "\n",
    "root = Path(\"/tmp/horse2zebra\")\n",
    "\n",
    "train_A = FilesDataset(root / \"trainA\")\n",
    "train_B = FilesDataset(root / \"trainB\")\n",
    "\n",
    "test_A = FilesDataset(root / \"testA\") \n",
    "test_B = FilesDataset(root / \"testB\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Kv1oXmSJFRaK"
   },
   "source": [
    "Some details on the datasets:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "UZ4rY5S9FRaL"
   },
   "source": [
    "print(\"Dataset sizes: \\ntrain A: {} | B: {}\\ntest A: {} | B: {}\\n\\t\".format(len(train_A), len(train_B), len(test_A), len(test_B)))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "R4lyZ6GSFRaR"
   },
   "source": [
    "print(\"Train random image sizes (A): {}, {}, {}, {}\".format(train_A[0].size, train_A[1].size, train_A[10].size, train_A[-1].size))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "e-61gbfVFRaX"
   },
   "source": [
    "print(\"Train random image sizes (B): {}, {}, {}, {}\".format(train_B[0].size, train_B[1].size, train_B[10].size, train_B[-1].size))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fAxkND8jFRac"
   },
   "source": [
    "import matplotlib.pylab as plt\n",
    "%matplotlib inline"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Xx7cAEI0FRah"
   },
   "source": [
    "plt.figure(figsize=(10, 5))\n",
    "plt.subplot(121)\n",
    "plt.title(\"Train dataset 'Horses'\")\n",
    "plt.imshow(train_A[10])\n",
    "plt.subplot(122)\n",
    "plt.title(\"Train dataset 'Zebras'\")\n",
    "plt.imshow(train_B[10])"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "S_v_FuvZFRao"
   },
   "source": [
    "plt.figure(figsize=(10, 5))\n",
    "plt.subplot(121)\n",
    "plt.title(\"Test dataset 'Horses'\")\n",
    "plt.imshow(test_A[0])\n",
    "plt.subplot(122)\n",
    "plt.title(\"Test dataset 'Zebras'\")\n",
    "plt.imshow(test_B[0])"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gc6ZuUWwFRas"
   },
   "source": [
    "Let's create a dataset composed of random image pairs of datasets A and B"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "B699THXXFRat"
   },
   "source": [
    "import random\n",
    "\n",
    "\n",
    "class Image2ImageDataset(Dataset):\n",
    "    \n",
    "    def __init__(self, ds_a, ds_b):\n",
    "        self.dataset_a = ds_a\n",
    "        self.dataset_b = ds_b\n",
    "    \n",
    "    def __len__(self):\n",
    "        return max(len(self.dataset_a), len(self.dataset_b))\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        dp_a = self.dataset_a[i % len(self.dataset_a)]\n",
    "        j = random.randint(0, len(self.dataset_b) - 1)\n",
    "        dp_b = self.dataset_b[j]\n",
    "        return {\n",
    "            'A': dp_a,\n",
    "            'B': dp_b\n",
    "        }\n",
    "\n",
    "\n",
    "class TransformedDataset(Dataset):\n",
    "        \n",
    "    def __init__(self, ds, transform):\n",
    "        self.dataset = ds\n",
    "        self.transform = transform\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.dataset)\n",
    "    \n",
    "    def __getitem__(self, i):\n",
    "        return {k: self.transform(v) for k, v in self.dataset[i].items()}"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "RF6OQwxoFRax"
   },
   "source": [
    "train_ab_ds = Image2ImageDataset(train_A, train_B)\n",
    "test_ab_ds = Image2ImageDataset(test_A, test_B)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "eIcsATwrFRa1"
   },
   "source": [
    "dp = train_ab_ds[20]\n",
    "\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.subplot(121)\n",
    "plt.title(\"Train dataset 'Horses'\")\n",
    "plt.imshow(dp['A'])\n",
    "plt.subplot(122)\n",
    "plt.title(\"Train dataset 'Zebras'\")\n",
    "plt.imshow(dp['B'])"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "X3zfGZoOFRa6"
   },
   "source": [
    "dp = test_ab_ds[20]\n",
    "\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.subplot(121)\n",
    "plt.title(\"Test dataset 'Horses'\")\n",
    "plt.imshow(dp['A'])\n",
    "plt.subplot(122)\n",
    "plt.title(\"Test dataset 'Zebras'\")\n",
    "plt.imshow(dp['B'])"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "q2Hda4tjFRa-"
   },
   "source": [
    "from torchvision.transforms import Compose, ColorJitter, RandomHorizontalFlip, ToTensor, Normalize, RandomCrop\n",
    "\n",
    "# To accelerate the training we reduce the image size to 200x200 pix instead of 256x256\n",
    "train_transform = Compose([\n",
    "    RandomCrop(200),\n",
    "    RandomHorizontalFlip(),\n",
    "    ColorJitter(),\n",
    "    ToTensor(),\n",
    "    Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n",
    "])\n",
    "transformed_train_ab_ds = TransformedDataset(train_ab_ds, transform=train_transform)\n",
    "\n",
    "# Please select appropriate batch_size value according to your infrastructure\n",
    "batch_size = 10\n",
    "train_ab_loader = DataLoader(transformed_train_ab_ds, batch_size=batch_size, shuffle=True, drop_last=True, pin_memory=True, num_workers=4)\n",
    "\n",
    "\n",
    "test_transform = Compose([\n",
    "    ToTensor(),\n",
    "    Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n",
    "])\n",
    "transformed_test_ab_ds = TransformedDataset(test_ab_ds, transform=test_transform)\n",
    "batch_size = 10\n",
    "test_ab_loader = DataLoader(transformed_test_ab_ds, batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True, num_workers=4)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "HtfXU9uaFRbB"
   },
   "source": [
    "import torchvision.utils as vutils\n",
    "\n",
    "# Plot some training images\n",
    "real_batch = next(iter(train_ab_loader))\n",
    "\n",
    "plt.figure(figsize=(16, 8))\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Training Images from A\")\n",
    "plt.imshow( \n",
    "    vutils.make_grid(real_batch['A'][:64], padding=2, normalize=True).cpu().numpy().transpose((1, 2, 0))\n",
    ")\n",
    "\n",
    "plt.figure(figsize=(16, 8))\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Training Images from B\")\n",
    "plt.imshow(\n",
    "    vutils.make_grid(real_batch['B'][:64], padding=2, normalize=True).cpu().numpy().transpose((1, 2, 0))\n",
    ")\n",
    "real_batch = None\n",
    "torch.cuda.empty_cache()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RlZ39Id-FRbF"
   },
   "source": [
    "## Generator and Discriminator networks\n",
    "\n",
    "- Generator network architecture contains 9 residual blocks\n",
    "Using paper's notations:\n",
    "```\n",
    "c7s1-64,d128,d256,R256,R256,R256,R256,R256,R256,R256,R256,R256,u128,u64,c7s1-3,tanh\n",
    "```\n",
    "where `c7s1-k` denotes a 7x7 Convolution-InstanceNorm-ReLU layer with `k` filters and stride 1. `dk` denotes a 3x3 Convolution-InstanceNorm-ReLU layer  with `k` filters and stride 2. Reflection padding was used to reduce artifacts. `Rk` denotes a residual block that contains two 3x3 convolutional layers with the same number of filters on both layer. `uk` denotes a 3x3 fractional-strided-Convolution-InstanceNorm-ReLU layer with `k` filters and stride 1/2."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Ri6_NvWfFRbG"
   },
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "def get_conv_inorm_relu(in_planes, out_planes, kernel_size, stride, reflection_pad=True, with_relu=True):\n",
    "    layers = []\n",
    "    padding = (kernel_size - 1) // 2\n",
    "    if reflection_pad:\n",
    "        layers.append(nn.ReflectionPad2d(padding=padding))\n",
    "        padding = 0\n",
    "    layers += [\n",
    "        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding),\n",
    "        nn.InstanceNorm2d(out_planes, affine=False, track_running_stats=False),\n",
    "    ]\n",
    "    if with_relu:\n",
    "        layers.append(nn.ReLU(inplace=True))\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "\n",
    "def get_conv_transposed_inorm_relu(in_planes, out_planes, kernel_size, stride):\n",
    "    return nn.Sequential(\n",
    "        nn.ConvTranspose2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=1, output_padding=1),\n",
    "        nn.InstanceNorm2d(out_planes, affine=False, track_running_stats=False),\n",
    "        nn.ReLU(inplace=True)\n",
    "    )\n",
    "\n",
    "\n",
    "class ResidualBlock(nn.Module):\n",
    "    \n",
    "    def __init__(self, in_planes):\n",
    "        super(ResidualBlock, self).__init__()\n",
    "        self.conv1 = get_conv_inorm_relu(in_planes, in_planes, kernel_size=3, stride=1)\n",
    "        self.conv2 = get_conv_inorm_relu(in_planes, in_planes, kernel_size=3, stride=1, with_relu=False)        \n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = x\n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)        \n",
    "        return x + residual\n",
    "\n",
    "\n",
    "class Generator(nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(Generator, self).__init__()\n",
    "        \n",
    "        self.c7s1_64 = get_conv_inorm_relu(3, 64, kernel_size=7, stride=1)\n",
    "        self.d128 = get_conv_inorm_relu(64, 128, kernel_size=3, stride=2, reflection_pad=False)\n",
    "        self.d256 = get_conv_inorm_relu(128, 256, kernel_size=3, stride=2, reflection_pad=False)\n",
    "\n",
    "        self.resnet9 = nn.Sequential(*[ResidualBlock(256) for i in range(9)])\n",
    "\n",
    "        self.u128 = get_conv_transposed_inorm_relu(256, 128, kernel_size=3, stride=2)\n",
    "        self.u64 = get_conv_transposed_inorm_relu(128, 64, kernel_size=3, stride=2)\n",
    "        self.c7s1_3 = get_conv_inorm_relu(64, 3, kernel_size=7, stride=1, with_relu=False)\n",
    "        # Replace instance norm by tanh activation\n",
    "        self.c7s1_3[-1] = nn.Tanh()\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Encoding\n",
    "        x = self.c7s1_64(x)\n",
    "        x = self.d128(x)\n",
    "        x = self.d256(x)\n",
    "        \n",
    "        # 9 residual blocks\n",
    "        x = self.resnet9(x)\n",
    "\n",
    "        # Decoding\n",
    "        x = self.u128(x)\n",
    "        x = self.u64(x)\n",
    "        y = self.c7s1_3(x)\n",
    "        return y\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "snM1h6BnFRbK"
   },
   "source": [
    "Let's check the network:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "rfOTRPt4FRbL"
   },
   "source": [
    "x = torch.rand(4, 3, 256, 256)\n",
    "g = Generator()\n",
    "y = g(x)\n",
    "y.shape"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Mb_6xzMvFRbP"
   },
   "source": [
    "- Discriminator network is a PatchGAN (with receptive field 70x70):\n",
    "```\n",
    "C64-C128-C256-C512\n",
    "```\n",
    "where `Ck` denote a 4x4 Convolution-InstanceNorm-LeakyReLU layer with `k` filters and stride 2. After the last layer, \n",
    "a convolution to produce a 1-dimensional output is applied. No `InstanceNorm` for the first `C64` layer. Leaky ReLUs are with a slope of `0.2`. \n",
    "\n",
    "Good explanation is given in [this comment](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/issues/39#issuecomment-305575964) on what is a PatchGAN:\n",
    "```\n",
    "In fact, a \"PatchGAN\" is just a convnet! Or you could say all convnets are patchnets: the power of convnets is that they process each image patch identically and independently, which makes things very cheap (# params, time, memory), and, amazingly, turns out to work.\n",
    "\n",
    "The difference between a PatchGAN and regular GAN discriminator is that rather the regular GAN maps from a 256x256 image to a single scalar output, which signifies \"real\" or \"fake\", whereas the PatchGAN maps from 256x256 to an NxN array of outputs X, where each X_ij signifies whether the patch ij in the image is real or fake. Which is patch ij in the input? Well, output X_ij is just a neuron in a convnet, and we can trace back its receptive field to see which input pixels it is sensitive to. In the CycleGAN architecture, the receptive fields of the discriminator turn out to be 70x70 patches in the input image!\n",
    "\n",
    "This is all mathematically equivalent to if we had manually chopped up the image into 70x70 overlapping patches, run a regular discriminator over each patch, and averaged the results.\n",
    "\n",
    "Maybe it would have been better if we called it a \"Fully Convolutional GAN\" like in FCNs... it's the same idea :)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "rdb5-RAYFRbR"
   },
   "source": [
    "def get_conv_inorm_lrelu(in_planes, out_planes, stride=2, negative_slope=0.2):\n",
    "    return nn.Sequential(\n",
    "        nn.Conv2d(in_planes, out_planes, kernel_size=4, stride=stride, padding=1),\n",
    "        nn.InstanceNorm2d(out_planes, affine=False, track_running_stats=False),\n",
    "        nn.LeakyReLU(negative_slope=negative_slope, inplace=True)\n",
    "    )\n",
    "\n",
    "\n",
    "class Discriminator(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.c64 = nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1)\n",
    "        self.relu = nn.LeakyReLU(0.2, inplace=True)\n",
    "        self.c128 = get_conv_inorm_lrelu(64, 128)\n",
    "        self.c256 = get_conv_inorm_lrelu(128, 256)\n",
    "        self.c512 = get_conv_inorm_lrelu(256, 512, stride=1)\n",
    "        self.last_conv = nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.c64(x)\n",
    "        x = self.relu(x)\n",
    "\n",
    "        x = self.c128(x)\n",
    "        x = self.c256(x)\n",
    "        x = self.c512(x)\n",
    "        y = self.last_conv(x)\n",
    "        return y\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Xu2K1U7JFRbU"
   },
   "source": [
    "Let's check the network:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "On4vpLc_FRbV"
   },
   "source": [
    "x = torch.rand(4, 3, 256, 256)\n",
    "d = Discriminator()\n",
    "y = d(x)\n",
    "y.shape"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "orcYUiDfFRbZ"
   },
   "source": [
    "According to the paper, the weights are initialized from a Gaussian distribution $\\mathcal{N}(0,0.02)$"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "2vId72FGFRba"
   },
   "source": [
    "def init_weights(module):\n",
    "    assert isinstance(module, nn.Module)\n",
    "    if hasattr(module, \"weight\") and module.weight is not None:\n",
    "        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "    if hasattr(module, \"bias\") and module.bias is not None:\n",
    "        torch.nn.init.constant_(module.bias, 0.0)\n",
    "    for c in module.children():\n",
    "        init_weights(c)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "IMoUXZ5yFRbd"
   },
   "source": [
    "g = None; d = None"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lUG71td9FRbg"
   },
   "source": [
    "## Networks training"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ozU5CQ9JFRbh"
   },
   "source": [
    "assert torch.backends.cudnn.enabled\n",
    "torch.backends.cudnn.benchmark = True"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "-zZ5ZuXGFRbl"
   },
   "source": [
    "device = \"cuda\"\n",
    "\n",
    "generator_A2B = Generator().to(device)\n",
    "init_weights(generator_A2B)\n",
    "\n",
    "discriminator_B = Discriminator().to(device)\n",
    "init_weights(discriminator_B)\n",
    "\n",
    "generator_B2A = Generator().to(device)\n",
    "init_weights(generator_B2A)\n",
    "discriminator_A = Discriminator().to(device)\n",
    "init_weights(discriminator_A)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f90Ew1kaFRbo"
   },
   "source": [
    "Similarly, we train the networks from scratch, with a learning rate of `0.0002`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "YbgTEAMCFRbo"
   },
   "source": [
    "from itertools import chain\n",
    "import torch.optim as optim\n",
    "\n",
    "lr = 0.0002\n",
    "beta1 = 0.5\n",
    "\n",
    "optimizer_G = optim.Adam(chain(generator_A2B.parameters(), generator_B2A.parameters()), lr=lr, betas=(beta1, 0.999))\n",
    "optimizer_D = optim.Adam(chain(discriminator_A.parameters(), discriminator_B.parameters()), lr=lr, betas=(beta1, 0.999))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Gs9la9yqFRbr"
   },
   "source": [
    "Let's define some helper functions:\n",
    "- to turn on/off gradients"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "p82seQ9JFRbs"
   },
   "source": [
    "def toggle_grad(model, on_or_off):\n",
    "    # https://github.com/ajbrock/BigGAN-PyTorch/blob/master/utils.py#L674\n",
    "    for param in model.parameters():\n",
    "        param.requires_grad = on_or_off"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wTF4p83kFRbv"
   },
   "source": [
    "### Fake images buffer trick\n",
    "\n",
    "According to the paper, to reduce model oscillation and the discriminators are updated using a history of generated images rather than the ones produced by the latest generators. There is an image buffer that stores the X previously created images."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "jyghMWUPFRbw"
   },
   "source": [
    "buffer_size = 50\n",
    "fake_a_buffer = []\n",
    "fake_b_buffer = []\n",
    "\n",
    "\n",
    "def buffer_insert_and_get(buffer, batch):\n",
    "    output_batch = []\n",
    "    for b in batch:\n",
    "        b = b.unsqueeze(0)\n",
    "        # if buffer is not fully filled:\n",
    "        if len(buffer) < buffer_size:\n",
    "            output_batch.append(b)\n",
    "            buffer.append(b.cpu())\n",
    "        elif random.uniform(0, 1) > 0.5:\n",
    "            # Add newly created image into the buffer and put ont from the buffer into the output\n",
    "            random_index = random.randint(0, buffer_size - 1)            \n",
    "            output_batch.append(buffer[random_index].clone().to(device))\n",
    "            buffer[random_index] = b.cpu()\n",
    "        else:\n",
    "            output_batch.append(b)\n",
    "    return torch.cat(output_batch, dim=0)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K4k46BvuFRb_"
   },
   "source": [
    "Next, let's define a single iteration via `update_fn` function. This function is then used by `ignite.engine.Engine` to update models while running over the input data.\n",
    "\n",
    "As suggested, we divide the objective by 2 while optimizing D, which slows down the rate at which D learns, relative to the rate of G. \n",
    "\n",
    "According to the paper:\n",
    "- generator A is trained  minimize $\\text{mean}_{x \\in A}[(D_B(G(x)) − 1)^2]$ and cycle loss $\\text{mean}_{x \\in A}\\left[ |F(G(x)) - x|_1 \\right]$\n",
    "- generator B is trained  minimize $\\text{mean}_{y \\in B}[(D_A(F(y)) − 1)^2]$ and cycle loss $\\text{mean}_{y \\in B}\\left[ |G(F(y)) - y|_1 \\right]$\n",
    "- discriminators A is trained to minimize $\\text{mean}_{x \\in A}[(D_A(x) − 1)^2] + \\text{mean}_{y \\in B}[D_A(F(y))^2]$.\n",
    "- discriminator B is trained to minimize $\\text{mean}_{y \\in B}[(D_B(y) − 1)^2] + \\text{mean}_{x \\in A}[D_B(G(x))^2]$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JE8dLeEfIl_Z"
   },
   "source": [
    "We will use [`torch.amp.autocast`](https://pytorch.org/docs/master/amp.html#torch.amp.autocast) and [`torch.cuda.amp.GradScaler`](https://pytorch.org/docs/master/amp.html#torch.cuda.amp.GradScaler) to perform automatic mixed precision training. Our code follows a [typical mixed precision training example](https://pytorch.org/docs/master/notes/amp_examples.html#typical-mixed-precision-training)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "vrJls4p-FRcA"
   },
   "source": [
    "from torch.cuda.amp import GradScaler\n",
    "from torch.amp import autocast\n",
    "\n",
    "from ignite.utils import convert_tensor\n",
    "import torch.nn.functional as F\n",
    "\n",
    "amp_enabled = True\n",
    "lambda_value = 10.0\n",
    "amp_scaler = GradScaler(enabled=amp_enabled)\n",
    "\n",
    "\n",
    "def discriminator_forward_pass(discriminator, batch_real, batch_fake, fake_buffer):\n",
    "    decision_real = discriminator(batch_real)\n",
    "    batch_fake = buffer_insert_and_get(fake_buffer, batch_fake)      \n",
    "    decision_fake = discriminator(batch_fake)\n",
    "    return decision_real, decision_fake\n",
    "\n",
    "\n",
    "def compute_loss_generator(batch_decision, batch_real, batch_rec, lambda_value):\n",
    "    # loss gan\n",
    "    target = torch.ones_like(batch_decision)\n",
    "    loss_gan = F.mse_loss(batch_decision, target)\n",
    "    # loss cycle\n",
    "    loss_cycle = F.l1_loss(batch_rec, batch_real) * lambda_value    \n",
    "    return loss_gan + loss_cycle\n",
    "\n",
    "\n",
    "def compute_loss_discriminator(decision_real, decision_fake):\n",
    "    # loss = mean (D_b(y) − 1)^2 + mean D_b(G(x))^2    \n",
    "    loss = F.mse_loss(decision_fake, torch.zeros_like(decision_fake))\n",
    "    loss += F.mse_loss(decision_real, torch.ones_like(decision_real))\n",
    "    return loss\n",
    "\n",
    "\n",
    "def update_fn(engine, batch):\n",
    "    generator_A2B.train()\n",
    "    generator_B2A.train()\n",
    "    discriminator_A.train()\n",
    "    discriminator_B.train()    \n",
    "\n",
    "    real_a = convert_tensor(batch['A'], device=device, non_blocking=True)\n",
    "    real_b = convert_tensor(batch['B'], device=device, non_blocking=True)\n",
    "\n",
    "    # Update generators\n",
    "\n",
    "    # Disable grads computation for the discriminators:\n",
    "    toggle_grad(discriminator_A, False)\n",
    "    toggle_grad(discriminator_B, False)    \n",
    "\n",
    "    with autocast(enabled=amp_enabled):\n",
    "        fake_b = generator_A2B(real_a)\n",
    "        rec_a = generator_B2A(fake_b)\n",
    "        fake_a = generator_B2A(real_b)\n",
    "        rec_b = generator_A2B(fake_a)\n",
    "        decision_fake_a = discriminator_A(fake_a)\n",
    "        decision_fake_b = discriminator_B(fake_b)\n",
    "\n",
    "        # Compute loss for generators and update generators\n",
    "        # loss_a2b = GAN loss: mean (D_b(G(x)) − 1)^2 + Forward cycle loss: || F(G(x)) - x ||_1    \n",
    "        loss_a2b = compute_loss_generator(decision_fake_b, real_a, rec_a, lambda_value)    \n",
    "\n",
    "        # loss_b2a = GAN loss: mean (D_a(F(x)) − 1)^2 + Backward cycle loss: || G(F(y)) - y ||_1\n",
    "        loss_b2a = compute_loss_generator(decision_fake_a, real_b, rec_b, lambda_value)\n",
    "\n",
    "        # total generators loss:\n",
    "        loss_generators = loss_a2b + loss_b2a\n",
    "\n",
    "    optimizer_G.zero_grad()\n",
    "    amp_scaler.scale(loss_generators).backward()\n",
    "    amp_scaler.step(optimizer_G)\n",
    "\n",
    "    decision_fake_a = rec_a = decision_fake_b = rec_b = None\n",
    "    \n",
    "    # Enable grads computation for the discriminators:\n",
    "    toggle_grad(discriminator_A, True)\n",
    "    toggle_grad(discriminator_B, True)    \n",
    "\n",
    "    with autocast(enabled=amp_enabled):\n",
    "        decision_real_a, decision_fake_a = discriminator_forward_pass(discriminator_A, real_a, fake_a.detach(), fake_a_buffer)    \n",
    "        decision_real_b, decision_fake_b = discriminator_forward_pass(discriminator_B, real_b, fake_b.detach(), fake_b_buffer)    \n",
    "        # Compute loss for discriminators and update discriminators\n",
    "        # loss_a = mean (D_a(y) − 1)^2 + mean D_a(F(x))^2\n",
    "        loss_a = compute_loss_discriminator(decision_real_a, decision_fake_a)\n",
    "\n",
    "        # loss_b = mean (D_b(y) − 1)^2 + mean D_b(G(x))^2\n",
    "        loss_b = compute_loss_discriminator(decision_real_b, decision_fake_b)\n",
    "      \n",
    "        # total discriminators loss:\n",
    "        loss_discriminators = 0.5 * (loss_a + loss_b)\n",
    "    \n",
    "    optimizer_D.zero_grad()\n",
    "    amp_scaler.scale(loss_discriminators).backward()\n",
    "    amp_scaler.step(optimizer_D)\n",
    "    amp_scaler.update()\n",
    "    \n",
    "    return {\n",
    "        \"loss_generators\": loss_generators.item(),\n",
    "        \"loss_generator_a2b\": loss_a2b.item(),\n",
    "        \"loss_generator_b2a\": loss_b2a.item(),\n",
    "        \"loss_discriminators\": loss_discriminators.item(),\n",
    "        \"loss_discriminator_a\": loss_a.item(),\n",
    "        \"loss_discriminator_b\": loss_b.item(),\n",
    "    }\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3DF1hNTrFRcD"
   },
   "source": [
    "Let's check `update_fn`"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "rLZextmDzzw_"
   },
   "source": [
    "real_batch = next(iter(train_ab_loader))\n",
    "\n",
    "res = update_fn(engine=None, batch=real_batch)\n",
    "\n",
    "real_batch = None\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "res"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SxVF6kGYFRcK"
   },
   "source": [
    "Now let's define a trainer and add some practical handlers:\n",
    "- log to tensorboard: losses, lr, generated images\n",
    "- progress bar\n",
    "- models/optimizers checkpointing\n",
    "\n",
    "Optionally, we also log to online platform [W&B](https://app.wandb.ai). Current logs can be seen at https://app.wandb.ai/vfdev-5/ignite-cyclegan-torch-amp\n",
    "```\n",
    "pip install --upgrade wandb\n",
    "wandb login your-token\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "camPNT4TcCFu"
   },
   "source": [
    "!pip install --upgrade wandb\n",
    "# !wandb login your-token"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "TO58FENsFRcM"
   },
   "source": [
    "from ignite.engine import Engine, Events\n",
    "from ignite.metrics import RunningAverage\n",
    "\n",
    "from ignite.handlers import TensorboardLogger, WandBLogger\n",
    "from ignite.handlers.tensorboard_logger import OutputHandler, OptimizerParamsHandler"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "jbeZjaBtFRcO"
   },
   "source": [
    "from functools import partial\n",
    "\n",
    "\n",
    "trainer = Engine(update_fn)\n",
    "\n",
    "metric_names = [\n",
    "    'loss_discriminators', \n",
    "    'loss_generators', \n",
    "    'loss_discriminator_a',\n",
    "    'loss_discriminator_b',\n",
    "    'loss_generator_a2b',\n",
    "    'loss_generator_b2a'    \n",
    "]\n",
    "\n",
    "def output_transform(out, name):\n",
    "    return out[name]\n",
    "\n",
    "for name in metric_names:\n",
    "    # here we cannot use lambdas as they do not store argument `name`\n",
    "    RunningAverage(output_transform=partial(output_transform, name=name)).attach(trainer, name)\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "qODNF0imFRcQ"
   },
   "source": [
    "from datetime import datetime\n",
    "\n",
    "exp_name = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "tb_logger = TensorboardLogger(log_dir=\"/tmp/cycle_gan_horse2zebra_tb_logs/{}\".format(exp_name))\n",
    "\n",
    "tb_logger.attach(trainer, \n",
    "                 log_handler=OutputHandler('training', metric_names), \n",
    "                 event_name=Events.ITERATION_COMPLETED)\n",
    "\n",
    "print(\"Experiment name: \", exp_name)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ljTkKgMYFRcT"
   },
   "source": [
    "from pathlib import Path\n",
    "\n",
    "try:\n",
    "    wb_run_name = \"cycle_gan_horse2zebra\"\n",
    "    wb_dir = Path(\"/tmp/cycle_gan_horse2zebra_wandb\")\n",
    "    if not wb_dir.exists():\n",
    "        wb_dir.mkdir()\n",
    "    wb_logger = WandBLogger(\n",
    "        project=\"ignite-cyclegan-torch-amp\",\n",
    "        name=wb_run_name,\n",
    "        sync_tensorboard=True,\n",
    "        dir=wb_dir.as_posix(),\n",
    "        reinit=True\n",
    "    )\n",
    "except RuntimeError:\n",
    "    wb_logger = None"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DbjedQjLFRcV"
   },
   "source": [
    "Let's create an evaluator to inference on train/test images and log the generated images to Tensorboard:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "rCcdC6q9FRcV"
   },
   "source": [
    "from ignite.engine import Engine\n",
    "\n",
    "\n",
    "def evaluate_fn(engine, batch):\n",
    "    generator_A2B.eval()\n",
    "    generator_B2A.eval()    \n",
    "    with torch.no_grad():\n",
    "        real_a = convert_tensor(batch['A'], device=device, non_blocking=True)\n",
    "        real_b = convert_tensor(batch['B'], device=device, non_blocking=True)\n",
    "        \n",
    "        fake_b = generator_A2B(real_a)\n",
    "        rec_a = generator_B2A(fake_b)\n",
    "\n",
    "        fake_a = generator_B2A(real_b)\n",
    "        rec_b = generator_A2B(fake_a)\n",
    "        \n",
    "    return {\n",
    "        'real_a': real_a,\n",
    "        'real_b': real_b,\n",
    "        'fake_a': fake_a,\n",
    "        'fake_b': fake_b,\n",
    "        'rec_a': rec_a,\n",
    "        'rec_b': rec_b,        \n",
    "    }\n",
    "\n",
    "\n",
    "evaluator = Engine(evaluate_fn)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "UIecBsPKFRcX"
   },
   "source": [
    "from torch.utils.data import Subset\n",
    "\n",
    "eval_batch_size = 5\n",
    "\n",
    "train_random_indices = [random.randint(0, len(train_ab_ds) - 1) for _ in range(eval_batch_size)]\n",
    "small_train_ds = Subset(train_ab_ds, train_random_indices)\n",
    "small_train_ds = TransformedDataset(small_train_ds, transform=test_transform)\n",
    "\n",
    "test_random_indices = [random.randint(0, len(test_ab_ds) - 1) for _ in range(eval_batch_size)]\n",
    "small_test_ds = Subset(test_ab_ds, test_random_indices)\n",
    "small_test_ds = TransformedDataset(small_test_ds, transform=test_transform)\n",
    "\n",
    "eval_train_loader = DataLoader(small_train_ds, batch_size=eval_batch_size, shuffle=False, drop_last=False, pin_memory=True, num_workers=4)\n",
    "eval_test_loader = DataLoader(small_test_ds, batch_size=eval_batch_size, shuffle=False, drop_last=False, pin_memory=True, num_workers=4)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "H6XkEchHFRca"
   },
   "source": [
    "@trainer.on(Events.EPOCH_STARTED)\n",
    "def run_evaluation(engine):\n",
    "    evaluator.run(eval_train_loader)\n",
    "    evaluator.run(eval_test_loader)\n",
    "\n",
    "\n",
    "def log_generated_images(engine, logger, event_name):\n",
    "\n",
    "    tag = \"Train\" if engine.state.dataloader == eval_train_loader else \"Test\"\n",
    "    output = engine.state.output\n",
    "    state = trainer.state\n",
    "    global_step = state.get_event_attrib_value(event_name)\n",
    "\n",
    "    # create a grid:\n",
    "    # [real a1, real a2, ...]\n",
    "    # [fake a1, fake a2, ...]\n",
    "    # [rec a1, rec a2, ...]\n",
    "    \n",
    "    s = output['real_a'].shape[0]\n",
    "    res_a = vutils.make_grid(torch.cat([\n",
    "        output['real_a'],\n",
    "        output['fake_b'],\n",
    "        output['rec_a'],\n",
    "    ]), padding=2, normalize=True, nrow=s).cpu()\n",
    "\n",
    "    logger.writer.add_image(tag=\"{} Horses2Zebras (real, fake, rec)\".format(tag), \n",
    "                            img_tensor=res_a, global_step=global_step, dataformats='CHW')\n",
    "\n",
    "    s = output['real_b'].shape[0]\n",
    "    res_b = vutils.make_grid(torch.cat([\n",
    "        output['real_b'],\n",
    "        output['fake_a'],\n",
    "        output['rec_b'],\n",
    "    ]), padding=2, normalize=True, nrow=s).cpu()\n",
    "    logger.writer.add_image(tag=\"{} Zebras2Horses (real, fake, rec)\".format(tag), \n",
    "                            img_tensor=res_b, global_step=global_step, dataformats='CHW')\n",
    "\n",
    "    \n",
    "tb_logger.attach(evaluator,\n",
    "                 log_handler=log_generated_images, \n",
    "                 event_name=Events.COMPLETED)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zrGYmv7RFRcc"
   },
   "source": [
    "We also follow suggested lr scheduling: the same learning rate for the first 100 epochs and linearly decay the rate to zero over the next 100 epochs."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "SSIcXzReFRcc"
   },
   "source": [
    "from ignite.handlers import PiecewiseLinear, ParamGroupScheduler\n",
    "\n",
    "lr = 0.0002\n",
    "\n",
    "milestones_values = [\n",
    "    (0, lr),\n",
    "    (100, lr),\n",
    "    (200, 0.0)\n",
    "]\n",
    "gen_lr_scheduler = PiecewiseLinear(optimizer_D, param_name='lr', milestones_values=milestones_values)\n",
    "desc_lr_scheduler = PiecewiseLinear(optimizer_G, param_name='lr', milestones_values=milestones_values)\n",
    "\n",
    "lr_scheduler = ParamGroupScheduler([gen_lr_scheduler, desc_lr_scheduler], \n",
    "                                   names=['gen_lr_scheduler', 'desc_lr_scheduler'])\n",
    "\n",
    "trainer.add_event_handler(Events.EPOCH_STARTED, lr_scheduler)\n",
    "\n",
    "\n",
    "tb_logger.attach(trainer,\n",
    "                 log_handler=OptimizerParamsHandler(optimizer_G, \"lr\"), \n",
    "                 event_name=Events.EPOCH_STARTED)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xWJcV2veFRcf"
   },
   "source": [
    "Let's setup models/optimizers checkpointing:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "7ZONS845FRcg"
   },
   "source": [
    "from ignite.handlers import ModelCheckpoint, TerminateOnNan"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "F-emWk-YFRci"
   },
   "source": [
    "!rm -rf \"/tmp/cycle_gan_checkpoints\" \n",
    "!mkdir \"/tmp/cycle_gan_checkpoints\"\n",
    "\n",
    "from ignite.handlers import ModelCheckpoint, TerminateOnNan\n",
    "\n",
    "\n",
    "checkpoint_handler = ModelCheckpoint(dirname=\"/tmp/cycle_gan_checkpoints\", filename_prefix=\"\")\n",
    "\n",
    "to_save = {\n",
    "    \"generator_A2B\": generator_A2B,\n",
    "    \"discriminator_B\": discriminator_B,\n",
    "    \"generator_B2A\": generator_B2A,\n",
    "    \"discriminator_A\": discriminator_A,\n",
    "    \n",
    "    \"optimizer_G\": optimizer_G,\n",
    "    \"optimizer_D\": optimizer_D,\n",
    "}\n",
    "\n",
    "trainer.add_event_handler(Events.ITERATION_COMPLETED(every=500), checkpoint_handler, to_save)\n",
    "trainer.add_event_handler(Events.ITERATION_COMPLETED, TerminateOnNan())"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "RtQKr6yxFRck"
   },
   "source": [
    "from ignite.handlers import ProgressBar\n",
    "\n",
    "# Iteration-wise progress bar\n",
    "ProgressBar(bar_format=\"\").attach(trainer)\n",
    "# Epoch-wise progress bar with display of training losses\n",
    "ProgressBar(persist=True, bar_format=\"\").attach(trainer, metric_names=['loss_discriminators', 'loss_generators'], \n",
    "                                                event_name=Events.EPOCH_STARTED, closing_event_name=Events.COMPLETED)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bM-6vr8pcmOW"
   },
   "source": [
    "# Display in Firefox may not work properly. Use Chrome.\n",
    "%load_ext tensorboard\n",
    "\n",
    "%tensorboard --logdir=/tmp/cycle_gan_horse2zebra_tb_logs"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "CcxhA9rHFRcn"
   },
   "source": [
    "trainer.run(train_ab_loader, max_epochs=200)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "XtXSfbHqFRct"
   },
   "source": [
    "tb_logger.close()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7CbI9ngDFRcx"
   },
   "source": [
    "### Inference with trained generator\n",
    "\n",
    "Let's display saved checkpoint, load weights for \"Horses to Zebras\" generator and run an inference on a test image"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "2xNAdA-WFRcx"
   },
   "source": [
    "!ls /tmp/cycle_gan_checkpoints/"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "TOhSsWRzFRcz"
   },
   "source": [
    "checkpoint_path = \"/tmp/cycle_gan_checkpoints/checkpoint_26500.pt\"\n",
    "\n",
    "# let's save this checkpoint to W&B\n",
    "if wb_logger is not None:\n",
    "    wb_logger.save(checkpoint_path)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "j5BFfZceFRc2"
   },
   "source": [
    "checkpoint_state_dict = torch.load(checkpoint_path)\n",
    "generator_A2B.load_state_dict(checkpoint_state_dict[\"generator_A2B\"])"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "VGLb14xyFRc4"
   },
   "source": [
    "def normalize(x):\n",
    "    vmin = x.min()\n",
    "    vmax = x.max()\n",
    "    x.clamp_(min=vmin, max=vmax)\n",
    "    x.add_(-vmin).div_(vmax - vmin + 1e-5)\n",
    "    return x"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "BO_DBOhSFRc7"
   },
   "source": [
    "i = random.randint(0, len(test_ab_ds) - 1)\n",
    "img = test_ab_ds[i]['A']\n",
    "x = test_transform(img)\n",
    "x = x.unsqueeze(0).to(device)\n",
    "\n",
    "\n",
    "with torch.no_grad():\n",
    "    y_pred = generator_A2B(x)\n",
    "    \n",
    "\n",
    "img_pred = (255 * normalize(y_pred[0, ...])).cpu().numpy().transpose((1, 2, 0)).astype('uint8')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "vseNrx2YFRc-"
   },
   "source": [
    "plt.figure(figsize=(10, 5))\n",
    "plt.subplot(121)\n",
    "plt.title(\"Horse\")\n",
    "plt.imshow(img)\n",
    "plt.subplot(122)\n",
    "plt.title(\"Generated zebra\")\n",
    "plt.imshow(img_pred)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VqFJ0FrPFRdA"
   },
   "source": [
    "Let's apply for fun our zebra-filter on an image with deep learning gurus:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "6BlVkvybFRdA"
   },
   "source": [
    "!wget https://www.kdnuggets.com/wp-content/uploads/photo.jpg -O/tmp/dl_durus.jpg"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "eB7aAeO1FRdC"
   },
   "source": [
    "from PIL import Image\n",
    "\n",
    "img = Image.open(\"/tmp/dl_durus.jpg\")\n",
    "x = test_transform(img)\n",
    "x = x.unsqueeze(0).to(device)\n",
    "\n",
    "\n",
    "with torch.no_grad():\n",
    "    y_pred = generator_A2B(x)\n",
    "\n",
    "\n",
    "img_pred = (255 * normalize(y_pred[0, ...])).cpu().numpy().transpose((1, 2, 0)).astype('uint8')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "8Mcypu0lFRdD"
   },
   "source": [
    "plt.figure(figsize=(15, 8))\n",
    "plt.subplot(121)\n",
    "plt.title(\"4 Deep-learning gurus\")\n",
    "plt.imshow(img)\n",
    "plt.subplot(122)\n",
    "plt.title(\"Zebras\")\n",
    "plt.imshow(img_pred)"
   ],
   "execution_count": null,
   "outputs": []
  }
 ]
}