notebooks/train/Checkpoints.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fe48dc5a-ef97-461a-82d5-69876781be0c",
   "metadata": {},
   "source": [
    "### Goal\n",
    "This code example shows you how to save intermediate model checkpoints and use these during training. Why is this useful? Say the training loss suddenly spikes during training. You may be tempted to restart training with a lower learning rate. However, it is more efficient to intermittently save the best-performing model during training - called a checkpoint. Now whenever the training loss suddenly increases, you can load the model checkpoint, decrease the learning rate, and continue training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8e4f0ab3-b30f-4892-854b-34647b843f4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "import devtorch\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bb99e371-f95e-42e1-87d4-70809708a192",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ANNClassifier(devtorch.DevModel):\n",
    "    \n",
    "    def __init__(self, n_in, n_hidden, n_out):\n",
    "        super().__init__()\n",
    "        self._n_in = n_in\n",
    "        self._n_hidden = n_hidden\n",
    "        self._n_out = n_out\n",
    "        self.layer1 = nn.Linear(n_in, n_hidden, bias=False)\n",
    "        self.layer2 = nn.Linear(n_hidden, n_out, bias=False)\n",
    "        self.init_weight(self.layer1.weight, \"glorot_uniform\")\n",
    "        self.init_weight(self.layer2.weight, \"glorot_uniform\")\n",
    "    \n",
    "    @property\n",
    "    def hyperparams(self):\n",
    "        return {**super().hyperparams, \"params\": {\"n_in\": self._n_in, \"n_hidden\": self._n_hidden, \"n_out\": self._n_out}}\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = F.leaky_relu(self.layer1(x.flatten(1, 3)))\n",
    "        return F.leaky_relu(self.layer2(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7e03722a-b295-45f9-84b6-23016016df7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CheckpointTrainer(devtorch.Trainer):\n",
    "    \n",
    "    def __init__(self, root, model_id, model, train_dataset, n_epochs=100, batch_size=128, lr=0.001, device=\"cuda\"):\n",
    "        super().__init__(root=root, id=model_id, model=model, train_dataset=train_dataset, n_epochs=n_epochs, batch_size=batch_size, lr=lr, device=device)\n",
    "        self._min_loss = np.inf\n",
    "    \n",
    "    @staticmethod\n",
    "    def load_model(root, model_id):\n",
    "        \n",
    "        def model_loader(hyperparams):\n",
    "            return ANNClassifier(**hyperparams[\"model\"][\"params\"])\n",
    "        \n",
    "        return devtorch.load_model(root, model_id, model_loader)\n",
    "    \n",
    "    def loss(self, output, target, model):\n",
    "        return F.cross_entropy(output, target.long())\n",
    "    \n",
    "    # Here we overwrite the on_epoch_complete hook\n",
    "    def on_epoch_complete(self, save):\n",
    "        train_loss = self.log[\"train_loss\"][-1]\n",
    "        \n",
    "        # If a new minimum loss was achieved we save the model\n",
    "        # otherwise if the loss spikes more than 5% compared to its minimum then we load\n",
    "        # load the checkpoint and reduce the learning rate\n",
    "        if train_loss < self._min_loss:\n",
    "            print(f\"Saving checkpoint train_loss={train_loss:.4f} < min_loss={self._min_loss:.4f}.\")\n",
    "            self._min_loss = train_loss\n",
    "            self.save_model()\n",
    "        elif train_loss > 1.05 * self._min_loss:\n",
    "            print(\"=========> Loading checkpoint and decaying lr <=========\")\n",
    "            self.lr *= 0.1\n",
    "            self.model = CheckpointTrainer.load_model(self.root, self.id)\n",
    "            self.optimizer = self.optimizer_func(self.model.parameters(), self.lr, **self.optimizer_kwargs)\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5a460107-2076-4f7e-9356-1904d6563162",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:trainer:Completed epoch 0 with loss 408.0909495726228 in 7.5060s\n",
      "Saving checkpoint train_loss=408.0909 < min_loss=inf.\n",
      "INFO:trainer:Completed epoch 1 with loss 92.43864496052265 in 7.4333s\n",
      "Saving checkpoint train_loss=92.4386 < min_loss=408.0909.\n",
      "INFO:trainer:Completed epoch 2 with loss 56.715450895018876 in 7.4289s\n",
      "Saving checkpoint train_loss=56.7155 < min_loss=92.4386.\n",
      "INFO:trainer:Completed epoch 3 with loss 45.24590137088671 in 7.4299s\n",
      "Saving checkpoint train_loss=45.2459 < min_loss=56.7155.\n",
      "INFO:trainer:Completed epoch 4 with loss 29.816457504639402 in 7.4320s\n",
      "Saving checkpoint train_loss=29.8165 < min_loss=45.2459.\n",
      "INFO:trainer:Completed epoch 5 with loss 22.350258031627163 in 7.4326s\n",
      "Saving checkpoint train_loss=22.3503 < min_loss=29.8165.\n",
      "INFO:trainer:Completed epoch 6 with loss 53.35907748359023 in 7.4237s\n",
      "=========> Loading checkpoint and decaying lr <=========\n",
      "INFO:trainer:Completed epoch 7 with loss 13.679868848761544 in 7.4394s\n",
      "Saving checkpoint train_loss=13.6799 < min_loss=22.3503.\n",
      "INFO:trainer:Completed epoch 8 with loss 11.271954031893983 in 7.4432s\n",
      "Saving checkpoint train_loss=11.2720 < min_loss=13.6799.\n",
      "INFO:trainer:Completed epoch 9 with loss 10.066672384506091 in 7.4440s\n",
      "Saving checkpoint train_loss=10.0667 < min_loss=11.2720.\n"
     ]
    }
   ],
   "source": [
    "root = \"../../data\"  # where to save the checkpoint\n",
    "model_id = \"ann\"  # the name of the checkpoint - if not is provided devtorch auto generates this.\n",
    "model = ANNClassifier(784, 2000, 10)\n",
    "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])\n",
    "train_dataset = datasets.MNIST(\"../../data\", train=True, download=True, transform=transform)\n",
    "test_dataset = datasets.MNIST(\"../../data\", train=False, download=True, transform=transform)\n",
    "\n",
    "trainer = CheckpointTrainer(root, model_id, model, train_dataset, n_epochs=10, batch_size=128, lr=0.01, device=\"cuda\")\n",
    "trainer.train(save=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cae2eef1-f236-4b85-b052-944bf92653d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy = 0.9639000296592712\n"
     ]
    }
   ],
   "source": [
    "def eval_metric(output, target):\n",
    "    return (torch.max(output, 1)[1] == target).sum().cpu().item()\n",
    "\n",
    "scores = devtorch.compute_metric(model, test_dataset, eval_metric, batch_size=256)\n",
    "print(f\"Accuracy = {torch.Tensor(scores).sum()/len(test_dataset)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6afcdd78-a536-4b4e-8b4b-134cd85b9fd0",
   "metadata": {},
   "source": [
    "**Exercise**: This tutorial just outlines one variant of how you might optimize the training process. Perhaps you can think of better and more exotic ways of doing this. As an exercise, you could try to extend the CheckpointTrainer to load model checkpoints and decay the learning rate whenever the training loss does not improve over a certain number of epochs."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:devtorch] *",
   "language": "python",
   "name": "conda-env-devtorch-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.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}