Pytorch/Simple Pytorch Training Loop Implementation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pytorch Training Loop Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# From \"Deep Learning Training Loop Implementation\" \n",
    "# https://github.com/devforfu/pytorch_playground/blob/master/loop.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict, OrderedDict\n",
    "import math\n",
    "from pathlib import Path\n",
    "import re\n",
    "import sys\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\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\n",
    "from torchvision import models\n",
    "from torchvision import transforms as T\n",
    "from torchvision.datasets import MNIST, CIFAR10\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "\n",
    "default_device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_model(params)\n",
    "phases = create_train_valid_data()\n",
    "opt = optim.SGD(model.params, lr=1e-3)\n",
    "\n",
    "model.to(device)\n",
    "    \n",
    "for epoch in range(1, epochs + 1):\n",
    "\n",
    "    for phase in phases:\n",
    "        n = len(phase.loader)\n",
    "        is_training = phase.grad\n",
    "        model.train(is_training)\n",
    "\n",
    "        for batch in phase.loader:\n",
    "            x, y = place_and_unwrap(batch, device)\n",
    "\n",
    "            with torch.set_grad_enabled(is_training):\n",
    "                out = model(x)\n",
    "                loss = loss_fn(out, y)\n",
    "\n",
    "            if is_training:\n",
    "                opt.zero_grad()\n",
    "                loss.backward()\n",
    "                opt.step()\n",
    "\n",
    "            phase.batch_loss = loss.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We don’t try to encapsulate all possible features into a single class or function, but delegate calls to subordinate modules. \n",
    "# Each module is responsible for reacting onto received notification properly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, opt, phases, callbacks=None, epochs=1, device=default_device, loss_fn=F.nll_loss):\n",
    "    model.to(device)\n",
    "    \n",
    "    cb = callbacks\n",
    "    \n",
    "    cb.training_started(phases=phases, optimizer=opt)\n",
    "    \n",
    "    for epoch in range(1, epochs + 1):\n",
    "        cb.epoch_started(epoch=epoch)\n",
    "\n",
    "        for phase in phases:\n",
    "            n = len(phase.loader)\n",
    "            cb.phase_started(phase=phase, total_batches=n)\n",
    "            is_training = phase.grad\n",
    "            model.train(is_training)\n",
    "\n",
    "            for batch in phase.loader:\n",
    "\n",
    "                phase.batch_index += 1\n",
    "                cb.batch_started(phase=phase, total_batches=n)\n",
    "                x, y = place_and_unwrap(batch, device)\n",
    "\n",
    "                with torch.set_grad_enabled(is_training):\n",
    "                    cb.before_forward_pass()\n",
    "                    out = model(x)\n",
    "                    cb.after_forward_pass()\n",
    "                    loss = loss_fn(out, y)\n",
    "\n",
    "                if is_training:\n",
    "                    opt.zero_grad()\n",
    "                    cb.before_backward_pass()\n",
    "                    loss.backward()\n",
    "                    cb.after_backward_pass()\n",
    "                    opt.step()\n",
    "\n",
    "                phase.batch_loss = loss.item()\n",
    "                cb.batch_ended(phase=phase, output=out, target=y)\n",
    "\n",
    "            cb.phase_ended(phase=phase)\n",
    "\n",
    "        cb.epoch_ended(phases=phases, epoch=epoch)\n",
    "\n",
    "    cb.training_ended(phases=phases)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Callbacks Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Loss\n",
    "# At the end of every batch, we’re computing a running loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RollingLoss(Callback):\n",
    "\n",
    "    def __init__(self, smooth=0.98):\n",
    "        self.smooth = smooth\n",
    "\n",
    "    def batch_ended(self, phase, **kwargs):\n",
    "        prev = phase.rolling_loss\n",
    "        a = self.smooth\n",
    "        avg_loss = a * prev + (1 - a) * phase.batch_loss\n",
    "        debias_loss = avg_loss / (1 - a ** phase.batch_index)\n",
    "        phase.rolling_loss = avg_loss\n",
    "        phase.update(debias_loss)\n",
    "\n",
    "    def epoch_ended(self, phases, **kwargs):\n",
    "        for phase in phases:\n",
    "            phase.update_metric('loss', phase.last_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2. Accuracy\n",
    "# Note that the callback receives notifications at the end of each batch, and the end of training epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def accuracy(out, y_true):\n",
    "    y_hat = out.argmax(dim=-1).view(y_true.size(0), -1)\n",
    "    y_true = y_true.view(y_true.size(0), -1)\n",
    "    match = y_hat == y_true\n",
    "    return match.float().mean()\n",
    "  \n",
    "\n",
    "class Accuracy(Callback):\n",
    "\n",
    "    def epoch_started(self, **kwargs):\n",
    "        self.values = defaultdict(int)\n",
    "        self.counts = defaultdict(int)\n",
    "\n",
    "    def batch_ended(self, phase, output, target, **kwargs):\n",
    "        acc = accuracy(output, target).detach().item()\n",
    "        self.counts[phase.name] += target.size(0)\n",
    "        self.values[phase.name] += target.size(0) * acc\n",
    "\n",
    "    def epoch_ended(self, phases, **kwargs):\n",
    "        for phase in phases:\n",
    "            metric = self.values[phase.name] / self.counts[phase.name]\n",
    "            phase.update_metric('accuracy', metric)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3. Parameter Scheduler\n",
    "# The idea is to use cyclic schedulers that adjust model’s optimizer parameters magnitudes during single or several training epochs. \n",
    "# Moreover, these schedulers not only decrease learning rates as a number of processed batches grows but also increase them for some number of steps or periodically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#  we are effectively getting a stochastic gradient with warm restarts that allows us to escape from local minima. \n",
    "# The following snippet shows how one can implement a cosine annealing learning rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CosineAnnealingSchedule:\n",
    "    \"\"\"\n",
    "    The schedule class that returns a learning rate multiplier from range [0.0, 1.0]\n",
    "    \"\"\"\n",
    "    def __init__(self, eta_min=0.0, eta_max=1.0, t_max=100, t_mult=2):\n",
    "        self.eta_min = eta_min\n",
    "        self.eta_max = eta_max\n",
    "        self.t_max = t_max\n",
    "        self.t_mult = t_mult\n",
    "        self.iter = 0\n",
    "\n",
    "    def update(self, **kwargs):\n",
    "        self.iter += 1\n",
    "\n",
    "        eta_min, eta_max, t_max = self.eta_min, self.eta_max, self.t_max\n",
    "\n",
    "        t = self.iter % t_max\n",
    "        eta = eta_min + 0.5 * (eta_max - eta_min) * (1 + math.cos(math.pi * t / t_max))\n",
    "        if t == 0:\n",
    "            self.iter = 0\n",
    "            self.t_max *= self.t_mult\n",
    "\n",
    "        return eta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stream Logger\n",
    "# The last thing we would like to add is some logging to see how well our model performs during the training process. \n",
    "# The most simplistic approach is to print stats into the standard output stream. \n",
    "# However, you could save it into CSV file or even send as a notification to your mobile phone instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_dicts(ds):\n",
    "    merged = OrderedDict()\n",
    "    for d in ds:\n",
    "        for k, v in d.items():\n",
    "            merged[k] = v\n",
    "    return merged"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StreamLogger(Callback):\n",
    "\n",
    "    def __init__(self, streams=None, log_every=1):\n",
    "        self.streams = streams or [sys.stdout]\n",
    "        self.log_every = log_every\n",
    "\n",
    "    def epoch_ended(self, phases, epoch, **kwargs):\n",
    "        if epoch % self.log_every != 0:\n",
    "            return\n",
    "        metrics = merge_dicts([phase.last_metrics for phase in phases])\n",
    "        values = [f'{k}={v:.4f}' for k, v in metrics.items()]\n",
    "        values_string = ', '.join(values)\n",
    "        string = f'Epoch: {epoch:4d} | {values_string}\\n'\n",
    "        for stream in self.streams:\n",
    "            stream.write(string)\n",
    "            stream.flush()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ProgressBar(Callback):\n",
    "\n",
    "    def training_started(self, phases, **kwargs):\n",
    "        bars = OrderedDict()\n",
    "        for phase in phases:\n",
    "            bars[phase.name] = tqdm(total=len(phase.loader), desc=phase.name)\n",
    "        self.bars = bars\n",
    "\n",
    "    def batch_ended(self, phase, **kwargs):\n",
    "        bar = self.bars[phase.name]\n",
    "        bar.set_postfix_str(f'loss: {phase.last_loss:.4f}')\n",
    "        bar.update(1)\n",
    "        bar.refresh()\n",
    "\n",
    "    def epoch_ended(self, **kwargs):\n",
    "        for bar in self.bars.values():\n",
    "            bar.n = 0\n",
    "            bar.refresh()\n",
    "\n",
    "    def training_ended(self, **kwargs):\n",
    "        for bar in self.bars.values():\n",
    "            bar.n = bar.total\n",
    "            bar.refresh()\n",
    "            bar.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we’re ready to start using our training loop!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## You definitely should user transfer learnin when working on your daily tasks. \n",
    "## It makes your network to converge much faster compared to the training from scratch. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def as_sequential(model: nn.Module):\n",
    "    \"\"\"Converts model with nested submodules into Sequential model.\"\"\"\n",
    "\n",
    "    return nn.Sequential(*list(model.children()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AdaptiveConcatPool2d(nn.Module):\n",
    "    \"\"\"Applies average and maximal adaptive pooling to the tensor and\n",
    "    concatenates results into a single tensor.\n",
    "\n",
    "    The idea is taken from fastai library.\n",
    "    \"\"\"\n",
    "    def __init__(self, size=1):\n",
    "        super().__init__()\n",
    "        self.avg = nn.AdaptiveAvgPool2d(size)\n",
    "        self.max = nn.AdaptiveMaxPool2d(size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.cat([self.max(x), self.avg(x)], 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Flatten(nn.Module):\n",
    "    \"\"\"Converts N-dimensional tensor into 'flat' one.\"\"\"\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x.view(x.size(0), -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResNetClassifier(nn.Module):\n",
    "\n",
    "    def __init__(self, n_classes, arch=models.resnet18):\n",
    "        super().__init__()\n",
    "\n",
    "        model = arch(True)\n",
    "        seq_model = as_sequential(model)\n",
    "        backbone, classifier = seq_model[:-2], seq_model[-2:]\n",
    "\n",
    "        self.backbone = backbone\n",
    "        self.top = nn.Sequential(\n",
    "            AdaptiveConcatPool2d(),\n",
    "            Flatten(),\n",
    "            \n",
    "            nn.Linear(1024, 512),\n",
    "            nn.Dropout(0.25),\n",
    "            nn.LeakyReLU(inplace=True),\n",
    "            nn.BatchNorm1d(512),\n",
    "            \n",
    "            nn.Linear(512, 256),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.LeakyReLU(inplace=True),\n",
    "            nn.BatchNorm1d(256),\n",
    "            \n",
    "            nn.Linear(256, n_classes)\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.top(self.backbone(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def classifier_weights(m: nn.Module, bn=(1, 1e-3)):\n",
    "    \"\"\"Initializes layers weights for a classification model.\"\"\"\n",
    "    \n",
    "    name = classname(m)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        if name.find('Conv') != -1:\n",
    "            nn.init.kaiming_normal_(m.weight, mode='fan_out')\n",
    "            if hasattr(m, 'bias') and m.bias is not None:\n",
    "                nn.init.zeros_(m.bias)\n",
    "\n",
    "        elif name.find('BatchNorm') != -1:\n",
    "            weight, bias = bn\n",
    "            nn.init.constant_(m.weight, weight)\n",
    "            nn.init.constant_(m.bias, bias)\n",
    "\n",
    "        elif name == 'Linear':\n",
    "            nn.init.kaiming_normal_(m.weight)\n",
    "            nn.init.zeros_(m.bias)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def freeze(m, freeze=True, bn=False):\n",
    "    for child in m.children():\n",
    "        name = classname(child)\n",
    "        if not bn and name.find('BatchNorm') != -1:\n",
    "            continue\n",
    "        for p in child.parameters():\n",
    "            p.requires_grad = not freeze"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = Path.home()/'data'/'cifar10'\n",
    "\n",
    "imagenet_stats = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "\n",
    "epochs = 10\n",
    "\n",
    "train_ds = CIFAR10(\n",
    "    data_path, \n",
    "    train=True, \n",
    "    download=True,\n",
    "    transform=T.Compose([\n",
    "        T.Resize(224),\n",
    "        T.RandomAffine(5, translate=(0.05, 0.05), scale=(0.9, 1.1)),\n",
    "        T.RandomHorizontalFlip(),\n",
    "        T.ToTensor(),\n",
    "        T.Normalize(*imagenet_stats)\n",
    "    ])\n",
    ")\n",
    "\n",
    "valid_ds = CIFAR10(\n",
    "    data_path, \n",
    "    train=False, \n",
    "    transform=T.Compose([\n",
    "        T.Resize(224),\n",
    "        T.ToTensor(),\n",
    "        T.Normalize(*imagenet_stats)\n",
    "    ])\n",
    ")\n",
    "\n",
    "phases = make_phases(train_ds, valid_ds, bs=200, n_jobs=8)\n",
    "\n",
    "model = ResNetClassifier(10)\n",
    "model.top.apply(classifier_weights)\n",
    "freeze(model.backbone)\n",
    "\n",
    "# Try AdamW later! for a better (correct) weight decay regularization\n",
    "opt = optim.SGD(model.parameters(), lr=1e-2, momentum=0.95, nesterov=True, weight_decay=1e-2)\n",
    "\n",
    "cb = CallbacksGroup([\n",
    "    RollingLoss(),\n",
    "    Accuracy(),\n",
    "    Scheduler(\n",
    "        OneCycleSchedule(t=len(phases[0].loader) * epochs),\n",
    "        params_conf=[\n",
    "            {'name': 'lr'},\n",
    "            {'name': 'weight_decay', 'inverse': True}\n",
    "        ],\n",
    "        mode='batch'\n",
    "    ),\n",
    "    StreamLogger(),\n",
    "    ProgressBar()\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train(model, opt, phases, cb, epochs=epochs, loss_fn=F.cross_entropy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr_history = pd.DataFrame(cb['scheduler'].parameter_history('lr'))\n",
    "ax = lr_history.plot(figsize=(8, 6))\n",
    "ax.set_xlabel('Training Batch Index')\n",
    "ax.set_ylabel('Learning Rate');"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}