examples/notebooks/customization_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install catalyst[ml]==22.02rc0 scikit-learn==1.0.0 optuna==2.7.0 --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Customization is all you need\n",
    "\n",
    "This demo shows several examples for Catalyst runs customization for a variety of needs. If you want to get used to [Catalyst](https://github.com/catalyst-team/catalyst) first, please follow the [minimal examples section](https://github.com/catalyst-team/catalyst#minimal-examples) first."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start from, let's prepare simple data to work with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import numpy as np\n",
    "from sklearn.datasets import make_moons, make_blobs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import *\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "\n",
    "from catalyst import dl, metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "matplotlib.use('TkAgg')\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make up a dataset\n",
    "def make_dataset(seed=42, n_samples=int(1e3)):\n",
    "    np.random.seed(seed)\n",
    "    random.seed(seed)\n",
    "    X, y = make_moons(n_samples=n_samples, noise=0.1)\n",
    "\n",
    "    y = y*2 - 1 # make y be -1 or 1\n",
    "    return X, y\n",
    "\n",
    "def visualize_dataset(X, y):\n",
    "    plt.figure(figsize=(5,5))\n",
    "    plt.scatter(X[:,0], X[:,1], c=y, s=20, cmap='jet')\n",
    "\n",
    "# let's create train data\n",
    "X_train, y_train = make_dataset()\n",
    "visualize_dataset(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# valid data\n",
    "X_valid, y_valid = make_dataset(seed=137)\n",
    "visualize_dataset(X_valid, y_valid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# and another train one (why not?)\n",
    "X_train2, y_train2 = make_dataset(seed=1337)\n",
    "visualize_dataset(X_train2, y_train2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize a model \n",
    "# 2-layer neural network\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(2, 16), nn.ReLU(), \n",
    "    nn.Linear(16, 16), nn.ReLU(), \n",
    "    nn.Linear(16, 1)\n",
    ")\n",
    "print(model)\n",
    "# print(\"number of parameters\", len(model.parameters()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_decision_boundary(X, y, model):\n",
    "    h = 0.25\n",
    "    x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1\n",
    "    y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1\n",
    "    xx, yy = np.meshgrid(np.arange(x_min, x_max, h),\n",
    "                         np.arange(y_min, y_max, h))\n",
    "    Xmesh = np.c_[xx.ravel(), yy.ravel()]\n",
    "    \n",
    "    inputs = torch.tensor([list(xrow) for xrow in Xmesh]).float()\n",
    "    scores = model(inputs)\n",
    "    \n",
    "    Z = np.array([s.data > 0 for s in scores])\n",
    "    Z = Z.reshape(xx.shape)\n",
    "\n",
    "    fig = plt.figure()\n",
    "    plt.contourf(xx, yy, Z, cmap=plt.cm.Spectral, alpha=0.8)\n",
    "    plt.scatter(X[:, 0], X[:, 1], c=y, s=40, cmap=plt.cm.Spectral)\n",
    "    plt.xlim(xx.min(), xx.max())\n",
    "    plt.ylim(yy.min(), yy.max())\n",
    "    plt.show()\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = visualize_decision_boundary(X_valid, y_valid, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "t1 = TensorDataset(torch.tensor(X_train).float(), torch.tensor(y_train > 0).float())\n",
    "t2 = TensorDataset(torch.tensor(X_train2).float(), torch.tensor(y_train2 > 0).float())\n",
    "v1 = TensorDataset(torch.tensor(X_valid).float(), torch.tensor(y_valid > 0).float())\n",
    "\n",
    "loaders = {\n",
    "    \"train_1\": DataLoader(t1, batch_size=32, num_workers=1), \n",
    "    \"train_2\": DataLoader(t2, batch_size=32, num_workers=1), \n",
    "    \"valid\": DataLoader(v1, batch_size=32, num_workers=1), \n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Act 1 - ``CustomRunner – batch handling by you own``"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First case – defining everything by hands and `.run`. PyTorch for-loop wrapper example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class CustomRunner(dl.IRunner):ß\n",
    "    @property\n",
    "    def num_epochs(self) -> int:\n",
    "        return 5\n",
    "\n",
    "    def get_engine(self) -> dl.Engine:\n",
    "        return dl.CPUEngine()    \n",
    "    \n",
    "    def get_loaders(self) -> \"OrderedDict[str, DataLoader]\":\n",
    "        return loaders\n",
    "    \n",
    "    def get_model(self):\n",
    "        return nn.Sequential(\n",
    "            nn.Linear(2, 16), nn.ReLU(), \n",
    "            nn.Linear(16, 16), nn.ReLU(), \n",
    "            nn.Linear(16, 1)\n",
    "        )\n",
    "\n",
    "    def get_optimizer(self, model):\n",
    "        return torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "    \n",
    "    def handle_batch(self, batch):\n",
    "        x, y = batch\n",
    "        y_hat = self.model(x)\n",
    "\n",
    "        loss = F.binary_cross_entropy_with_logits(y_hat.view(-1), y)\n",
    "        self.batch_metrics = {\"loss\": loss}\n",
    "        if self.loader_batch_step % 10 == 0:\n",
    "            print(\n",
    "                f\"{self.loader_key} ({self.loader_batch_step}/{self.loader_batch_len}:\" \n",
    "                f\"loss {loss.item()}\"\n",
    "            )\n",
    "\n",
    "        if self.is_train_loader:\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "            self.optimizer.zero_grad()\n",
    "\n",
    "runner = CustomRunner().run()\n",
    "model = runner.model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = visualize_decision_boundary(X_valid, y_valid, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Act 2 - ``SupervisedRunner – Runner with Callbacks``"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Same example, but with extra loggers and callbacks to simplify the run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class CustomSupervisedRunner(dl.IRunner):\n",
    "    @property\n",
    "    def num_epochs(self) -> int:\n",
    "        return 5\n",
    "\n",
    "    def get_engine(self) -> dl.Engine:\n",
    "        return dl.CPUEngine()\n",
    "    \n",
    "    def get_loggers(self):\n",
    "        return {\n",
    "            \"console\": dl.ConsoleLogger(),\n",
    "    #         \"csv\": dl.LogdirLogger(logdir=\"./logdir02\"),\n",
    "            \"tensorboard\": dl.TensorboardLogger(logdir=\"./logdir02/tb\"),\n",
    "        }\n",
    "    \n",
    "    def get_loaders(self) -> \"OrderedDict[str, DataLoader]\":\n",
    "        return loaders\n",
    "    \n",
    "    def get_model(self):\n",
    "        return nn.Sequential(\n",
    "            nn.Linear(2, 16), nn.ReLU(), \n",
    "            nn.Linear(16, 16), nn.ReLU(), \n",
    "            nn.Linear(16, 1)\n",
    "        )\n",
    "\n",
    "    def get_criterion(self):\n",
    "        return nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def get_optimizer(self, model):\n",
    "        return torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "    def get_scheduler(self, optimizer):\n",
    "        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [2, 4])\n",
    "    \n",
    "    def get_callbacks(self):\n",
    "        return {\n",
    "            # Let's use AUC metric as an example – it's loader-based, so we shouldn't compute it on each batch\n",
    "            \"auc\": dl.LoaderMetricCallback(\n",
    "                metric=metrics.AUCMetric(),\n",
    "                input_key=\"scores\", target_key=\"targets\", \n",
    "            ), \n",
    "            # To wrap the criterion step logic, you could use CriterionCallback:\n",
    "            \"criterion\": dl.CriterionCallback(\n",
    "                metric_key=\"loss\", \n",
    "                input_key=\"logits\", \n",
    "                target_key=\"targets\"\n",
    "            ), \n",
    "            # To wrap the optimizer step logic, you could use BackwardCallback & OptimizerCallback:\n",
    "            \"backward\": dl.BackwardCallback(metric_key=\"loss\"), \n",
    "            \"optimizer\": dl.OptimizerCallback(metric_key=\"loss\"), \n",
    "            # The same case with the scheduler:\n",
    "            \"scheduler\": dl.SchedulerCallback(\n",
    "                loader_key=\"valid\", metric_key=\"loss\"\n",
    "            ),\n",
    "            # We could also use lrfinder for lr scheduling:\n",
    "#             \"lr-finder\": dl.LRFinder(\n",
    "#                 final_lr=1.0,\n",
    "#                 scale=\"log\",\n",
    "#                 num_steps=None,\n",
    "#                 optimizer_key=None,\n",
    "#             ),\n",
    "            # You can select any number of metrics to checkpoint on:\n",
    "            \"checkpoint1\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir02/auc\",\n",
    "                loader_key=\"valid\", metric_key=\"auc\", \n",
    "                minimize=False, topk=3\n",
    "            ),\n",
    "            \"checkpoint2\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir02/loss\",\n",
    "                loader_key=\"valid\", metric_key=\"loss\", \n",
    "                minimize=True, topk=1\n",
    "            ),\n",
    "            # Or turn on/off tqdm verbose during loader run:\n",
    "            \"verbose\": dl.TqdmCallback(),\n",
    "        }\n",
    "    \n",
    "    def handle_batch(self, batch):\n",
    "        x, y = batch\n",
    "        y_hat = self.model(x)\n",
    "        \n",
    "        self.batch = {\n",
    "            \"features\": x,\n",
    "            \"targets\": y,\n",
    "            \"logits\": y_hat.view(-1),\n",
    "            \"scores\": torch.sigmoid(y_hat.view(-1)),\n",
    "        }\n",
    "\n",
    "runner = CustomSupervisedRunner().run()\n",
    "model = runner.model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = visualize_decision_boundary(X_valid, y_valid, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Act 3 - ``CustomMetric``"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose you would like to add some custom metric to your pipeline..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomAccuracyMetric(metrics.ICallbackBatchMetric, metrics.AdditiveMetric):\n",
    "    def update(self, scores: torch.Tensor, targets: torch.Tensor) -> float:\n",
    "        value = ((scores > 0.5) == targets).float().mean().item()\n",
    "        value = super().update(value, len(targets))\n",
    "        return value\n",
    "    \n",
    "    def update_key_value(self, scores: torch.Tensor, targets: torch.Tensor) -> Dict[str, float]:\n",
    "        value = self.update(scores, targets)\n",
    "        return {\"accuracy\": value}\n",
    "\n",
    "    def compute_key_value(self) -> Dict[str, float]:\n",
    "        mean, std = super().compute()\n",
    "        return {\"accuracy\": mean, \"accuracy/std\": std}\n",
    "\n",
    "    \n",
    "class CustomSupervisedRunner(dl.IRunner):\n",
    "    @property\n",
    "    def num_epochs(self) -> int:\n",
    "        return 5\n",
    "\n",
    "    def get_engine(self) -> dl.Engine:\n",
    "        return dl.CPUEngine()\n",
    "    \n",
    "    def get_loggers(self):\n",
    "        return {\n",
    "            \"console\": dl.ConsoleLogger(),\n",
    "            \"tensorboard\": dl.TensorboardLogger(logdir=\"./logdir03/tb\"),\n",
    "        }\n",
    "    \n",
    "    def get_loaders(self) -> \"OrderedDict[str, DataLoader]\":\n",
    "        return loaders\n",
    "    \n",
    "    def get_model(self):\n",
    "        return nn.Sequential(\n",
    "            nn.Linear(2, 16), nn.ReLU(), \n",
    "            nn.Linear(16, 16), nn.ReLU(), \n",
    "            nn.Linear(16, 1)\n",
    "        )\n",
    "\n",
    "    def get_criterion(self):\n",
    "        return nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def get_optimizer(self, model):\n",
    "        return torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "    def get_scheduler(self, optimizer):\n",
    "        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [2, 4])\n",
    "    \n",
    "    def get_callbacks(self):\n",
    "        return {\n",
    "            \"accuracy\": dl.BatchMetricCallback(\n",
    "                metric=CustomAccuracyMetric(), log_on_batch=True,\n",
    "                input_key=\"scores\", target_key=\"targets\", \n",
    "            ),\n",
    "            \"auc\": dl.LoaderMetricCallback(\n",
    "                metric=metrics.AUCMetric(),\n",
    "                input_key=\"scores\", target_key=\"targets\", \n",
    "            ), \n",
    "            \"criterion\": dl.CriterionCallback(\n",
    "                metric_key=\"loss\", \n",
    "                input_key=\"logits\", \n",
    "                target_key=\"targets\"\n",
    "            ), \n",
    "            \"backward\": dl.BackwardCallback(metric_key=\"loss\"), \n",
    "            \"optimizer\": dl.OptimizerCallback(metric_key=\"loss\"), \n",
    "            \"scheduler\": dl.SchedulerCallback(\n",
    "                loader_key=\"valid\", metric_key=\"loss\"\n",
    "            ),\n",
    "            \"checkpoint1\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir03/accuracy\",\n",
    "                loader_key=\"valid\", metric_key=\"accuracy\", \n",
    "                minimize=False, topk=3\n",
    "            ),\n",
    "            \"checkpoint2\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir03/loss\",\n",
    "                loader_key=\"valid\", metric_key=\"loss\", \n",
    "                minimize=True, topk=1\n",
    "            ),\n",
    "    #         \"verbose\": dl.TqdmCallback(),\n",
    "        }\n",
    "    \n",
    "    def handle_batch(self, batch):\n",
    "        x, y = batch\n",
    "        y_hat = self.model(x)\n",
    "        \n",
    "        self.batch = {\n",
    "            \"features\": x,\n",
    "            \"targets\": y,\n",
    "            \"logits\": y_hat.view(-1),\n",
    "            \"scores\": torch.sigmoid(y_hat.view(-1)),\n",
    "        }\n",
    "\n",
    "runner = CustomSupervisedRunner().run()\n",
    "model = runner.model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = visualize_decision_boundary(X_valid, y_valid, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Act 4 - ``CustomCallback``"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or some custom functionality – Callback."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Let's plot the decision doundary after each epoch:\n",
    "class VisualizationCallback(dl.Callback):\n",
    "    def __init__(self):\n",
    "        super().__init__(order=dl.CallbackOrder.External)\n",
    "\n",
    "    def on_epoch_end(self, runner):\n",
    "        img = visualize_decision_boundary(X_valid, y_valid, runner.model)\n",
    "\n",
    "\n",
    "class CustomSupervisedRunner(dl.IRunner):\n",
    "    @property\n",
    "    def num_epochs(self) -> int:\n",
    "        return 5\n",
    "\n",
    "    def get_engine(self) -> dl.Engine:\n",
    "        return dl.CPUEngine()\n",
    "    \n",
    "    def get_loggers(self):\n",
    "        return {\n",
    "            \"console\": dl.ConsoleLogger(),\n",
    "            \"tensorboard\": dl.TensorboardLogger(logdir=\"./logdir04/tb\"),\n",
    "        }\n",
    "    \n",
    "    def get_loaders(self) -> \"OrderedDict[str, DataLoader]\":\n",
    "        return loaders\n",
    "    \n",
    "    def get_model(self):\n",
    "        return nn.Sequential(\n",
    "            nn.Linear(2, 16), nn.ReLU(), \n",
    "            nn.Linear(16, 16), nn.ReLU(), \n",
    "            nn.Linear(16, 1)\n",
    "        )\n",
    "\n",
    "    def get_criterion(self):\n",
    "        return nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def get_optimizer(self, model):\n",
    "        return torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "    def get_scheduler(self, optimizer):\n",
    "        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [2, 4])\n",
    "    \n",
    "    def get_callbacks(self):\n",
    "        return {\n",
    "            \"criterion\": dl.CriterionCallback(\n",
    "                metric_key=\"loss\", \n",
    "                input_key=\"logits\", \n",
    "                target_key=\"targets\"\n",
    "            ), \n",
    "            \"backward\": dl.BackwardCallback(metric_key=\"loss\"), \n",
    "            \"optimizer\": dl.OptimizerCallback(metric_key=\"loss\"), \n",
    "            \"scheduler\": dl.SchedulerCallback(\n",
    "                loader_key=\"valid\", metric_key=\"loss\"\n",
    "            ),\n",
    "            \"checkpoint\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir04/loss\",\n",
    "                loader_key=\"valid\", metric_key=\"loss\", \n",
    "                minimize=True, topk=1\n",
    "            ),\n",
    "            # And include it into callbacks:        \n",
    "            \"visualization\": VisualizationCallback()\n",
    "        }\n",
    "    \n",
    "    def handle_batch(self, batch):\n",
    "        x, y = batch\n",
    "        y_hat = self.model(x)\n",
    "        \n",
    "        self.batch = {\n",
    "            \"features\": x,\n",
    "            \"targets\": y,\n",
    "            \"logits\": y_hat.view(-1),\n",
    "            \"scores\": torch.sigmoid(y_hat.view(-1)),\n",
    "        }\n",
    "\n",
    "runner = CustomSupervisedRunner().run()\n",
    "model = runner.model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Act 5 - ``CustomLogger``"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But what if you also want custom logging? For your private pipeline storage, maybe?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import cv2\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def get_img_from_fig(fig, dpi=180):\n",
    "    buf = io.BytesIO()\n",
    "    fig.savefig(buf, format=\"png\", dpi=dpi)\n",
    "    buf.seek(0)\n",
    "    \n",
    "    img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)\n",
    "    buf.close()\n",
    "    img = cv2.imdecode(img_arr, 1)\n",
    "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We need to add only a few lines to log the image to all runner's loggers\n",
    "class VisualizationCallback(dl.Callback):\n",
    "    def __init__(self):\n",
    "        super().__init__(order=dl.CallbackOrder.External)\n",
    "\n",
    "    def on_epoch_end(self, runner):\n",
    "        image = visualize_decision_boundary(X_valid, y_valid, runner.model)\n",
    "        image = get_img_from_fig(image)\n",
    "        # runner will propagate it to all loggers\n",
    "        runner.log_image(tag=\"decision_boundary\", image=image, scope=\"epoch\")\n",
    "\n",
    "\n",
    "# Let's also add our own Logger to store image on the disk\n",
    "class VisualizationLogger(dl.ILogger):\n",
    "    def __init__(self, logdir: str):\n",
    "        self.logdir = logdir\n",
    "        os.makedirs(self.logdir, exist_ok=True)\n",
    "        \n",
    "    def log_image(\n",
    "        self,\n",
    "        tag: str,\n",
    "        image: np.ndarray,\n",
    "        runner,\n",
    "        scope: str = None,\n",
    "    ) -> None:\n",
    "        if scope == \"epoch\":\n",
    "            plt.imsave(\n",
    "                os.path.join(self.logdir, f\"{tag}_{runner.epoch_step}.png\"),\n",
    "                image,\n",
    "            )\n",
    "\n",
    "\n",
    "class CustomSupervisedRunner(dl.IRunner):\n",
    "    @property\n",
    "    def num_epochs(self) -> int:\n",
    "        return 5\n",
    "\n",
    "    def get_engine(self) -> dl.Engine:\n",
    "        return dl.CPUEngine()\n",
    "    \n",
    "    def get_loggers(self):\n",
    "        return {\n",
    "            \"console\": dl.ConsoleLogger(),\n",
    "            \"visualization\": VisualizationLogger(logdir=\"./logdir05/visualization\"),\n",
    "            \"tensorboard\": dl.TensorboardLogger(logdir=\"./logdir05/tb\"),\n",
    "        }\n",
    "    \n",
    "    def get_loaders(self) -> \"OrderedDict[str, DataLoader]\":\n",
    "        return loaders\n",
    "    \n",
    "    def get_model(self):\n",
    "        return nn.Sequential(\n",
    "            nn.Linear(2, 16), nn.ReLU(), \n",
    "            nn.Linear(16, 16), nn.ReLU(), \n",
    "            nn.Linear(16, 1)\n",
    "        )\n",
    "\n",
    "    def get_criterion(self):\n",
    "        return nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def get_optimizer(self, model):\n",
    "        return torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "    def get_scheduler(self, optimizer):\n",
    "        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [2, 4])\n",
    "    \n",
    "    def get_callbacks(self):\n",
    "        return {\n",
    "            \"criterion\": dl.CriterionCallback(\n",
    "                metric_key=\"loss\", \n",
    "                input_key=\"logits\", \n",
    "                target_key=\"targets\"\n",
    "            ), \n",
    "            \"backward\": dl.BackwardCallback(metric_key=\"loss\"), \n",
    "            \"optimizer\": dl.OptimizerCallback(metric_key=\"loss\"), \n",
    "            \"scheduler\": dl.SchedulerCallback(\n",
    "                loader_key=\"valid\", metric_key=\"loss\"\n",
    "            ),\n",
    "            \"checkpoint\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir05/loss\",\n",
    "                loader_key=\"valid\", metric_key=\"loss\", \n",
    "                minimize=True, topk=1\n",
    "            ),\n",
    "            \"visualization\": VisualizationCallback()\n",
    "        }\n",
    "    \n",
    "    def handle_batch(self, batch):\n",
    "        x, y = batch\n",
    "        y_hat = self.model(x)\n",
    "        \n",
    "        self.batch = {\n",
    "            \"features\": x,\n",
    "            \"targets\": y,\n",
    "            \"logits\": y_hat.view(-1),\n",
    "            \"scores\": torch.sigmoid(y_hat.view(-1)),\n",
    "        }\n",
    "\n",
    "runner = CustomSupervisedRunner().run()\n",
    "model = runner.model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! ls ./logdir05\n",
    "! ls ./logdir05/loss\n",
    "! ls ./logdir05/tb\n",
    "! ls ./logdir05/visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Act 6 - Confusion Matrix logging - IMetric+ICallback+ILogger"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the end, let's review a simple example with a bunch of metrics, loggers, and confusion matrix computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# useful in colab:\n",
    "# %load_ext tensorboard\n",
    "# tensorboard --logdir=./logdir06/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from catalyst import dl, metrics, utils\n",
    "\n",
    "# sample data\n",
    "num_samples, num_features, num_classes = int(1e4), int(1e1), 6\n",
    "num_epochs = 6\n",
    "\n",
    "class CustomSupervisedRunner(dl.IRunner):\n",
    "    @property\n",
    "    def num_epochs(self) -> int:\n",
    "        return num_epochs\n",
    "\n",
    "    def get_engine(self) -> dl.Engine:\n",
    "        return dl.CPUEngine()\n",
    "    \n",
    "    def get_loggers(self):\n",
    "        return {\n",
    "            \"console\": dl.ConsoleLogger(),\n",
    "            \"csv\": dl.CSVLogger(logdir=\"./logdir06\"),\n",
    "            \"tensorboard\": dl.TensorboardLogger(logdir=\"./logdir06/tb\"),\n",
    "        }\n",
    "    \n",
    "    def get_loaders(self) -> \"OrderedDict[str, DataLoader]\":\n",
    "        # sample data\n",
    "        num_samples, num_features, num_classes = int(1e4), int(1e1), 6\n",
    "        X = torch.rand(num_samples, num_features)\n",
    "        y = (torch.rand(num_samples, ) * num_classes).to(torch.int64)\n",
    "\n",
    "        # pytorch loaders\n",
    "        dataset = TensorDataset(X, y)\n",
    "        loader = DataLoader(dataset, batch_size=32, num_workers=1)\n",
    "        loaders = {\"train\": loader, \"valid\": loader}\n",
    "        return loaders\n",
    "    \n",
    "    def get_model(self):\n",
    "        return torch.nn.Linear(num_features, num_classes)\n",
    "\n",
    "    def get_criterion(self):\n",
    "        return torch.nn.CrossEntropyLoss()\n",
    "\n",
    "    def get_optimizer(self, model):\n",
    "        return torch.optim.Adam(model.parameters())\n",
    "\n",
    "    def get_scheduler(self, optimizer):\n",
    "        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [2])\n",
    "    \n",
    "    def get_callbacks(self):\n",
    "        return {\n",
    "            \"accuracy\": dl.BatchMetricCallback(\n",
    "                metric=metrics.AccuracyMetric(num_classes=num_classes),\n",
    "                input_key=\"probs\", target_key=\"targets\", \n",
    "            ),\n",
    "            \"auc\": dl.LoaderMetricCallback(\n",
    "                metric=metrics.AUCMetric(),\n",
    "                input_key=\"scores\", target_key=\"targets\", \n",
    "            ), \n",
    "            \"criterion\": dl.CriterionCallback(\n",
    "                metric_key=\"loss\", \n",
    "                input_key=\"logits\", \n",
    "                target_key=\"targets\",\n",
    "            ), \n",
    "            \"backward\": dl.BackwardCallback(metric_key=\"loss\"), \n",
    "            \"optimizer\": dl.OptimizerCallback(metric_key=\"loss\"), \n",
    "            \"scheduler\": dl.SchedulerCallback(\n",
    "                loader_key=\"valid\", metric_key=\"loss\"\n",
    "            ),\n",
    "            \"checkpoint1\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir06/loss\",\n",
    "                loader_key=\"valid\", metric_key=\"loss\", \n",
    "                minimize=False, topk=3\n",
    "            ),\n",
    "            \"checkpoint2\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir06/auc\",\n",
    "                loader_key=\"valid\", metric_key=\"auc\", \n",
    "                minimize=True, topk=1\n",
    "            ),\n",
    "            \"checkpoint3\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir06/accuracy\",\n",
    "                loader_key=\"valid\", metric_key=\"accuracy01\", \n",
    "                minimize=True, topk=1,\n",
    "            ),\n",
    "            \"verbose\": dl.TqdmCallback(),\n",
    "            \"confusion_matrix\": dl.ConfusionMatrixCallback(\n",
    "                input_key=\"probs\", \n",
    "                target_key=\"targets\",\n",
    "                prefix=\"confusion_matrix\",\n",
    "                num_classes=num_classes,\n",
    "            )\n",
    "        }\n",
    "    \n",
    "    def handle_batch(self, batch):\n",
    "        x, y = batch\n",
    "        y_hat = self.model(x)\n",
    "        \n",
    "        self.batch = {\n",
    "            \"features\": x,\n",
    "            \"targets\": y,\n",
    "            \"logits\": y_hat,\n",
    "            \"scores\": torch.sigmoid(y_hat),\n",
    "            \"probs\": torch.softmax(y_hat, dim=1),\n",
    "        }\n",
    "\n",
    "runner = CustomSupervisedRunner().run()\n",
    "model = runner.model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congrats! You have finished the Catalyst customization tutorial and are now ready to create your own different pipelines and R&D breakthroughs. See you at the next stage!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:py37-dev]",
   "language": "python",
   "name": "conda-env-py37-dev-py"
  },
  "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}