examples/notebooks/customization_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install catalyst[ml]==21.04.2 scikit-learn==0.24.1 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, utils"
   ]
  },
  {
   "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",
    "    def get_engine(self) -> dl.IEngine:\n",
    "        return dl.DeviceEngine(\"cpu\")\n",
    "    \n",
    "    @property\n",
    "    def stages(self) -> Iterable[str]:\n",
    "        return [\"train\"]\n",
    "    \n",
    "    def get_stage_len(self, stage: str) -> int:\n",
    "        return 5\n",
    "    \n",
    "    def get_loaders(self, stage: str) -> \"OrderedDict[str, DataLoader]\":\n",
    "        return loaders\n",
    "    \n",
    "    def get_model(self, stage: str):\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, stage: str):\n",
    "        return None\n",
    "\n",
    "    def get_optimizer(self, stage: str, model):\n",
    "        return torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "    def get_scheduler(self, stage: str, optimizer):\n",
    "        return None\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",
    "    def get_engine(self) -> dl.IEngine:\n",
    "        return dl.DeviceEngine(\"cpu\")\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",
    "    @property\n",
    "    def stages(self) -> Iterable[str]:\n",
    "        return [\"train\"]\n",
    "    \n",
    "    def get_stage_len(self, stage: str) -> int:\n",
    "        return 5\n",
    "    \n",
    "    def get_loaders(self, stage: str) -> \"OrderedDict[str, DataLoader]\":\n",
    "        return loaders\n",
    "    \n",
    "    def get_model(self, stage: str):\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, stage: str):\n",
    "        return nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def get_optimizer(self, stage: str, model):\n",
    "        return torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "    def get_scheduler(self, stage: str, optimizer):\n",
    "        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [2, 4])\n",
    "    \n",
    "    def get_callbacks(self, stage: str):\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 OptimizerCallback:\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, save_n_best=3\n",
    "            ),\n",
    "            \"checkpoint2\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir02/loss\",\n",
    "                loader_key=\"valid\", metric_key=\"loss\", \n",
    "                minimize=True, save_n_best=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.AdditiveValueMetric):\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",
    "    def get_engine(self) -> dl.IEngine:\n",
    "        return dl.DeviceEngine(\"cpu\")\n",
    "    \n",
    "    def get_loggers(self):\n",
    "        return {\n",
    "            \"console\": dl.ConsoleLogger(),\n",
    "            \"tensorboard\": dl.TensorboardLogger(logdir=\"./logdir03/tb\"),\n",
    "        }\n",
    "    \n",
    "    @property\n",
    "    def stages(self) -> Iterable[str]:\n",
    "        return [\"train\"]\n",
    "    \n",
    "    def get_stage_len(self, stage: str) -> int:\n",
    "        return 5\n",
    "    \n",
    "    def get_loaders(self, stage: str) -> \"OrderedDict[str, DataLoader]\":\n",
    "        return loaders\n",
    "    \n",
    "    def get_model(self, stage: str):\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, stage: str):\n",
    "        return nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def get_optimizer(self, stage: str, model):\n",
    "        return torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "    def get_scheduler(self, stage: str, optimizer):\n",
    "        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [2, 4])\n",
    "    \n",
    "    def get_callbacks(self, stage: str):\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",
    "            \"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, save_n_best=3\n",
    "            ),\n",
    "            \"checkpoint2\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir03/loss\",\n",
    "                loader_key=\"valid\", metric_key=\"loss\", \n",
    "                minimize=True, save_n_best=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",
    "    def get_engine(self) -> dl.IEngine:\n",
    "        return dl.DeviceEngine(\"cpu\")\n",
    "    \n",
    "    def get_loggers(self):\n",
    "        return {\n",
    "            \"console\": dl.ConsoleLogger(),\n",
    "            \"tensorboard\": dl.TensorboardLogger(logdir=\"./logdir04/tb\"),\n",
    "        }\n",
    "    \n",
    "    @property\n",
    "    def stages(self) -> Iterable[str]:\n",
    "        return [\"train\"]\n",
    "    \n",
    "    def get_stage_len(self, stage: str) -> int:\n",
    "        return 5\n",
    "    \n",
    "    def get_loaders(self, stage: str) -> \"OrderedDict[str, DataLoader]\":\n",
    "        return loaders\n",
    "    \n",
    "    def get_model(self, stage: str):\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, stage: str):\n",
    "        return nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def get_optimizer(self, stage: str, model):\n",
    "        return torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "    def get_scheduler(self, stage: str, optimizer):\n",
    "        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [2, 4])\n",
    "    \n",
    "    def get_callbacks(self, stage: str):\n",
    "        return {\n",
    "            \"criterion\": dl.CriterionCallback(\n",
    "                metric_key=\"loss\", \n",
    "                input_key=\"logits\", \n",
    "                target_key=\"targets\"\n",
    "            ), \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, save_n_best=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",
    "        scope: str = None,\n",
    "        # experiment info\n",
    "        run_key: str = None,\n",
    "        global_epoch_step: int = 0,\n",
    "        global_batch_step: int = 0,\n",
    "        global_sample_step: int = 0,\n",
    "        # stage info\n",
    "        stage_key: str = None,\n",
    "        stage_epoch_len: int = 0,\n",
    "        stage_epoch_step: int = 0,\n",
    "        stage_batch_step: int = 0,\n",
    "        stage_sample_step: int = 0,\n",
    "        # loader info\n",
    "        loader_key: str = None,\n",
    "        loader_batch_len: int = 0,\n",
    "        loader_sample_len: int = 0,\n",
    "        loader_batch_step: int = 0,\n",
    "        loader_sample_step: int = 0,\n",
    "    ) -> None:\n",
    "        if scope == \"epoch\":\n",
    "            plt.imsave(\n",
    "                os.path.join(self.logdir, f\"{tag}_{stage_key}_{stage_epoch_step}.png\"),\n",
    "                image,\n",
    "            )\n",
    "\n",
    "\n",
    "class CustomSupervisedRunner(dl.IRunner):\n",
    "    def get_engine(self) -> dl.IEngine:\n",
    "        return dl.DeviceEngine(\"cpu\")\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",
    "    @property\n",
    "    def stages(self) -> Iterable[str]:\n",
    "        return [\"train\"]\n",
    "    \n",
    "    def get_stage_len(self, stage: str) -> int:\n",
    "        return 5\n",
    "    \n",
    "    def get_loaders(self, stage: str) -> \"OrderedDict[str, DataLoader]\":\n",
    "        return loaders\n",
    "    \n",
    "    def get_model(self, stage: str):\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, stage: str):\n",
    "        return nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def get_optimizer(self, stage: str, model):\n",
    "        return torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "    def get_scheduler(self, stage: str, optimizer):\n",
    "        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [2, 4])\n",
    "    \n",
    "    def get_callbacks(self, stage: str):\n",
    "        return {\n",
    "            \"criterion\": dl.CriterionCallback(\n",
    "                metric_key=\"loss\", \n",
    "                input_key=\"logits\", \n",
    "                target_key=\"targets\"\n",
    "            ), \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, save_n_best=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": [
    "### Act 6 - ``Multistage Run``"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Last but not least, let's make things more complex:\n",
    "- multistage pipeline\n",
    "- with different train data"
   ]
  },
  {
   "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": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "loaders = {\n",
    "    \"stage_1\": {\n",
    "        \"train_1\": DataLoader(t1, batch_size=32, num_workers=1), \n",
    "        \"valid\": DataLoader(v1, batch_size=32, num_workers=1), \n",
    "    },\n",
    "    \"stage_2\": {\n",
    "        \"train_2\": DataLoader(t2, batch_size=32, num_workers=1), \n",
    "        \"valid\": DataLoader(v1, batch_size=32, num_workers=1), \n",
    "    },\n",
    "}\n",
    "\n",
    "    \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",
    "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",
    "        scope: str = None,\n",
    "        # experiment info\n",
    "        experiment_key: str = None,\n",
    "        global_epoch_step: int = 0,\n",
    "        global_batch_step: int = 0,\n",
    "        global_sample_step: int = 0,\n",
    "        # stage info\n",
    "        stage_key: str = None,\n",
    "        stage_epoch_len: int = 0,\n",
    "        stage_epoch_step: int = 0,\n",
    "        stage_batch_step: int = 0,\n",
    "        stage_sample_step: int = 0,\n",
    "        # loader info\n",
    "        loader_key: str = None,\n",
    "        loader_batch_len: int = 0,\n",
    "        loader_sample_len: int = 0,\n",
    "        loader_batch_step: int = 0,\n",
    "        loader_sample_step: int = 0,\n",
    "    ) -> None:\n",
    "        if scope == \"epoch\":\n",
    "            plt.imsave(\n",
    "                os.path.join(self.logdir, f\"{tag}_{stage_key}_{stage_epoch_step}.png\"),\n",
    "                image,\n",
    "            )\n",
    "\n",
    "class CustomSupervisedRunner(dl.IRunner):\n",
    "    def get_engine(self) -> dl.IEngine:\n",
    "        return dl.DeviceEngine(\"cpu\")\n",
    "    \n",
    "    def get_loggers(self):\n",
    "        return {\n",
    "            \"console\": dl.ConsoleLogger(),\n",
    "            \"visualization\": VisualizationLogger(logdir=\"./logdir06/visualization\"),\n",
    "            \"tensorboard\": dl.TensorboardLogger(logdir=\"./logdir06/tb\"),\n",
    "        }\n",
    "    \n",
    "    @property\n",
    "    def stages(self) -> Iterable[str]:\n",
    "        return loaders.keys()\n",
    "    \n",
    "    def get_stage_len(self, stage: str) -> int:\n",
    "        return 5\n",
    "    \n",
    "    def get_loaders(self, stage: str) -> \"OrderedDict[str, DataLoader]\":\n",
    "        return loaders[stage]\n",
    "    \n",
    "    def get_model(self, stage: str):\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, stage: str):\n",
    "        return nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def get_optimizer(self, stage: str, model):\n",
    "        return torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "    def get_scheduler(self, stage: str, optimizer):\n",
    "        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [2, 4])\n",
    "    \n",
    "    def get_callbacks(self, stage: str):\n",
    "        return {\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",
    "            \"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/auc\",\n",
    "                loader_key=\"valid\", metric_key=\"auc\", \n",
    "                minimize=False, save_n_best=3\n",
    "            ),\n",
    "            \"checkpoint2\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir06/loss\",\n",
    "                loader_key=\"valid\", metric_key=\"loss\", \n",
    "                minimize=True, save_n_best=1\n",
    "            ),\n",
    "            \"visualization\": VisualizationCallback(),\n",
    "    #         \"verbose\": TqdmCallback(),\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.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 ./logdir06"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "_ = visualize_decision_boundary(X_valid, y_valid, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Act 7 - ``CustomRunner``"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sorry, I mean **complex**:\n",
    "- multistage pipeline\n",
    "- with different train data\n",
    "- **and different optimizers/schedulers per each stage**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class CustomSupervisedRunner(dl.IRunner):\n",
    "    def get_engine(self) -> dl.IEngine:\n",
    "        return dl.DeviceEngine(\"cpu\")\n",
    "    \n",
    "    def get_loggers(self):\n",
    "        return {\n",
    "            \"console\": dl.ConsoleLogger(),\n",
    "            \"tensorboard\": dl.TensorboardLogger(logdir=\"./logdir07/tb\"),\n",
    "        }\n",
    "    \n",
    "    @property\n",
    "    def seed(self) -> int:\n",
    "        return 73\n",
    "\n",
    "    @property\n",
    "    def name(self) -> str:\n",
    "        return \"experiment73\"\n",
    "    \n",
    "    @property\n",
    "    def stages(self) -> Iterable[str]:\n",
    "        return [\"stage_1\", \"stage_2\"]\n",
    "    \n",
    "    def get_stage_len(self, stage: str) -> int:\n",
    "        return 5\n",
    "    \n",
    "    def get_loaders(self, stage: str) -> \"OrderedDict[str, DataLoader]\":\n",
    "        if stage == \"stage_1\":\n",
    "            return {\n",
    "                \"train_1\": DataLoader(t1, batch_size=32, num_workers=1), \n",
    "                \"valid\": DataLoader(v1, batch_size=32, num_workers=1), \n",
    "            }\n",
    "        elif stage == \"stage_2\":\n",
    "            return {\n",
    "                \"train_2\": DataLoader(t2, batch_size=32, num_workers=1), \n",
    "                \"valid\": DataLoader(v1, batch_size=32, num_workers=1), \n",
    "            }\n",
    "        else:\n",
    "            raise NotImplemented()\n",
    "    \n",
    "    def get_model(self, stage: str):\n",
    "        if self.model is not None:\n",
    "            return self.model\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, stage: str):\n",
    "        return nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def get_optimizer(self, stage: str, model):\n",
    "        if stage == \"stage_1\":\n",
    "            return torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "        elif stage == \"stage_2\":\n",
    "            return torch.optim.SGD(model.parameters(), lr=0.01)\n",
    "        else:\n",
    "            raise NotImplemented()\n",
    "\n",
    "    def get_scheduler(self, stage: str, optimizer):\n",
    "        if stage == \"stage_1\":\n",
    "            return torch.optim.lr_scheduler.MultiStepLR(optimizer, [3, 8])\n",
    "        elif stage == \"stage_2\":\n",
    "            return torch.optim.lr_scheduler.MultiStepLR(optimizer, [3, 6])\n",
    "        else:\n",
    "            raise NotImplemented()\n",
    "    \n",
    "    def get_callbacks(self, stage: str):\n",
    "        if stage == \"stage_1\":\n",
    "            return {\n",
    "                \"criterion\": dl.CriterionCallback(\n",
    "                    metric_key=\"loss\", \n",
    "                    input_key=\"logits\", \n",
    "                    target_key=\"targets\"\n",
    "                ), \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=\"./logdir07/loss\",\n",
    "                    loader_key=\"valid\", metric_key=\"loss\", \n",
    "                    minimize=True, save_n_best=3\n",
    "                ),\n",
    "            }\n",
    "        elif stage == \"stage_2\":\n",
    "            return {\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",
    "                \"optimizer\": dl.OptimizerCallback(metric_key=\"loss\"), \n",
    "                \"scheduler\": dl.SchedulerCallback(\n",
    "                    loader_key=\"valid\", metric_key=\"loss\"\n",
    "                ),\n",
    "                \"checkpoint_auc\": dl.CheckpointCallback(\n",
    "                    logdir=\"./logdir07/auc\",\n",
    "                    loader_key=\"valid\", metric_key=\"auc\", \n",
    "                    minimize=False, save_n_best=3\n",
    "                ),\n",
    "            }\n",
    "        else:\n",
    "            raise NotImplemented()\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 8 - integration with hyperparameter search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could also wrap everything for hyperparameter search quickly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "import optuna    \n",
    "\n",
    "def objective(trial):\n",
    "    num_epochs = 6\n",
    "    num_hidden1 = int(trial.suggest_loguniform(\"num_hidden1\", 2, 16))\n",
    "    num_hidden2 = int(trial.suggest_loguniform(\"num_hidden2\", 2, 16))\n",
    "    logdir = f\"./logdir08/{datetime.now().strftime('%Y%m%d-%H%M%S')}\"\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",
    "    }\n",
    "\n",
    "    class CustomRunner(dl.IRunner):\n",
    "        def get_trial(self):\n",
    "            return trial\n",
    "\n",
    "        def get_engine(self) -> dl.IEngine:\n",
    "            return dl.DeviceEngine(\"cpu\")\n",
    "\n",
    "        def get_loggers(self):\n",
    "            return {\n",
    "                \"console\": dl.ConsoleLogger(),\n",
    "                \"tensorboard\": dl.TensorboardLogger(logdir=f\"{logdir}/tb\"),\n",
    "            }\n",
    "\n",
    "        @property\n",
    "        def seed(self) -> int:\n",
    "            return 73\n",
    "\n",
    "        @property\n",
    "        def name(self) -> str:\n",
    "            return \"experiment73\"\n",
    "\n",
    "        @property\n",
    "        def stages(self) -> Iterable[str]:\n",
    "            return [\"stage_1\", \"stage_2\"]\n",
    "\n",
    "        def get_stage_len(self, stage: str) -> int:\n",
    "            return num_epochs\n",
    "\n",
    "        def get_loaders(self, stage: str) -> \"OrderedDict[str, DataLoader]\":\n",
    "            return loaders\n",
    "\n",
    "        def get_model(self, stage: str):\n",
    "            return nn.Sequential(\n",
    "                nn.Linear(2, num_hidden1), nn.ReLU(), \n",
    "                nn.Linear(num_hidden1, num_hidden2), nn.ReLU(), \n",
    "                nn.Linear(num_hidden2, 1)\n",
    "            )\n",
    "\n",
    "        def get_criterion(self, stage: str):\n",
    "            return nn.BCEWithLogitsLoss()\n",
    "\n",
    "        def get_optimizer(self, stage: str, model):\n",
    "            return torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "        def get_scheduler(self, stage: str, optimizer):\n",
    "            return torch.optim.lr_scheduler.MultiStepLR(optimizer, [3, 6])\n",
    "\n",
    "        def get_callbacks(self, stage: str):\n",
    "            return {\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",
    "                \"optimizer\": dl.OptimizerCallback(metric_key=\"loss\"), \n",
    "                \"scheduler\": dl.SchedulerCallback(\n",
    "                    loader_key=\"valid\", metric_key=\"loss\"\n",
    "                ),\n",
    "                \"checkpoint\": dl.CheckpointCallback(\n",
    "                    logdir=f\"{logdir}/auc\",\n",
    "                    loader_key=\"valid\", metric_key=\"auc\", \n",
    "                    minimize=False, save_n_best=3\n",
    "                ),\n",
    "                \"optuna\": dl.OptunaPruningCallback(loader_key=\"valid\", metric_key=\"auc\", minimize=False)\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 = CustomRunner()\n",
    "    runner.run()\n",
    "    score = runner.callbacks[\"optuna\"].best_score\n",
    "    \n",
    "    return score\n",
    "\n",
    "study = optuna.create_study(\n",
    "    direction=\"maximize\",\n",
    "#     direction=\"minimize\",\n",
    "    pruner=optuna.pruners.MedianPruner(\n",
    "        n_startup_trials=0, n_warmup_steps=0, interval_steps=1\n",
    "    ),\n",
    ")\n",
    "study.optimize(objective, n_trials=5, timeout=300)\n",
    "print(study.best_value, study.best_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Act 9 - 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=./logdir09/"
   ]
  },
  {
   "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",
    "    def get_engine(self) -> dl.IEngine:\n",
    "        return dl.DeviceEngine(\"cpu\")\n",
    "    \n",
    "    def get_loggers(self):\n",
    "        return {\n",
    "            \"console\": dl.ConsoleLogger(),\n",
    "            \"csv\": dl.CSVLogger(logdir=\"./logdir09\"),\n",
    "            \"tensorboard\": dl.TensorboardLogger(logdir=\"./logdir09/tb\"),\n",
    "        }\n",
    "    \n",
    "    @property\n",
    "    def stages(self) -> Iterable[str]:\n",
    "        return [\"train\"]\n",
    "    \n",
    "    def get_stage_len(self, stage: str) -> int:\n",
    "        return num_epochs\n",
    "    \n",
    "    def get_loaders(self, stage: str) -> \"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, stage: str):\n",
    "        return torch.nn.Linear(num_features, num_classes)\n",
    "\n",
    "    def get_criterion(self, stage: str):\n",
    "        return torch.nn.CrossEntropyLoss()\n",
    "\n",
    "    def get_optimizer(self, stage: str, model):\n",
    "        return torch.optim.Adam(model.parameters())\n",
    "\n",
    "    def get_scheduler(self, stage: str, optimizer):\n",
    "        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [2])\n",
    "    \n",
    "    def get_callbacks(self, stage: str):\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",
    "            \"optimizer\": dl.OptimizerCallback(metric_key=\"loss\"), \n",
    "            \"scheduler\": dl.SchedulerCallback(\n",
    "                loader_key=\"valid\", metric_key=\"loss\"\n",
    "            ),\n",
    "            \"checkpoint1\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir09/loss\",\n",
    "                loader_key=\"valid\", metric_key=\"loss\", \n",
    "                minimize=False, save_n_best=3\n",
    "            ),\n",
    "            \"checkpoint2\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir09/auc\",\n",
    "                loader_key=\"valid\", metric_key=\"auc\", \n",
    "                minimize=True, save_n_best=1\n",
    "            ),\n",
    "            \"checkpoint3\": dl.CheckpointCallback(\n",
    "                logdir=\"./logdir9/accuracy\",\n",
    "                loader_key=\"valid\", metric_key=\"accuracy\", \n",
    "                minimize=True, save_n_best=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] *",
   "language": "python",
   "name": "conda-env-py37-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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}