dev/02c_training.ipynb

{
 "cells": [
  {
   "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": [
    "#export\n",
    "import sys\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "\n",
    "from loop import callbacks as C\n",
    "from loop.config import defaults\n",
    "from loop.metrics import accuracy\n",
    "from loop.modules import TinyNet\n",
    "from loop.phase import Phase\n",
    "from loop.utils import unwrap_if_single"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def default_optimizer(model, **params):\n",
    "    if 'lr' not in params:\n",
    "        params['lr'] = 0.001\n",
    "    return optim.Adam(model.parameters(), **params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def create_callbacks(cbs, default: bool=True):\n",
    "    defaults = [C.RollingLoss(), C.History(), C.StreamLogger()] if default else []\n",
    "    cbs = list(cbs or [])\n",
    "    cbs += defaults\n",
    "    return C.Group(cbs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "_err_stream = sys.stderr\n",
    "def report_error(exc):\n",
    "    _err_stream.write(str(exc))\n",
    "    _err_stream.flush()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class Loop:\n",
    "    def __init__(self, model: nn.Module, cbs: list=None,\n",
    "                 default_cb: bool=True, opt_fn: 'callable'=default_optimizer,\n",
    "                 opt_params: dict=None, device: 'device'=defaults.device,\n",
    "                 loss_fn: 'callable'=defaults.loss_fn):\n",
    "        \n",
    "        model.to(device)\n",
    "        opt = opt_fn(model, **(opt_params or {}))\n",
    "        cb = create_callbacks(cbs, default_cb)\n",
    "        cb.model = model\n",
    "        \n",
    "        self.model = model\n",
    "        self.opt = opt\n",
    "        self.cb = cb\n",
    "        self.loss_fn = loss_fn \n",
    "        self.device = device\n",
    "        \n",
    "    def fit_datasets(self, trn_ds, val_ds, epochs: int=1, batch_size: int=defaults.bs):\n",
    "        phases = Phase.make_train_valid(\n",
    "            trn_ds, val_ds, bs=batch_size, num_workers=defaults.n_jobs)\n",
    "        self.train(phases, epochs)\n",
    "        \n",
    "    def train(self, phases: list, epochs: int=1):\n",
    "        try:\n",
    "            self.cb.training_started(phases=phases)\n",
    "            for epoch in range(1, epochs + 1):\n",
    "                self.train_one_epoch(phases, epoch)\n",
    "            self.cb.training_ended(phases=phases)\n",
    "            \n",
    "        except TrainingInterrupted as e:\n",
    "            self.cb.training_ended(phases=phases)\n",
    "            self.cb.interrupted(exc=e)\n",
    "            \n",
    "        except Exception as e:\n",
    "            report_error(e)\n",
    "            \n",
    "        finally:\n",
    "            self.cb.cleanup()\n",
    "    \n",
    "    def train_one_epoch(self, phases: list, curr_epoch: int=1):\n",
    "        cb, model, opt = self.cb, self.model, self.opt\n",
    "        \n",
    "        cb.epoch_started(epoch=curr_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",
    "                phase.batch_index += 1\n",
    "                cb.batch_started(phase=phase, total_batches=n)\n",
    "                x, y = place_and_unwrap(batch, self.device)\n",
    "                \n",
    "                with torch.set_grad_enabled(is_training):\n",
    "                    cb.before_forward()\n",
    "                    out = model(x)\n",
    "                    cb.after_forward()\n",
    "                    loss = self.loss_fn(out, y)\n",
    "                \n",
    "                if is_training:\n",
    "                    opt.zero_grad()\n",
    "                    cb.before_backward()\n",
    "                    loss.backward()\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=curr_epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TrainingInterrupted(Exception):\n",
    "    def __init__(self, context=None):\n",
    "        self.context = context\n",
    "    def __str__(self):\n",
    "        return str(self.context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def raise_interruption(context):\n",
    "    raise TrainingInterrupted(context=context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def place_and_unwrap(batch, device):\n",
    "    \"\"\"Places tensors from batch onto proper device and converts targets\n",
    "    into proper shape depending on number of tensors.\n",
    "    \"\"\"\n",
    "    batch = [t.to(device) for t in batch]\n",
    "    x, *ys = batch\n",
    "    return x, unwrap_if_single(ys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import functional as F\n",
    "from loop.testing import get_mnist\n",
    "loop = Loop(TinyNet(1), cbs=[C.Average(accuracy)], loss_fn=F.cross_entropy)\n",
    "loop.fit_datasets(*get_mnist(), epochs=3, batch_size=2048)\n",
    "loop.cb['history'].plot();"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fastai (cuda 10)",
   "language": "python",
   "name": "fastai"
  },
  "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.1"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}