examples/notebooks/customizing_what_happens_in_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Catalyst - customizing what happens in `train()`\n",
    "based on `Keras customizing what happens in fit`"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "poK9I4m9wSTW"
   }
  },
  {
   "cell_type": "markdown",
   "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."
   ],
   "metadata": {
    "colab_type": "text",
    "id": "NJb_eBKEwXS6"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Setup"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "E6R34jh5xKkW"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!pip install catalyst[ml]==21.8\n",
    "# don't forget to restart runtime for correct `PIL` work with Colab"
   ],
   "outputs": [],
   "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"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import catalyst\n",
    "from catalyst import dl, metrics, utils\n",
    "catalyst.__version__"
   ],
   "outputs": [],
   "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"
   }
  },
  {
   "cell_type": "markdown",
   "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 `handle_batch(self, batch)` method for custom train step logic \n",
    "- And update `on_loader_start`/`on_loader_start` handlers for correct custom metrics aggregation.\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 batch-based metrics via `self.batch_metrics`**, which passes them to the loggers.\n",
    "\n",
    "Addiionally, we have to use [`AdditiveMetric`](https://catalyst-team.github.io/catalyst/api/metrics.html#AdditiveMetric) during `on_loader_start` and `on_loader_start` for correct metrics aggregation for the whole loader. Importantly, **we log loader-based metrics via `self.loader_metrics`**, which passes them to the loggers."
   ],
   "metadata": {
    "colab_type": "text",
    "id": "5F8q4oByxt2T"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import torch\n",
    "from torch.nn import functional as F\n",
    "\n",
    "class CustomRunner(dl.Runner):\n",
    "    \n",
    "    def on_loader_start(self, runner):\n",
    "        super().on_loader_start(runner)\n",
    "        self.meters = {\n",
    "            key: metrics.AdditiveMetric(compute_on_call=False)\n",
    "            for key in [\"loss\", \"mae\"]\n",
    "        }\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",
    "        for key in [\"loss\", \"mae\"]:\n",
    "            self.meters[key].update(self.batch_metrics[key].item(), self.batch_size)\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()\n",
    "    \n",
    "    def on_loader_end(self, runner):\n",
    "        for key in [\"loss\", \"mae\"]:\n",
    "            self.loader_metrics[key] = self.meters[key].compute()[0]\n",
    "        super().on_loader_end(runner)"
   ],
   "outputs": [],
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MbTkRLQUxQmC"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Let's try this out:"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "nAEiVP4IzNj-"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "# and use `train`\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",
    ")"
   ],
   "outputs": [],
   "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"
   }
  },
  {
   "cell_type": "markdown",
   "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:"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "NGJgVd9lzkQc"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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 on_loader_start(self, runner):\n",
    "        super().on_loader_start(runner)\n",
    "        self.meters = {\n",
    "            key: metrics.AdditiveMetric(compute_on_call=False)\n",
    "            for key in [\"loss\", \"mae\"]\n",
    "        }\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",
    "        for key in [\"loss\", \"mae\"]:\n",
    "            self.meters[key].update(self.batch_metrics[key].item(), self.batch_size)\n",
    "\n",
    "    def on_loader_end(self, runner):\n",
    "        for key in [\"loss\", \"mae\"]:\n",
    "            self.loader_metrics[key] = self.meters[key].compute()[0]\n",
    "        super().on_loader_end(runner)\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_fn=None,     # or you can use `grad_clip_fn=nn.utils.clip_grad_norm_`\n",
    "      grad_clip_params=None, #   with `grad_clip_params={max_norm=1, norm_type=2}`\n",
    "                             # or `grad_clip_fn=nn.utils.clip_grad_value_`\n",
    "                             #   with `grad_clip_params={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",
    ")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "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 all batch-based artefacts to `self.batch`, so Callbacks can find it.\n",
    "- Add extra callbacks, that will use data from `runner.batch` during training.\n",
    "\n",
    "That's it. That's the list. Let's see the example:"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "vVZtc6P61icn"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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 all batch-based artefacts to `self.batch`\n",
    "        # we recommend to use key-value storage to make it Callbacks-friendly\n",
    "        self.batch = {\"features\": x, \"targets\": y, \"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",
    "    # alias for \n",
    "    # `runner.batch_metrics[metric_key] = \\\n",
    "    #     runner.criterion[criterion_key](runner.batch[input_key], runner.batch[target_key])`\n",
    "    \"criterion\": dl.CriterionCallback(  # special Callback for criterion computation\n",
    "      input_key=\"logits\",               # `input_key` specifies model predictions (`y_pred`) from `runner.batch`\n",
    "      target_key=\"targets\",             # `target_key` specifies correct labels (or `y_true`) from `runner.batch`       \n",
    "      metric_key=\"loss\",                # `metric_key` - key to use with `runner.batch_metrics`\n",
    "      criterion_key=None,               # `criterion_key` specifies criterion in case of key-value runner.criterion\n",
    "                                        #   if `criterion_key=None`, runner.criterion used for computation\n",
    "    ), \n",
    "    # alias for \n",
    "    # `runner.batch_metrics[metric_key] = \\\n",
    "    #     metric_fn(runner.batch[input_key], runner.batch[target_key])`\n",
    "    \"metric\": dl.FunctionalMetricCallback( # special Callback for metrics computation\n",
    "      input_key=\"logits\",                  # the same logic as with `CriterionCallback`\n",
    "      target_key=\"targets\",                # the same logic as with `CriterionCallback`\n",
    "      metric_key=\"loss_mae\",               # the same logic as with `CriterionCallback`\n",
    "      metric_fn=F.l1_loss,                 # metric function to use\n",
    "    ),  \n",
    "    \"optimizer\": dl.OptimizerCallback(\n",
    "      metric_key=\"loss\", \n",
    "      accumulation_steps=1,\n",
    "      grad_clip_fn=None,\n",
    "      grad_clip_params=None,\n",
    "    )\n",
    "  }\n",
    ")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "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!"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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, target_key: y} and stored to runner.batch\n",
    "    target_key=\"targets\",      # then the model will be used like\n",
    "    loss_key=\"loss\",           # runner.batch[runner.output_key] = model(runner.batch[input_key])\n",
    ")                              # loss computation suppose to looks like\n",
    "                               # loss = criterion(runner.batch[runner.output_key], runner.batch[runner.target_key])\n",
    "                               # and stored to `runner.batch_metrics[runner.loss_key]`\n",
    "\n",
    "# thanks to prespecified `input_key`, `output_key`, `target_key` and `loss_key`\n",
    "#   `SupervisedRunner` automatically adds required `CriterionCallback` and `OptimizerCallback`\n",
    "# moreover, with specified `logdir`, `valid_loader` and `valid_metric`\n",
    "#   `SupervisedRunner` automatically adds `CheckpointCallback` and tracks best performing based on selected metric\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",
    "  valid_loader=\"valid\",        # `loader_key` from loaders to use for model selection\n",
    "  valid_metric=\"loss\",         # `metric_key` to use for model selection\n",
    "  logdir=\"./logs_supervised\",  # logdir to store models checkpoints\n",
    "  callbacks={\n",
    "#     \"criterion_mse\": dl.CriterionCallback(\n",
    "#       input_key=\"logits\",\n",
    "#       target_key=\"targets\",\n",
    "#       metric_key=\"loss\",\n",
    "#     ),\n",
    "    \"criterion_mae\": dl.FunctionalMetricCallback(\n",
    "      input_key=\"logits\",\n",
    "      target_key=\"targets\",\n",
    "      metric_key=\"mae\",\n",
    "      metric_fn=F.l1_loss,\n",
    "    ),\n",
    "#     \"optimizer\": dl.OptimizerCallback(\n",
    "#       metric_key=\"loss\", \n",
    "#       accumulation_steps=1,\n",
    "#       grad_clip_fn=None,\n",
    "#       grad_clip_params=None,\n",
    "#     )\n",
    "  }\n",
    ")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "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:"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "rn1q6NCP2dtR"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "    # our first time example\n",
    "    def on_loader_start(self, runner):\n",
    "        super().on_loader_start(runner)\n",
    "        self.meters = {\n",
    "            key: metrics.AdditiveMetric(compute_on_call=False)\n",
    "            for key in [\"loss\", \"mae\"]\n",
    "        }\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",
    "        for key in [\"loss\", \"mae\"]:\n",
    "            self.meters[key].update(self.batch_metrics[key].item(), self.batch_size)\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()\n",
    "    \n",
    "    def on_loader_end(self, runner):\n",
    "        for key in [\"loss\", \"mae\"]:\n",
    "            self.loader_metrics[key] = self.meters[key].compute()[0]\n",
    "        super().on_loader_end(runner)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "  valid_loader=\"valid\",   # `loader_key` from loaders to use for model selection\n",
    "  valid_metric=\"loss\",    # `metric_key` to use for model selection\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"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "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!"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "features_batch = next(iter(loaders[\"valid\"]))[0].to(runner.device)\n",
    "# model stochastic weight averaging\n",
    "model.load_state_dict(utils.get_averaged_weights_by_path_mask(logdir=\"./logs\", path_mask=\"*.pth\"))\n",
    "# model tracing\n",
    "utils.trace_model(model=runner.model, batch=features_batch)\n",
    "# model quantization\n",
    "utils.quantize_model(model=runner.model)\n",
    "# model pruning\n",
    "utils.prune_model(model=runner.model, pruning_fn=\"l1_unstructured\", amount=0.8)\n",
    "# onnx export, catalyst[onnx] or catalyst[onnx-gpu] required\n",
    "# utils.onnx_export(model=runner.model, batch=features_batch, file=\"./logs/mnist.onnx\", verbose=True)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "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"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "9-CMpP5a3Wcp"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from catalyst.contrib.nn.modules import Flatten, GlobalMaxPool2d, 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 = {\"generator\": generator, \"discriminator\": discriminator}\n",
    "criterion = {\"generator\": nn.BCEWithLogitsLoss(), \"discriminator\": nn.BCEWithLogitsLoss()}\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",
    "}"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "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`:"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "POY42XRf5Jbd"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "class GANRunner(dl.Runner):\n",
    "  \n",
    "    def __init__(self, latent_dim: int):\n",
    "        super().__init__()\n",
    "        self.latent_dim = latent_dim\n",
    "\n",
    "    def predict_batch(self, batch):\n",
    "        batch_size = 1\n",
    "        # Sample random points in the latent space\n",
    "        random_latent_vectors = torch.randn(batch_size, self.latent_dim).to(self.device)\n",
    "        # Decode them to fake images\n",
    "        generated_images = self.model[\"generator\"](random_latent_vectors).detach()\n",
    "        return generated_images\n",
    "\n",
    "    def handle_batch(self, batch):\n",
    "        real_images, _ = batch\n",
    "        batch_size = real_images.shape[0]\n",
    "\n",
    "        # Sample random points in the latent space\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 = \\\n",
    "            torch.cat([torch.ones((batch_size, 1)), torch.zeros((batch_size, 1))]).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",
    "        # Discriminator forward\n",
    "        combined_predictions = self.model[\"discriminator\"](combined_images)\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",
    "        # Generator forward\n",
    "        generated_images = self.model[\"generator\"](random_latent_vectors)\n",
    "        generated_predictions = self.model[\"discriminator\"](generated_images)\n",
    "\n",
    "        self.batch = {\n",
    "            \"combined_predictions\": combined_predictions,\n",
    "            \"labels\": labels,\n",
    "            \"generated_predictions\": generated_predictions,\n",
    "            \"misleading_labels\": misleading_labels,\n",
    "        }"
   ],
   "outputs": [],
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iyKOtjfn5RL3"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Let's test-drive it:"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "zYGZRIJh6ZYu"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import os\n",
    "from torch.utils.data import DataLoader\n",
    "from catalyst.contrib.datasets import MNIST\n",
    "from catalyst.data import ToTensor\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",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    loaders=loaders,\n",
    "    callbacks=[\n",
    "        dl.CriterionCallback(\n",
    "            input_key=\"combined_predictions\",\n",
    "            target_key=\"labels\",\n",
    "            metric_key=\"loss_discriminator\",\n",
    "            criterion_key=\"discriminator\",\n",
    "        ),\n",
    "        dl.CriterionCallback(\n",
    "            input_key=\"generated_predictions\",\n",
    "            target_key=\"misleading_labels\",\n",
    "            metric_key=\"loss_generator\",\n",
    "            criterion_key=\"generator\",\n",
    "        ),\n",
    "        dl.OptimizerCallback(\n",
    "            model_key=\"generator\", \n",
    "            optimizer_key=\"generator\", \n",
    "            metric_key=\"loss_generator\"\n",
    "        ),\n",
    "        dl.OptimizerCallback(\n",
    "            model_key=\"discriminator\", \n",
    "            optimizer_key=\"discriminator\", \n",
    "            metric_key=\"loss_discriminator\"\n",
    "        ),\n",
    "    ],\n",
    "    valid_loader=\"train\",\n",
    "    valid_metric=\"loss_generator\",\n",
    "    minimize_valid_metric=True,\n",
    "    num_epochs=20,\n",
    "    verbose=True,\n",
    "    logdir=\"./logs_gan\",\n",
    ")"
   ],
   "outputs": [],
   "metadata": {
    "scrolled": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "The idea behind deep learning are simple, so why should their implementation be painful?"
   ],
   "metadata": {
    "colab_type": "text",
    "id": "M9Fz5_u68FqW"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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())"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%load_ext tensorboard\n",
    "%tensorboard --logdir ./logs_gan"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {}
  }
 ],
 "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.9"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}