examples/notebooks/demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demo\n",
    "\n",
    "Minimal working examples with Catalyst.\n",
    "- ML - Projector, aka \"Linear regression is my profession\"\n",
    "- CV - mnist classification, autoencoder, variational autoencoder\n",
    "- GAN - mnist again :)\n",
    "- NLP - sentiment analysis\n",
    "- RecSys - movie recommendations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -U torch==1.4.0 torchvision==0.5.0 torchtext==0.5.0 catalyst==20.10.1 pandas==1.0.1 tqdm==4.43"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for tensorboard integration\n",
    "# !pip install tensorflow\n",
    "# %load_ext tensorboard\n",
    "# %tensorboard --logdir ./logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchtext\n",
    "import catalyst\n",
    "\n",
    "print(\n",
    "    \"torch\", torch.__version__, \"\\n\",\n",
    "    \"torchvision\", torchvision.__version__, \"\\n\",\n",
    "    \"torchtext\", torchtext.__version__, \"\\n\",\n",
    "    \"catalyst\", catalyst.__version__,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ML - Projector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from catalyst.dl import SupervisedRunner\n",
    "\n",
    "# experiment setup\n",
    "logdir = \"./logdir\"\n",
    "num_epochs = 8\n",
    "\n",
    "# data\n",
    "num_samples, num_features = int(1e4), int(1e1)\n",
    "X, y = torch.rand(num_samples, num_features), torch.rand(num_samples)\n",
    "dataset = TensorDataset(X, y)\n",
    "loader = DataLoader(dataset, batch_size=32, num_workers=1)\n",
    "loaders = {\"train\": loader, \"valid\": loader}\n",
    "\n",
    "# model, criterion, optimizer, scheduler\n",
    "model = torch.nn.Linear(num_features, 1)\n",
    "criterion = torch.nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [3, 6])\n",
    "\n",
    "# model training\n",
    "runner = SupervisedRunner()\n",
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    loaders=loaders,\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST - classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "from catalyst.data.cv import ToTensor\n",
    "from catalyst.contrib.datasets import MNIST\n",
    "\n",
    "model = torch.nn.Linear(28 * 28, 10)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "loaders = {\n",
    "    \"train\": DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=ToTensor()), batch_size=32),\n",
    "    \"valid\": DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=ToTensor()), batch_size=32),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst import dl, metrics\n",
    "\n",
    "\n",
    "class CustomRunner(dl.Runner):\n",
    "    def _handle_batch(self, batch):\n",
    "        x, y = batch\n",
    "        y_hat = self.model(x.view(x.size(0), -1))\n",
    "        loss = F.cross_entropy(y_hat, y)\n",
    "        accuracy01, accuracy03, accuracy05 = metrics.accuracy(y_hat, y, topk=(1, 3, 5))\n",
    "        \n",
    "        self.batch_metrics = {\n",
    "            \"loss\": loss,\n",
    "            \"accuracy01\": accuracy01,\n",
    "            \"accuracy03\": accuracy03,\n",
    "            \"accuracy05\": accuracy05,\n",
    "        }\n",
    "        \n",
    "        if self.is_train_loader:\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "            self.optimizer.zero_grad()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "runner = CustomRunner()\n",
    "runner.train(\n",
    "    model=model, \n",
    "    optimizer=optimizer, \n",
    "    loaders=loaders, \n",
    "    num_epochs=1,\n",
    "    verbose=True,\n",
    "    timeit=False,\n",
    "    logdir=\"./logs_custom\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST - classification with AutoEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "from catalyst.data.cv import ToTensor\n",
    "from catalyst.contrib.datasets import MNIST\n",
    "\n",
    "class ClassifyAE(nn.Module):\n",
    "    def __init__(self, in_features, hid_features, out_features):\n",
    "        super().__init__()\n",
    "        self.encoder = nn.Sequential(nn.Linear(in_features, hid_features), nn.Tanh())\n",
    "        self.decoder = nn.Sequential(nn.Linear(hid_features, in_features), nn.Sigmoid())\n",
    "        self.clf = nn.Linear(hid_features, out_features)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        z = self.encoder(x)\n",
    "        y_hat = self.clf(z)\n",
    "        x_ = self.decoder(z)\n",
    "        return y_hat, x_\n",
    "        \n",
    "\n",
    "model = ClassifyAE(28 * 28, 128, 10)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "loaders = {\n",
    "    \"train\": DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=ToTensor()), batch_size=32),\n",
    "    \"valid\": DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=ToTensor()), batch_size=32),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst import dl, metrics\n",
    "\n",
    "\n",
    "class CustomRunner(dl.Runner):\n",
    "    def _handle_batch(self, batch):\n",
    "        x, y = batch\n",
    "        x = x.view(x.size(0), -1)\n",
    "        y_hat, x_ = self.model(x)\n",
    "        loss_clf = F.cross_entropy(y_hat, y)\n",
    "        loss_ae = F.mse_loss(x_, x)\n",
    "        loss = loss_clf + loss_ae\n",
    "        accuracy01, accuracy03, accuracy05 = metrics.accuracy(y_hat, y, topk=(1, 3, 5))\n",
    "        \n",
    "        self.batch_metrics = {\n",
    "            \"loss_clf\": loss_clf,\n",
    "            \"loss_ae\": loss_ae,\n",
    "            \"loss\": loss,\n",
    "            \"accuracy01\": accuracy01,\n",
    "            \"accuracy03\": accuracy03,\n",
    "            \"accuracy05\": accuracy05,\n",
    "        }\n",
    "        \n",
    "        if self.is_train_loader:\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "            self.optimizer.zero_grad()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "runner = CustomRunner()\n",
    "runner.train(\n",
    "    model=model, \n",
    "    optimizer=optimizer, \n",
    "    loaders=loaders,\n",
    "    num_epochs=1,\n",
    "    verbose=True,\n",
    "    timeit=False,\n",
    "    logdir=\"./logs_custom_ae\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST - classification with Variational AutoEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "from catalyst.data.cv import ToTensor\n",
    "from catalyst.contrib.datasets import MNIST\n",
    "\n",
    "\n",
    "LOG_SCALE_MAX = 2\n",
    "LOG_SCALE_MIN = -10\n",
    "\n",
    "\n",
    "def normal_sample(mu, sigma):\n",
    "    \"\"\"\n",
    "    Sample from multivariate Gaussian distribution z ~ N(z|mu,sigma)\n",
    "    while supporting backpropagation through its mean and variance.\n",
    "    \"\"\"\n",
    "    return mu + sigma * torch.randn_like(sigma)\n",
    "\n",
    "\n",
    "def normal_logprob(mu, sigma, z):\n",
    "    \"\"\"\n",
    "    Probability density function of multivariate Gaussian distribution\n",
    "    N(z|mu,sigma).\n",
    "    \"\"\"\n",
    "    normalization_constant = (-sigma.log() - 0.5 * np.log(2 * np.pi))\n",
    "    square_term = -0.5 * ((z - mu) / sigma)**2\n",
    "    logprob_vec = normalization_constant + square_term\n",
    "    logprob = logprob_vec.sum(1)\n",
    "    return logprob\n",
    "\n",
    "\n",
    "class ClassifyVAE(torch.nn.Module):\n",
    "    def __init__(self, in_features, hid_features, out_features):\n",
    "        super().__init__()\n",
    "        self.encoder = torch.nn.Linear(in_features, hid_features * 2)\n",
    "        self.decoder = nn.Sequential(nn.Linear(hid_features, in_features), nn.Sigmoid())\n",
    "        self.clf = torch.nn.Linear(hid_features, out_features)\n",
    "    \n",
    "    def forward(self, x, deterministic=False):\n",
    "        z = self.encoder(x)\n",
    "        bs, z_dim = z.shape\n",
    "\n",
    "        loc, log_scale = z[:, :z_dim // 2], z[:, z_dim // 2:]\n",
    "        log_scale = torch.clamp(log_scale, LOG_SCALE_MIN, LOG_SCALE_MAX)\n",
    "        scale = torch.exp(log_scale)\n",
    "        z_ = loc if deterministic else normal_sample(loc, scale)\n",
    "        z_logprob = normal_logprob(loc, scale, z_)\n",
    "        z_ = z_.view(bs, -1)\n",
    "        x_ = self.decoder(z_)\n",
    "        y_hat = self.clf(z_)\n",
    "\n",
    "        return y_hat, x_, z_logprob, loc, log_scale\n",
    "        \n",
    "\n",
    "model = ClassifyVAE(28 * 28, 64, 10)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "loaders = {\n",
    "    \"train\": DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=ToTensor()), batch_size=32),\n",
    "    \"valid\": DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=ToTensor()), batch_size=32),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst import dl, metrics\n",
    "\n",
    "\n",
    "class CustomRunner(dl.Runner):\n",
    "    def _handle_batch(self, batch):\n",
    "        kld_regularization = 0.1\n",
    "        logprob_regularization = 0.01\n",
    "        \n",
    "        x, y = batch\n",
    "        x = x.view(x.size(0), -1)\n",
    "        y_hat, x_, z_logprob, loc, log_scale = self.model(x)\n",
    "        \n",
    "        loss_clf = F.cross_entropy(y_hat, y)\n",
    "        loss_ae = F.mse_loss(x_, x)\n",
    "        loss_kld = -0.5 * torch.mean(\n",
    "            1 + log_scale - loc.pow(2) - log_scale.exp()\n",
    "        ) * kld_regularization\n",
    "        loss_logprob = torch.mean(z_logprob) * logprob_regularization\n",
    "        loss = loss_clf + loss_ae + loss_kld + loss_logprob\n",
    "        accuracy01, accuracy03, accuracy05 = metrics.accuracy(y_hat, y, topk=(1, 3, 5))\n",
    "        \n",
    "        self.batch_metrics = {\n",
    "            \"loss_clf\": loss_clf,\n",
    "            \"loss_ae\": loss_ae,\n",
    "            \"loss_kld\": loss_kld,\n",
    "            \"loss_logprob\": loss_logprob,\n",
    "            \"loss\": loss,\n",
    "            \"accuracy01\": accuracy01,\n",
    "            \"accuracy03\": accuracy03,\n",
    "            \"accuracy05\": accuracy05,\n",
    "        }\n",
    "        \n",
    "        if self.is_train_loader:\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "            self.optimizer.zero_grad()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "runner = CustomRunner()\n",
    "runner.train(\n",
    "    model=model, \n",
    "    optimizer=optimizer, \n",
    "    loaders=loaders,\n",
    "    num_epochs=1,\n",
    "    verbose=True,\n",
    "    timeit=False,\n",
    "    logdir=\"./logs_custom_vae\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST - segmentation with classification auxiliary task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "from catalyst.data.cv import ToTensor\n",
    "from catalyst.contrib.datasets import MNIST\n",
    "\n",
    "class ClassifyUnet(nn.Module):\n",
    "    def __init__(self, in_channels, in_hw, out_features):\n",
    "        super().__init__()\n",
    "        self.encoder = nn.Sequential(nn.Conv2d(in_channels, in_channels, 3, 1, 1), nn.Tanh())\n",
    "        self.decoder = nn.Conv2d(in_channels, in_channels, 3, 1, 1)\n",
    "        self.clf = nn.Linear(in_channels * in_hw * in_hw, out_features)\n",
    "\n",
    "    def forward(self, x):\n",
    "        z = self.encoder(x)\n",
    "        z_ = z.view(z.size(0), -1)\n",
    "        y_hat = self.clf(z_)\n",
    "        x_ = self.decoder(z)\n",
    "        return y_hat, x_\n",
    "        \n",
    "\n",
    "model = ClassifyUnet(1, 28, 10)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.02)\n",
    "\n",
    "loaders = {\n",
    "    \"train\": DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=ToTensor()), batch_size=32),\n",
    "    \"valid\": DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=ToTensor()), batch_size=32),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from catalyst import dl, metrics\n",
    "\n",
    "\n",
    "class CustomRunner(dl.Runner):\n",
    "    def _handle_batch(self, batch):\n",
    "        x, y = batch\n",
    "        x_noise = (x + torch.rand_like(x)).clamp_(0, 1)\n",
    "        y_hat, x_ = self.model(x_noise)\n",
    "\n",
    "        loss_clf = F.cross_entropy(y_hat, y)\n",
    "        iou = metrics.iou(x_, x)\n",
    "        loss_iou = 1 - iou\n",
    "        loss = loss_clf + loss_iou\n",
    "        accuracy01, accuracy03, accuracy05 = metrics.accuracy(y_hat, y, topk=(1, 3, 5))\n",
    "        \n",
    "        self.batch_metrics = {\n",
    "            \"loss_clf\": loss_clf,\n",
    "            \"loss_iou\": loss_iou,\n",
    "            \"loss\": loss,\n",
    "            \"iou\": iou,\n",
    "            \"accuracy01\": accuracy01,\n",
    "            \"accuracy03\": accuracy03,\n",
    "            \"accuracy05\": accuracy05,\n",
    "        }\n",
    "        \n",
    "        if self.is_train_loader:\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "            self.optimizer.zero_grad()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "runner = CustomRunner()\n",
    "runner.train(\n",
    "    model=model, \n",
    "    optimizer=optimizer, \n",
    "    loaders=loaders,\n",
    "    num_epochs=1,\n",
    "    verbose=True,\n",
    "    timeit=False,\n",
    "    logdir=\"./logs_custom_unet\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from catalyst.contrib.nn import GlobalMaxPool2d, Flatten, Lambda\n",
    "\n",
    "# Create the discriminator\n",
    "discriminator = nn.Sequential(\n",
    "    nn.Conv2d(1, 64, (3, 3), stride=(2, 2), padding=1),\n",
    "    nn.LeakyReLU(0.2, inplace=True),\n",
    "    nn.Conv2d(64, 128, (3, 3), stride=(2, 2), padding=1),\n",
    "    nn.LeakyReLU(0.2, inplace=True),\n",
    "    GlobalMaxPool2d(),\n",
    "    Flatten(),\n",
    "    nn.Linear(128, 1)\n",
    ")\n",
    "\n",
    "# Create the generator\n",
    "latent_dim = 128\n",
    "generator = nn.Sequential(\n",
    "    # We want to generate 128 coefficients to reshape into a 7x7x128 map\n",
    "    nn.Linear(128, 128 * 7 * 7),\n",
    "    nn.LeakyReLU(0.2, inplace=True),\n",
    "    Lambda(lambda x: x.view(x.size(0), 128, 7, 7)),\n",
    "    nn.ConvTranspose2d(128, 128, (4, 4), stride=(2, 2), padding=1),\n",
    "    nn.LeakyReLU(0.2, inplace=True),\n",
    "    nn.ConvTranspose2d(128, 128, (4, 4), stride=(2, 2), padding=1),\n",
    "    nn.LeakyReLU(0.2, inplace=True),\n",
    "    nn.Conv2d(128, 1, (7, 7), padding=3),\n",
    "    nn.Sigmoid(),\n",
    ")\n",
    "\n",
    "# Final model\n",
    "model = {\n",
    "    \"generator\": generator,\n",
    "    \"discriminator\": discriminator,\n",
    "}\n",
    "\n",
    "optimizer = {\n",
    "    \"generator\": torch.optim.Adam(generator.parameters(), lr=0.0003, betas=(0.5, 0.999)),\n",
    "    \"discriminator\": torch.optim.Adam(discriminator.parameters(), lr=0.0003, betas=(0.5, 0.999)),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst import dl\n",
    "\n",
    "class CustomRunner(dl.Runner):\n",
    "\n",
    "    def _handle_batch(self, batch):\n",
    "        real_images, _ = batch\n",
    "        batch_metrics = {}\n",
    "        \n",
    "        # Sample random points in the latent space\n",
    "        batch_size = real_images.shape[0]\n",
    "        random_latent_vectors = torch.randn(batch_size, latent_dim).to(self.device)\n",
    "        \n",
    "        # Decode them to fake images\n",
    "        generated_images = self.model[\"generator\"](random_latent_vectors).detach()\n",
    "        # Combine them with real images\n",
    "        combined_images = torch.cat([generated_images, real_images])\n",
    "        \n",
    "        # Assemble labels discriminating real from fake images\n",
    "        labels = torch.cat([\n",
    "            torch.ones((batch_size, 1)), torch.zeros((batch_size, 1))\n",
    "        ]).to(self.device)\n",
    "        # Add random noise to the labels - important trick!\n",
    "        labels += 0.05 * torch.rand(labels.shape).to(self.device)\n",
    "        \n",
    "        # Train the discriminator\n",
    "        predictions = self.model[\"discriminator\"](combined_images)\n",
    "        batch_metrics[\"loss_discriminator\"] = \\\n",
    "          F.binary_cross_entropy_with_logits(predictions, labels)\n",
    "        \n",
    "        # Sample random points in the latent space\n",
    "        random_latent_vectors = torch.randn(batch_size, latent_dim).to(self.device)\n",
    "        # Assemble labels that say \"all real images\"\n",
    "        misleading_labels = torch.zeros((batch_size, 1)).to(self.device)\n",
    "        \n",
    "        # Train the generator\n",
    "        generated_images = self.model[\"generator\"](random_latent_vectors)\n",
    "        predictions = self.model[\"discriminator\"](generated_images)\n",
    "        batch_metrics[\"loss_generator\"] = \\\n",
    "          F.binary_cross_entropy_with_logits(predictions, misleading_labels)\n",
    "        \n",
    "        self.batch_metrics.update(**batch_metrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from torch.utils.data import DataLoader\n",
    "from catalyst.data.cv import ToTensor\n",
    "from catalyst.contrib.datasets import MNIST\n",
    "\n",
    "\n",
    "loaders = {\n",
    "    \"train\": DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=ToTensor()), batch_size=32),\n",
    "}\n",
    "\n",
    "runner = CustomRunner()\n",
    "runner.train(\n",
    "    model=model, \n",
    "    optimizer=optimizer,\n",
    "    loaders=loaders,\n",
    "    callbacks=[\n",
    "        dl.OptimizerCallback(\n",
    "            optimizer_key=\"generator\", \n",
    "            metric_key=\"loss_generator\"\n",
    "        ),\n",
    "        dl.OptimizerCallback(\n",
    "            optimizer_key=\"discriminator\", \n",
    "            metric_key=\"loss_discriminator\"\n",
    "        ),\n",
    "    ],\n",
    "    main_metric=\"loss_generator\",\n",
    "    num_epochs=20,\n",
    "    verbose=True,\n",
    "    logdir=\"./logs_gan\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F \n",
    "\n",
    "import torchtext\n",
    "from torchtext.datasets import text_classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NGRAMS = 2\n",
    "import os\n",
    "if not os.path.isdir('./data'):\n",
    "    os.mkdir('./data')\n",
    "if not os.path.isdir('./data/nlp'):\n",
    "    os.mkdir('./data/nlp')\n",
    "train_dataset, valid_dataset = text_classification.DATASETS['AG_NEWS'](\n",
    "    root='./data/nlp', ngrams=NGRAMS, vocab=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "VOCAB_SIZE = len(train_dataset.get_vocab())\n",
    "EMBED_DIM = 32\n",
    "NUM_CLASS = len(train_dataset.get_labels())\n",
    "BATCH_SIZE = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_batch(batch):\n",
    "    label = torch.tensor([entry[0] for entry in batch])\n",
    "    text = [entry[1] for entry in batch]\n",
    "    offsets = [0] + [len(entry) for entry in text]\n",
    "    # torch.Tensor.cumsum returns the cumulative sum\n",
    "    # of elements in the dimension dim.\n",
    "    # torch.Tensor([1.0, 2.0, 3.0]).cumsum(dim=0)\n",
    "\n",
    "    offsets = torch.tensor(offsets[:-1]).cumsum(dim=0)\n",
    "    text = torch.cat(text)\n",
    "    output = {\n",
    "        \"text\": text,\n",
    "        \"offsets\": offsets,\n",
    "        \"label\": label\n",
    "    }\n",
    "    return output\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    train_dataset, \n",
    "    batch_size=BATCH_SIZE, \n",
    "    shuffle=True,\n",
    "    collate_fn=generate_batch,\n",
    ")\n",
    "\n",
    "valid_loader = torch.utils.data.DataLoader(\n",
    "    valid_dataset, \n",
    "    batch_size=BATCH_SIZE, \n",
    "    shuffle=False,\n",
    "    collate_fn=generate_batch,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TextSentiment(nn.Module):\n",
    "    def __init__(self, vocab_size, embed_dim, num_class):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.EmbeddingBag(vocab_size, embed_dim, sparse=True)\n",
    "        self.fc = nn.Linear(embed_dim, num_class)\n",
    "        self.init_weights()\n",
    "\n",
    "    def init_weights(self):\n",
    "        initrange = 0.5\n",
    "        self.embedding.weight.data.uniform_(-initrange, initrange)\n",
    "        self.fc.weight.data.uniform_(-initrange, initrange)\n",
    "        self.fc.bias.data.zero_()\n",
    "\n",
    "    def forward(self, text, offsets):\n",
    "        embedded = self.embedding(text, offsets)\n",
    "        return self.fc(embedded)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = TextSentiment(VOCAB_SIZE, EMBED_DIM, NUM_CLASS)\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=4.0)\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1, gamma=0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst.dl import SupervisedRunner, \\\n",
    "    CriterionCallback, AccuracyCallback\n",
    "\n",
    "# input_keys - which key from dataloader we need to pass to the model\n",
    "runner = SupervisedRunner(input_key=[\"text\", \"offsets\"])\n",
    "\n",
    "runner.train(\n",
    "    model=model, \n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer, \n",
    "    scheduler=scheduler,\n",
    "    loaders={'train': train_loader, 'valid': valid_loader},\n",
    "    logdir=\"./logs_nlp\",\n",
    "    num_epochs=3,\n",
    "    verbose=True,\n",
    "    # input_key - which key from dataloader we need to pass to criterion as target label\n",
    "    callbacks=[\n",
    "        CriterionCallback(input_key=\"label\"),\n",
    "        AccuracyCallback(input_key=\"label\")\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RecSys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import os\n",
    "import requests\n",
    "import tqdm\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scipy.sparse as sp\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F \n",
    "import torch.utils.data as td\n",
    "import torch.optim as to\n",
    "\n",
    "import matplotlib.pyplot as pl\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration\n",
    "\n",
    "# The directory to store the data\n",
    "data_dir = \"data/recsys\"\n",
    "\n",
    "train_rating = \"ml-1m.train.rating\"\n",
    "test_negative = \"ml-1m.test.negative\"\n",
    "\n",
    "# NCF config\n",
    "train_negative_samples = 4\n",
    "test_negative_samples = 99\n",
    "embedding_dim = 64\n",
    "hidden_dim = 32\n",
    "\n",
    "# Training config\n",
    "batch_size = 256\n",
    "epochs = 10  # Original implementation uses 20\n",
    "top_k=10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.isdir('./data'):\n",
    "    os.mkdir('./data')\n",
    "if not os.path.isdir('./data/recsys'):\n",
    "    os.mkdir('./data/recsys')\n",
    "    \n",
    "for file_name in [train_rating, test_negative]:\n",
    "    file_path = os.path.join(data_dir, file_name)\n",
    "    if os.path.exists(file_path):\n",
    "        print(\"Skip loading \" + file_name)\n",
    "        continue\n",
    "    with open(file_path, \"wb\") as tf:\n",
    "        print(\"Load \" + file_name)\n",
    "        r = requests.get(\"https://raw.githubusercontent.com/hexiangnan/neural_collaborative_filtering/master/Data/\" + file_name, allow_redirects=True)\n",
    "        tf.write(r.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_train():\n",
    "    train_data = pd.read_csv(os.path.join(data_dir, train_rating), sep='\\t', header=None, names=['user', 'item'], usecols=[0, 1], dtype={0: np.int32, 1: np.int32})\n",
    "\n",
    "    user_num = train_data['user'].max() + 1\n",
    "    item_num = train_data['item'].max() + 1\n",
    "\n",
    "    train_data = train_data.values.tolist()\n",
    "\n",
    "    # Convert ratings as a dok matrix\n",
    "    train_mat = sp.dok_matrix((user_num, item_num), dtype=np.float32)\n",
    "    for user, item in train_data:\n",
    "        train_mat[user, item] = 1.0\n",
    "        \n",
    "    return train_data, train_mat, user_num, item_num\n",
    "\n",
    "\n",
    "train_data, train_mat, user_num, item_num = preprocess_train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_test():\n",
    "    test_data = []\n",
    "    with open(os.path.join(data_dir, test_negative)) as tnf:\n",
    "        for line in tnf:\n",
    "            parts = line.split('\\t')\n",
    "            assert len(parts) == test_negative_samples + 1\n",
    "            \n",
    "            user, positive = eval(parts[0])\n",
    "            test_data.append([user, positive])\n",
    "            \n",
    "            for negative in parts[1:]:\n",
    "                test_data.append([user, int(negative)])\n",
    "\n",
    "    return test_data\n",
    "\n",
    "\n",
    "valid_data = preprocess_test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NCFDataset(td.Dataset):\n",
    "    \n",
    "    def __init__(self, positive_data, item_num, positive_mat, negative_samples=0):\n",
    "        super(NCFDataset, self).__init__()\n",
    "        self.positive_data = positive_data\n",
    "        self.item_num = item_num\n",
    "        self.positive_mat = positive_mat\n",
    "        self.negative_samples = negative_samples\n",
    "        \n",
    "        self.reset()\n",
    "        \n",
    "    def reset(self):\n",
    "        print(\"Resetting dataset\")\n",
    "        if self.negative_samples > 0:\n",
    "            negative_data = self.sample_negatives()\n",
    "            data = self.positive_data + negative_data\n",
    "            labels = [1] * len(self.positive_data) + [0] * len(negative_data)\n",
    "        else:\n",
    "            data = self.positive_data\n",
    "            labels = [0] * len(self.positive_data)\n",
    "            \n",
    "        self.data = np.concatenate([\n",
    "            np.array(data), \n",
    "            np.array(labels)[:, np.newaxis]], \n",
    "            axis=1\n",
    "        )\n",
    "        \n",
    "\n",
    "    def sample_negatives(self):\n",
    "        negative_data = []\n",
    "        for user, positive in self.positive_data:\n",
    "            for _ in range(self.negative_samples):\n",
    "                negative = np.random.randint(self.item_num)\n",
    "                while (user, negative) in self.positive_mat:\n",
    "                    negative = np.random.randint(self.item_num)\n",
    "                    \n",
    "                negative_data.append([user, negative])\n",
    "\n",
    "        return negative_data\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        user, item, label = self.data[idx]\n",
    "        output = {\n",
    "            \"user\": user,\n",
    "            \"item\": item,\n",
    "            \"label\": np.float32(label),\n",
    "        }\n",
    "        return output\n",
    "\n",
    "    \n",
    "class SamplerWithReset(td.RandomSampler):\n",
    "    def __iter__(self):\n",
    "        self.data_source.reset()\n",
    "        return super().__iter__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = NCFDataset(\n",
    "    train_data, \n",
    "    item_num, \n",
    "    train_mat, \n",
    "    train_negative_samples\n",
    ")\n",
    "train_loader = td.DataLoader(\n",
    "    train_dataset, \n",
    "    batch_size=batch_size, \n",
    "    shuffle=False, \n",
    "    num_workers=4,\n",
    "    sampler=SamplerWithReset(train_dataset)\n",
    ")\n",
    "\n",
    "valid_dataset = NCFDataset(valid_data, item_num, train_mat)\n",
    "valid_loader = td.DataLoader(\n",
    "    valid_dataset, \n",
    "    batch_size=test_negative_samples+1, \n",
    "    shuffle=False, \n",
    "    num_workers=0\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Ncf(nn.Module):\n",
    "    \n",
    "    def __init__(self, user_num, item_num, embedding_dim, hidden_dim):\n",
    "        super(Ncf, self).__init__()\n",
    "        \n",
    "        self.user_embeddings = nn.Embedding(user_num, embedding_dim)\n",
    "        self.item_embeddings = nn.Embedding(item_num, embedding_dim)\n",
    "\n",
    "        self.layers = nn.Sequential(\n",
    "            nn.Linear(2 * embedding_dim, hidden_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_dim, hidden_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_dim, 1)\n",
    "        )\n",
    "\n",
    "        self.initialize()\n",
    "\n",
    "    def initialize(self):\n",
    "        nn.init.normal_(self.user_embeddings.weight, std=0.01)\n",
    "        nn.init.normal_(self.item_embeddings.weight, std=0.01)\n",
    "\n",
    "        for layer in self.layers:\n",
    "            if isinstance(layer, nn.Linear):\n",
    "                nn.init.xavier_uniform_(layer.weight)\n",
    "                layer.bias.data.zero_()\n",
    "            \n",
    "    def forward(self, user, item):\n",
    "        user_embedding = self.user_embeddings(user)\n",
    "        item_embedding = self.item_embeddings(item)\n",
    "        concat = torch.cat((user_embedding, item_embedding), -1)\n",
    "        return self.layers(concat).view(-1)\n",
    "    \n",
    "    def name(self):\n",
    "        return \"Ncf\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hit_metric(recommended, actual):\n",
    "    return int(actual in recommended)\n",
    "\n",
    "\n",
    "def dcg_metric(recommended, actual):\n",
    "    if actual in recommended:\n",
    "        index = recommended.index(actual)\n",
    "        return np.reciprocal(np.log2(index + 2))\n",
    "    return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Ncf(user_num, item_num, embedding_dim, hidden_dim)\n",
    "criterion = nn.BCEWithLogitsLoss()\n",
    "optimizer = to.Adam(model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst.dl import Callback, CallbackOrder, IRunner\n",
    "\n",
    "class NdcgLoaderMetricCallback(Callback):\n",
    "    def __init__(self):\n",
    "        super().__init__(CallbackOrder.Metric)\n",
    "\n",
    "    def on_batch_end(self, runner: IRunner):\n",
    "        item = runner.input[\"item\"]\n",
    "        predictions = runner.output[\"logits\"]\n",
    "\n",
    "        _, indices = torch.topk(predictions, top_k)\n",
    "        recommended = torch.take(item, indices).cpu().numpy().tolist()\n",
    "\n",
    "        item = item[0].item()\n",
    "        runner.batch_metrics[\"hits\"] = hit_metric(recommended, item)\n",
    "        runner.batch_metrics[\"dcgs\"] = dcg_metric(recommended, item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from catalyst.dl import SupervisedRunner, CriterionCallback\n",
    "\n",
    "# input_keys - which key from dataloader we need to pass to the model\n",
    "runner = SupervisedRunner(input_key=[\"user\", \"item\"])\n",
    "\n",
    "runner.train(\n",
    "    model=model, \n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer, \n",
    "    loaders={'train': train_loader, 'valid': valid_loader},\n",
    "    logdir=\"./logs_recsys\",\n",
    "    num_epochs=3,\n",
    "    verbose=True,\n",
    "    # input_key - which key from dataloader we need to pass to criterion as target label\n",
    "    callbacks=[\n",
    "        CriterionCallback(input_key=\"label\"),\n",
    "        NdcgLoaderMetricCallback()\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}