examples/notebooks/customizing_what_happens_in_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "poK9I4m9wSTW"
   },
   "source": [
    "# Catalyst - customizing what happens in `train()`\n",
    "based on `Keras customizing what happens in fit`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NJb_eBKEwXS6"
   },
   "source": [
    "## Introduction\n",
    "\n",
    "When you're doing supervised learning, you can use `train()` and everything works smoothly.\n",
    "\n",
    "A core principle of Catalyst is **progressive disclosure of complexity**. You should always be able to get into lower-level workflows in a gradual way. You shouldn't fall off a cliff if the high-level functionality doesn't exactly match your use case. You should be able to gain more control over the small details while retaing a commensurate amount of high-level convenience. \n",
    "\n",
    "When you need to customize what `train()` does, you should **override the `_handle_batch` function of the `Runner` class**. This is the function that is called by `train()` for every batch of data. You will then be able to call `train()` as usual -- and it will be running your own learning algorithm.\n",
    "\n",
    "Note that this pattern does not prevent you from building models with the Functional API. You can do this with **any** PyTorch model.\n",
    "\n",
    "Let's see how that works."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "E6R34jh5xKkW"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 9418,
     "status": "ok",
     "timestamp": 1587014070467,
     "user": {
      "displayName": "Sergey Kolesnikov",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjYXcxiGiIDQYhW2wTkdrNLwx68llP5BzH91oGlAQ=s64",
      "userId": "07081474162282073276"
     },
     "user_tz": -180
    },
    "id": "S1rkIIaKaG2O",
    "outputId": "5c4d2dc1-74a0-4b04-f3f5-9a08e8efc28c"
   },
   "outputs": [],
   "source": [
    "!pip install catalyst==20.10.1\n",
    "# don't forget to restart runtime for correct `PIL` work with Colab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 255
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 1572,
     "status": "ok",
     "timestamp": 1587014153359,
     "user": {
      "displayName": "Sergey Kolesnikov",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjYXcxiGiIDQYhW2wTkdrNLwx68llP5BzH91oGlAQ=s64",
      "userId": "07081474162282073276"
     },
     "user_tz": -180
    },
    "id": "3eLP4fR6wCYc",
    "outputId": "1e36766e-6d62-46da-894f-8d2e4967544c"
   },
   "outputs": [],
   "source": [
    "import catalyst\n",
    "from catalyst import dl, utils\n",
    "catalyst.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5F8q4oByxt2T"
   },
   "source": [
    "## A first simple example\n",
    "\n",
    "Let's start from a simple example:\n",
    "\n",
    "- We create a new runner that subclasses `dl.Runner`.\n",
    "- We just override the method `_handle_batch(self, batch)`.\n",
    "- We do our train step with any possible custom logic.\n",
    "\n",
    "The input argument `batch` is what gets passed to fit as training data. If you pass a `torch.utils.data.DataLoader`, by calling `train(loaders={\"train\": loader, \"valid\": loader}, ...)`, then `batch` will be what gets yielded by `loader` at each batch.\n",
    "\n",
    "In the body of the `_handle_batch` method, we implement a regular training update, similar to what you are already familiar with. Importantly, **we log metrics via `self.batch_metrics`**, which passes them to the loggers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MbTkRLQUxQmC"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.nn import functional as F\n",
    "\n",
    "class CustomRunner(dl.Runner):\n",
    "\n",
    "  def _handle_batch(self, batch):\n",
    "    # Unpack the data. Its structure depends on your model and\n",
    "    # on what you pass to `train()`.\n",
    "    x, y = batch\n",
    "\n",
    "    y_pred = self.model(x) # Forward pass\n",
    "\n",
    "    # Compute the loss value\n",
    "    loss = F.mse_loss(y_pred, y)\n",
    "\n",
    "    # Update metrics (includes the metric that tracks the loss)\n",
    "    self.batch_metrics.update({\"loss\": loss, \"mae\": F.l1_loss(y_pred, y)})\n",
    "\n",
    "    if self.is_train_loader:\n",
    "      # Compute gradients\n",
    "      loss.backward()\n",
    "      # Update weights\n",
    "      # (the optimizer is stored in `self.state`)\n",
    "      self.optimizer.step()\n",
    "      self.optimizer.zero_grad()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "nAEiVP4IzNj-"
   },
   "source": [
    "Let's try this out:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 562
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 17386,
     "status": "error",
     "timestamp": 1587015544733,
     "user": {
      "displayName": "Sergey Kolesnikov",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjYXcxiGiIDQYhW2wTkdrNLwx68llP5BzH91oGlAQ=s64",
      "userId": "07081474162282073276"
     },
     "user_tz": -180
    },
    "id": "AlUHnIG6zPV9",
    "outputId": "bcc53cac-174d-4a1e-c3e8-441c102609cb"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "\n",
    "# Construct custom data\n",
    "num_samples, num_features = int(1e4), int(1e1)\n",
    "X, y = torch.rand(num_samples, num_features), torch.rand(num_samples, 1)\n",
    "dataset = TensorDataset(X, y)\n",
    "loader = DataLoader(dataset, batch_size=32, num_workers=1)\n",
    "loaders = {\"train\": loader, \"valid\": loader}\n",
    "\n",
    "# and model\n",
    "model = torch.nn.Linear(num_features, 1)\n",
    "criterion = torch.nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "\n",
    "# Just use `train` as usual\n",
    "runner = CustomRunner()\n",
    "runner.train(\n",
    "  model=model, \n",
    "  optimizer=optimizer, \n",
    "  loaders=loaders, \n",
    "  num_epochs=3,\n",
    "  verbose=True, # you can pass True for more precise training process logging\n",
    "  timeit=False, # you can pass True to measure execution time of different parts of train process\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NGJgVd9lzkQc"
   },
   "source": [
    "## Going high-level\n",
    "\n",
    "Naturally, you could skip a loss function backward in `_handle_batch()`, and instead do everything with `Callbacks` in `train` params. Likewise for metrics. Here's a high-level example, that only uses `_handle_batch()` for model forward pass and metrics computation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "\n",
    "\n",
    "class CustomRunner(dl.Runner):\n",
    "\n",
    "  def _handle_batch(self, batch):\n",
    "    # Unpack the data. Its structure depends on your model and\n",
    "    # on what you pass to `train()`.\n",
    "    x, y = batch\n",
    "\n",
    "    y_pred = self.model(x) # Forward pass\n",
    "\n",
    "    # Compute the loss value\n",
    "    # (the criterion is stored in `self.state` also)\n",
    "    loss = self.criterion(y_pred, y)\n",
    "\n",
    "    # Update metrics (includes the metric that tracks the loss)\n",
    "    self.batch_metrics.update({\"loss\": loss, \"mae\": F.l1_loss(y_pred, y)})\n",
    "\n",
    "\n",
    "# Construct custom data\n",
    "num_samples, num_features = int(1e4), int(1e1)\n",
    "X, y = torch.rand(num_samples, num_features), torch.rand(num_samples, 1)\n",
    "dataset = TensorDataset(X, y)\n",
    "loader = DataLoader(dataset, batch_size=32, num_workers=1)\n",
    "loaders = {\"train\": loader, \"valid\": loader}\n",
    "\n",
    "# and model\n",
    "model = torch.nn.Linear(num_features, 1)\n",
    "criterion = torch.nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "\n",
    "# Just use `train` as usual\n",
    "runner = CustomRunner()\n",
    "runner.train(\n",
    "  model=model, \n",
    "  optimizer=optimizer,\n",
    "  criterion=criterion,       # you could also pass any PyTorch criterion for loss computation\n",
    "  scheduler=None,            # or scheduler, but let's simplify the train loop for now :)\n",
    "  loaders=loaders, \n",
    "  num_epochs=3,\n",
    "  verbose=True,\n",
    "  timeit=False,\n",
    "  callbacks={\n",
    "    \"optimizer\": dl.OptimizerCallback(\n",
    "      metric_key=\"loss\",     # you can also pass 'mae' to optimize it instead\n",
    "                             # generaly, you can optimize any differentiable metric from `runner.batch_metrics`\n",
    "      accumulation_steps=1,  # also you can pass any number of steps for gradient accumulation\n",
    "      grad_clip_params=None, # or yor use `{\"func\": \"clip_grad_norm_\", max_norm=1, norm_type=2}`\n",
    "                             #         or `{\"func\": \"clip_grad_value_\", clip_value=1}`\n",
    "                             # for gradient clipping during training!\n",
    "                             # for more information about gradient clipping please follow pytorch docs\n",
    "                             # https://pytorch.org/docs/stable/nn.html#clip-grad-norm\n",
    "    )\n",
    "  }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vVZtc6P61icn"
   },
   "source": [
    "## Metrics support through Callbacks\n",
    "\n",
    "Let's go even deeper! Could we transfer different metrics/criterions computation to `Callbacks` too? Of course! If you want to support different losses, you'd simply do the following:\n",
    "\n",
    "- Do your model forward pass as usual.\n",
    "- Save model input to `runner.input` and model output to `runner.output`, so Callbacks can find it.\n",
    "- Add extra callbacks, that will use data from `runner.input` and `runner.output` for computation.\n",
    "\n",
    "That's it. That's the list. Let's see the example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "\n",
    "\n",
    "class CustomRunner(dl.Runner):\n",
    "\n",
    "  def _handle_batch(self, batch):\n",
    "    # Unpack the data. Its structure depends on your model and\n",
    "    # on what you pass to `train()`.\n",
    "    x, y = batch\n",
    "\n",
    "    y_pred = self.model(x) # Forward pass\n",
    "    \n",
    "    # pass network input to state `input`\n",
    "    self.input = {\"features\": x, \"targets\": y}\n",
    "    # and network output to state `output`\n",
    "    # we recommend to use key-value storage to make it Callbacks-friendly\n",
    "    self.output = {\"logits\": y_pred}\n",
    "\n",
    "\n",
    "# Construct custom data\n",
    "num_samples, num_features = int(1e4), int(1e1)\n",
    "X, y = torch.rand(num_samples, num_features), torch.rand(num_samples, 1)\n",
    "dataset = TensorDataset(X, y)\n",
    "loader = DataLoader(dataset, batch_size=32, num_workers=1)\n",
    "loaders = {\"train\": loader, \"valid\": loader}\n",
    "\n",
    "# and model\n",
    "model = torch.nn.Linear(num_features, 1)\n",
    "criterion = torch.nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "\n",
    "# Just use `train` as usual\n",
    "runner = CustomRunner()\n",
    "runner.train(\n",
    "  model=model, \n",
    "  optimizer=optimizer,\n",
    "  criterion=criterion,\n",
    "  scheduler=None,\n",
    "  loaders=loaders, \n",
    "  num_epochs=3,\n",
    "  verbose=True,\n",
    "  timeit=False,\n",
    "  callbacks={\n",
    "    \"criterion\": dl.CriterionCallback(  # special Callback for criterion computation\n",
    "      input_key=\"targets\",              # `input_key` specifies correct labels (or `y_true`) from `runner.input` \n",
    "      output_key=\"logits\",              # `output_key` specifies model predictions (`y_pred`) from `runner.output`\n",
    "      prefix=\"loss\",                    # `prefix` - key to use with `runner.batch_metrics`\n",
    "    ),  # alias for `runner.batch_metrics[prefix] = runner.criterion(runner.output[output_key], runner.input[input_key])`\n",
    "    \"metric\": dl.MetricCallback(        # special Callback for metrics computation\n",
    "      input_key=\"targets\",              # shares logic with `CriterionCallback`\n",
    "      output_key=\"logits\",\n",
    "      prefix=\"loss_mae\",\n",
    "      metric_fn=F.l1_loss,              # metric function to use\n",
    "    ),  # alias for `runner.batch_metrics[prefix] = metric_fn(runner.output[output_key], runner.input[input_key])`\n",
    "    \"optimizer\": dl.OptimizerCallback(\n",
    "      metric_key=\"loss\", \n",
    "      accumulation_steps=1,\n",
    "      grad_clip_params=None,\n",
    "    )\n",
    "  }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simplify it a bit - SupervisedRunner\n",
    "\n",
    "But can we simplify last example a bit? <br/>\n",
    "What if we know, that we are going to train `Supervised` model, that will take some `features` in and output some `logits` back? <br/>\n",
    "Looks like commom case... could we automate it? Let's check it out!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "\n",
    "# Construct custom data\n",
    "num_samples, num_features = int(1e4), int(1e1)\n",
    "X, y = torch.rand(num_samples, num_features), torch.rand(num_samples, 1)\n",
    "dataset = TensorDataset(X, y)\n",
    "loader = DataLoader(dataset, batch_size=32, num_workers=1)\n",
    "loaders = {\"train\": loader, \"valid\": loader}\n",
    "\n",
    "# and model\n",
    "model = torch.nn.Linear(num_features, 1)\n",
    "criterion = torch.nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "\n",
    "# Just use `train` as usual\n",
    "runner = dl.SupervisedRunner(  # `SupervisedRunner` works with any model like `some_output = model(some_input)`\n",
    "  input_key=\"features\",        # if your dataloader yields (x, y) tuple, it will be transformed to \n",
    "  output_key=\"logits\",         # {input_key: x, input_target_key: y} and stored to runner.input\n",
    "  input_target_key=\"targets\",  # then the model will be used like\n",
    ")                              # runner.output = model(runner.input[input_key])\n",
    "                               # loss computation suppose to looks like\n",
    "                               # loss = criterion(runner.output[input_target_key], runner.output[output_key])\n",
    "                               # and stored to `runner.batch_metrics['loss']`\n",
    "\n",
    "runner.train(\n",
    "  model=model, \n",
    "  optimizer=optimizer,\n",
    "  criterion=criterion,\n",
    "  scheduler=None,\n",
    "  loaders=loaders, \n",
    "  num_epochs=3,\n",
    "  verbose=True,\n",
    "  timeit=False,\n",
    "  callbacks={\n",
    "    \"criterion_mse\": dl.CriterionCallback(\n",
    "      input_key=\"targets\",\n",
    "      output_key=\"logits\",\n",
    "      prefix=\"loss\",\n",
    "    ),\n",
    "    \"criterion_mae\": dl.MetricCallback(\n",
    "      input_key=\"targets\",\n",
    "      output_key=\"logits\",\n",
    "      prefix=\"mae\",\n",
    "      metric_fn=F.l1_loss,\n",
    "    ),\n",
    "    \"optimizer\": dl.OptimizerCallback(\n",
    "      metric_key=\"loss\", \n",
    "      accumulation_steps=1,\n",
    "      grad_clip_params=None,\n",
    "    )\n",
    "  }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rn1q6NCP2dtR"
   },
   "source": [
    "## Providing your own inference step\n",
    "\n",
    "But let's return to the basics.\n",
    "\n",
    "What if you want to do the same customization for calls to `runner.predict_*()`? Then you would override `predict_batch` in exactly the same way. Here's what it looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.nn import functional as F\n",
    "\n",
    "class CustomRunner(dl.Runner):\n",
    "    \n",
    "  def predict_batch(self, batch):                 # here is the trick\n",
    "    return self.model(batch[0].to(self.device))   # you can write any prediciton logic here\n",
    "\n",
    "  def _handle_batch(self, batch):                 # our first time example\n",
    "    # Unpack the data. Its structure depends on your model and\n",
    "    # on what you pass to `train()`.\n",
    "    x, y = batch\n",
    "\n",
    "    y_pred = self.model(x) # Forward pass\n",
    "\n",
    "    # Compute the loss value\n",
    "    loss = F.mse_loss(y_pred, y)\n",
    "\n",
    "    # Update metrics (includes the metric that tracks the loss)\n",
    "    self.batch_metrics.update({\"loss\": loss, \"mae\": F.l1_loss(y_pred, y)})\n",
    "\n",
    "    if self.is_train_loader:\n",
    "      # Compute gradients\n",
    "      loss.backward()\n",
    "      # Update weights\n",
    "      # (the optimizer is stored in `self.state`)\n",
    "      self.optimizer.step()\n",
    "      self.optimizer.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "\n",
    "# Construct custom data\n",
    "num_samples, num_features = int(1e4), int(1e1)\n",
    "X, y = torch.rand(num_samples, num_features), torch.rand(num_samples, 1)\n",
    "dataset = TensorDataset(X, y)\n",
    "loader = DataLoader(dataset, batch_size=32, num_workers=1)\n",
    "loaders = {\"train\": loader, \"valid\": loader}\n",
    "\n",
    "# and model\n",
    "model = torch.nn.Linear(num_features, 1)\n",
    "criterion = torch.nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "\n",
    "# Just use `train` as usual\n",
    "runner = CustomRunner()\n",
    "runner.train(\n",
    "  model=model, \n",
    "  optimizer=optimizer, \n",
    "  loaders=loaders, \n",
    "  num_epochs=3,\n",
    "  verbose=True,\n",
    "  timeit=False,\n",
    "  load_best_on_end=True, # flag to load best model at the end of the training process\n",
    "  logdir=\"./logs\",       # logdir to store models checkpoints (required for `load_best_on_end`)\n",
    ")\n",
    "# and use `batch` prediciton\n",
    "prediction = runner.predict_batch(next(iter(loader))) # let's sample first batch from loader\n",
    "# or `loader` prediction\n",
    "for prediction in runner.predict_loader(loader=loader):\n",
    "  assert prediction.detach().cpu().numpy().shape[-1] == 1 # as we have 1-class regression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, after model training and evaluation, it's time to prepare it for deployment. PyTorch upport model tracing for production-friendly Deep Leanring models deployment.\n",
    "\n",
    "Could we make it quick with Catalyst? Sure!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# you can trace your model through batch 'mode'\n",
    "traced_model = runner.trace(batch=next(iter(loader)))\n",
    "# or loader 'mode' - it will take first batch automatically\n",
    "traced_model = runner.trace(loader=loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9-CMpP5a3Wcp"
   },
   "source": [
    "## Wrapping up: an end-to-end GAN example\n",
    "\n",
    "Let's walk through an end-to-end example that leverages everything you just learned.\n",
    "\n",
    "Let's consider:\n",
    "\n",
    "- A generator network meant to generate 28x28x1 images.\n",
    "- A discriminator network meant to classify 28x28x1 images into two classes (\"fake\" - 1 and \"real\" - 0).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from catalyst.contrib.nn import GlobalMaxPool2d, Flatten, Lambda\n",
    "\n",
    "# Create the discriminator\n",
    "discriminator = nn.Sequential(\n",
    "    nn.Conv2d(1, 64, (3, 3), stride=(2, 2), padding=1),\n",
    "    nn.LeakyReLU(0.2, inplace=True),\n",
    "    nn.Conv2d(64, 128, (3, 3), stride=(2, 2), padding=1),\n",
    "    nn.LeakyReLU(0.2, inplace=True),\n",
    "    GlobalMaxPool2d(),\n",
    "    Flatten(),\n",
    "    nn.Linear(128, 1)\n",
    ")\n",
    "\n",
    "# Create the generator\n",
    "latent_dim = 128\n",
    "generator = nn.Sequential(\n",
    "    # We want to generate 128 coefficients to reshape into a 7x7x128 map\n",
    "    nn.Linear(128, 128 * 7 * 7),\n",
    "    nn.LeakyReLU(0.2, inplace=True),\n",
    "    Lambda(lambda x: x.view(x.size(0), 128, 7, 7)),\n",
    "    nn.ConvTranspose2d(128, 128, (4, 4), stride=(2, 2), padding=1),\n",
    "    nn.LeakyReLU(0.2, inplace=True),\n",
    "    nn.ConvTranspose2d(128, 128, (4, 4), stride=(2, 2), padding=1),\n",
    "    nn.LeakyReLU(0.2, inplace=True),\n",
    "    nn.Conv2d(128, 1, (7, 7), padding=3),\n",
    "    nn.Sigmoid(),\n",
    ")\n",
    "\n",
    "# Final model\n",
    "model = {\n",
    "    \"generator\": generator,\n",
    "    \"discriminator\": discriminator,\n",
    "}\n",
    "\n",
    "optimizer = {\n",
    "    \"generator\": torch.optim.Adam(generator.parameters(), lr=0.0003, betas=(0.5, 0.999)),\n",
    "    \"discriminator\": torch.optim.Adam(discriminator.parameters(), lr=0.0003, betas=(0.5, 0.999)),\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "POY42XRf5Jbd"
   },
   "source": [
    "Here's a feature-complete `GANRunner`, overriding `predict_batch()` to use its own signature, and implementing the entire GAN algorithm in 16 lines in `_handle_batch`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iyKOtjfn5RL3"
   },
   "outputs": [],
   "source": [
    "class GANRunner(dl.Runner):\n",
    "  \n",
    "  def _init(self, latent_dim: int):\n",
    "    self.latent_dim = latent_dim\n",
    "    self.experiment = None  # spoiler for next lesson ;)\n",
    "\n",
    "  def predict_batch(self, batch):\n",
    "    random_latent_vectors = torch.randn(1, self.latent_dim).to(self.device)\n",
    "    generated_images = self.model[\"generator\"](random_latent_vectors)\n",
    "    return generated_images\n",
    "\n",
    "  def _handle_batch(self, batch):\n",
    "    real_images, _ = batch\n",
    "    batch_metrics = {}\n",
    "    \n",
    "    # Sample random points in the latent space\n",
    "    batch_size = real_images.shape[0]\n",
    "    random_latent_vectors = torch.randn(batch_size, self.latent_dim).to(self.device)\n",
    "    \n",
    "    # Decode them to fake images\n",
    "    generated_images = self.model[\"generator\"](random_latent_vectors).detach()\n",
    "    # Combine them with real images\n",
    "    combined_images = torch.cat([generated_images, real_images])\n",
    "    \n",
    "    # Assemble labels discriminating real from fake images\n",
    "    labels = torch.cat([\n",
    "        torch.ones((batch_size, 1)), torch.zeros((batch_size, 1))\n",
    "    ]).to(self.device)\n",
    "    # Add random noise to the labels - important trick!\n",
    "    labels += 0.05 * torch.rand(labels.shape).to(self.device)\n",
    "    \n",
    "    # Train the discriminator\n",
    "    predictions = self.model[\"discriminator\"](combined_images)\n",
    "    batch_metrics[\"loss_discriminator\"] = \\\n",
    "      F.binary_cross_entropy_with_logits(predictions, labels)\n",
    "    \n",
    "    # Sample random points in the latent space\n",
    "    random_latent_vectors = torch.randn(batch_size, self.latent_dim).to(self.device)\n",
    "    # Assemble labels that say \"all real images\"\n",
    "    misleading_labels = torch.zeros((batch_size, 1)).to(self.device)\n",
    "    \n",
    "    # Train the generator\n",
    "    generated_images = self.model[\"generator\"](random_latent_vectors)\n",
    "    predictions = self.model[\"discriminator\"](generated_images)\n",
    "    batch_metrics[\"loss_generator\"] = \\\n",
    "      F.binary_cross_entropy_with_logits(predictions, misleading_labels)\n",
    "    \n",
    "    self.batch_metrics.update(**batch_metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zYGZRIJh6ZYu"
   },
   "source": [
    "Let's test-drive it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from torch.utils.data import DataLoader\n",
    "from catalyst.data.cv import ToTensor\n",
    "from catalyst.contrib.datasets import MNIST\n",
    "\n",
    "loaders = {\n",
    "  \"train\": DataLoader(\n",
    "    MNIST(os.getcwd(), train=True, download=True, transform=ToTensor()), \n",
    "    batch_size=64),\n",
    "}\n",
    "\n",
    "runner = GANRunner(latent_dim=latent_dim)\n",
    "runner.train(\n",
    "    model=model, \n",
    "    optimizer=optimizer,\n",
    "    loaders=loaders,\n",
    "    callbacks=[\n",
    "        dl.OptimizerCallback(\n",
    "            optimizer_key=\"generator\", \n",
    "            metric_key=\"loss_generator\"\n",
    "        ),\n",
    "        dl.OptimizerCallback(\n",
    "            optimizer_key=\"discriminator\", \n",
    "            metric_key=\"loss_discriminator\"\n",
    "        ),\n",
    "    ],\n",
    "    main_metric=\"loss_generator\",\n",
    "    num_epochs=20,\n",
    "    verbose=True,\n",
    "    logdir=\"./logs_gan\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "M9Fz5_u68FqW"
   },
   "source": [
    "The idea behind deep learning are simple, so why should their implementation be painful?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "utils.set_global_seed(42)\n",
    "generated_image = runner.predict_batch(None)\n",
    "plt.imshow(generated_image[0, 0].detach().cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard\n",
    "%tensorboard --logdir ./logs_gan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Сustomizing_what_happens_in_train.ipynb",
   "provenance": []
  },
  "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.7"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}