pl_save_checkpoint.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e9a5b445",
   "metadata": {},
   "source": [
    "## Save and Load using PL example\n",
    "\n",
    "This is written for testing how overwrting `on_save_checkpoint()` works.\n",
    "\n",
    "In results, this script will save and load `my_checkpoint.ckpt` in current directory."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0648d23",
   "metadata": {},
   "source": [
    "## 1. Import Library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bf2bd855",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from copy import deepcopy\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.optim import SGD\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from torchvision import models, transforms\n",
    "from torchvision.datasets import CIFAR10\n",
    "\n",
    "from pytorch_lightning import LightningModule, LightningDataModule, Trainer\n",
    "from torchmetrics import Accuracy, MetricCollection\n",
    "\n",
    "os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba8baf64",
   "metadata": {},
   "source": [
    "## 2. Define dataloader "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cfdf9f4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CIFAR(LightningDataModule):\n",
    "    def __init__(self, img_size=32, batch_size=32):\n",
    "        super().__init__()\n",
    "        self.img_size = img_size if isinstance(img_size, tuple) else (img_size, img_size)\n",
    "        self.batch_size = batch_size\n",
    "        self.train_transforms = transforms.Compose([\n",
    "            transforms.Resize(self.img_size),\n",
    "            transforms.Pad(4, padding_mode='reflect'),\n",
    "            transforms.RandomCrop(self.img_size),\n",
    "            transforms.RandomHorizontalFlip(),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n",
    "        ])\n",
    "\n",
    "        self.test_transforms = transforms.Compose([\n",
    "            transforms.Resize(self.img_size),\n",
    "            transforms.CenterCrop(self.img_size),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n",
    "        ])\n",
    "\n",
    "    def prepare_data(self) -> None:\n",
    "        CIFAR10(root='data', train=True, download=True)\n",
    "        CIFAR10(root='data', train=False, download=True)\n",
    "    \n",
    "    def setup(self, stage=None):\n",
    "        self.train_ds = CIFAR10(root='data', train=True, download=False, transform=self.train_transforms)\n",
    "        self.valid_ds = CIFAR10(root='data', train=False, download=False, transform=self.test_transforms)\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(self.train_ds, num_workers=4, batch_size=self.batch_size, shuffle=True)\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(self.valid_ds, num_workers=4, batch_size=self.batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4194e991",
   "metadata": {},
   "source": [
    "## 3. Define Model with on_save_checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ae111c90",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BasicModule(LightningModule):\n",
    "    def __init__(self, lr=0.01):\n",
    "        super().__init__()\n",
    "        self.model = models.resnet18(pretrained=False)\n",
    "        self.criterion = nn.CrossEntropyLoss()\n",
    "        self.lr = lr\n",
    "        \n",
    "        metric = MetricCollection({'top@1': Accuracy(top_k=1), 'top@5': Accuracy(top_k=5)})\n",
    "        self.train_metric = metric.clone(prefix='train_')\n",
    "        self.valid_metric = metric.clone(prefix='valid_')\n",
    "    \n",
    "    def training_step(self, batch, batch_idx, optimizer_idx=None):\n",
    "        return self.shared_step(*batch, self.train_metric)\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        return self.shared_step(*batch, self.valid_metric)\n",
    "\n",
    "    def shared_step(self, x, y, metric):\n",
    "        y_hat = self.model(x)\n",
    "        loss = self.criterion(y_hat, y)\n",
    "        self.log_dict(metric(y_hat, y), prog_bar=True)\n",
    "        return loss\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return SGD(self.model.parameters(), lr=self.lr)\n",
    "    \n",
    "    def on_save_checkpoint(self, checkpoint):\n",
    "        with open('my_checkpoint.ckpt', 'wb') as f:\n",
    "            torch.save(self.model.state_dict(), f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97218c1b",
   "metadata": {},
   "source": [
    "## 4. Train only 2 epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cfb057e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hankyul/.local/lib/python3.7/site-packages/pytorch_lightning/core/datamodule.py:122: LightningDeprecationWarning: DataModule property `train_transforms` was deprecated in v1.5 and will be removed in v1.7.\n",
      "  \"DataModule property `train_transforms` was deprecated in v1.5 and will be removed in v1.7.\"\n",
      "/home/hankyul/.local/lib/python3.7/site-packages/pytorch_lightning/core/datamodule.py:160: LightningDeprecationWarning: DataModule property `test_transforms` was deprecated in v1.5 and will be removed in v1.7.\n",
      "  \"DataModule property `test_transforms` was deprecated in v1.5 and will be removed in v1.7.\"\n",
      "Using 16bit native Automatic Mixed Precision (AMP)\n",
      "GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hankyul/.local/lib/python3.7/site-packages/pytorch_lightning/core/datamodule.py:115: LightningDeprecationWarning: DataModule property `train_transforms` was deprecated in v1.5 and will be removed in v1.7.\n",
      "  \"DataModule property `train_transforms` was deprecated in v1.5 and will be removed in v1.7.\"\n",
      "/home/hankyul/.local/lib/python3.7/site-packages/pytorch_lightning/core/datamodule.py:153: LightningDeprecationWarning: DataModule property `test_transforms` was deprecated in v1.5 and will be removed in v1.7.\n",
      "  \"DataModule property `test_transforms` was deprecated in v1.5 and will be removed in v1.7.\"\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7,8]\n",
      "\n",
      "  | Name         | Type             | Params\n",
      "--------------------------------------------------\n",
      "0 | model        | ResNet           | 11.7 M\n",
      "1 | criterion    | CrossEntropyLoss | 0     \n",
      "2 | train_metric | MetricCollection | 0     \n",
      "3 | valid_metric | MetricCollection | 0     \n",
      "--------------------------------------------------\n",
      "11.7 M    Trainable params\n",
      "0         Non-trainable params\n",
      "11.7 M    Total params\n",
      "23.379    Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation sanity check: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6cc02b6d13f748d7b4859d54672234d1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data = CIFAR(batch_size=512)\n",
    "model = BasicModule(lr=0.01)\n",
    "trainer = Trainer(max_epochs=2, gpus='0,', precision=16)\n",
    "trainer.fit(model, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47383090",
   "metadata": {},
   "source": [
    "## 5. Load state dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c12562e7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = models.resnet18(pretrained=False)\n",
    "model.load_state_dict(torch.load('my_checkpoint.ckpt'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}