script_1_2_model_save_load_stat.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# __1.2. Test save/load checkpoint__\n",
    "\n",
    "Goal:\n",
    "- Test how checkpoint can be saved and loaded for retraining.\n",
    "\n",
    "Log:\n",
    "\n",
    "- 11/21/23\n",
    "  - Spent another 3 hours trying to figure out what's wrong:\n",
    "    - Setting resume_from_checkpoint=ckpt_path did not work. The doc said this should be: \"local path to a saved checkpoint as saved by a previous instance of Trainer\". \n",
    "    - Reading more about setting resume_from_checkpoint=True, it said \"If a bool and equals True, load the last checkpoint in args.output_dir as saved by a previous instance of Trainer.\" \n",
    "    - The output_dir is also mentioned in the TrainingArgument section as \"where the model predictions and checkpoints will be written\". In my case, this is the model_dir in the config file which is where the check point folders are located. \n",
    "    - Still does not work. Give ValueError: checkpoint_path is not in list.\n",
    " - Ok, I found out why.. Should have look into the error message more closely. \n",
    "   - Turned out the issue is in the \"best_model_checkpoint\" value in trainer_state.json. I ran 10 epochs when setting up the dual GPU server and #   has the project directory as:\n",
    "     - `/home/shiulab/docs_shius/plantbert``\n",
    "   - But this should be now:\n",
    "     - `/home/shius/projects/plantbert`\n",
    "   - Darn... a simple mistake that stumbed me for a week...\n",
    "- 11/20/23\n",
    "  - Resume pretraining on 11/17 but get cuda out of memory so the process was killed. This is because I change the batch size from 20 to 21... Change it back.\n",
    "  - Resume pretraining again using checkpoint-11500. But loss of the next check point rose unexpected before starting to go down rapidly. \n",
    "    - [Similar issue](https://github.com/Lightning-AI/lightning/issues/4045) has been reporte.\n",
    "    - [Another post](https://github.com/huggingface/transformers/issues/23099) with similar issue due to learning rate reset. This is because `trainer.train()` was being called without specify that the passed argument is for `resume_from_checkpoint`. Modified 1_1_bert_retrain_v3 to reflect this.\n",
    "- 11/17/23\n",
    "  - Realize that `transformers.Trainer` class already have abstraction to deal with save and load for retrain:\n",
    "  - E.g., see [this post](https://stackoverflow.com/questions/64663385/saving-and-reload-huggingface-fine-tuned-transformer) for example\n",
    "  - [Transformer doc on Trainer](https://huggingface.co/docs/transformers/main_classes/trainer)\n",
    "    - Particularly, look into `train(resume_from_checkpoint=XXX)`\n",
    "  - So the pytorch-based info is not as useful. See below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ___Setup___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, yaml, tarfile, json\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\n",
    "from transformers import BertTokenizerFast, BertConfig, BertForMaskedLM, \\\n",
    "                         DataCollatorForLanguageModeling, TrainingArguments, \\\n",
    "                         Trainer\n",
    "from tokenizers import BertWordPieceTokenizer\n",
    "from torchinfo import summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ___Pytorch save/load___\n",
    "\n",
    "[Tutorial on saving/loading checkpoint](\n",
    "https://pytorch.org/tutorials/recipes/recipes/saving_and_loading_a_general_checkpoint.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Background\n",
    "\n",
    "To save multiple checkpoints, you must\n",
    "- organize them in a dictionary\n",
    "- use `torch.save()` to serialize the dictionary\n",
    "\n",
    "A common PyTorch convention is to:\n",
    "- save these checkpoints using the `.tar`` file extension.\n",
    "- load the items by:\n",
    "  - initialize the model and optimizer\n",
    "  - load the dictionary locally using `torch.load()`\n",
    "  - access the saved items by simply querying the dictionary as you would expect."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Steps\n",
    "\n",
    "1. Import all necessary libraries for loading our data\n",
    "1. Define and initialize the neural network\n",
    "1. Initialize the optimizer\n",
    "1. Save the general checkpoint\n",
    "1. Load the general checkpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define and initialize the neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Net(\n",
      "  (conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))\n",
      "  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))\n",
      "  (fc1): Linear(in_features=400, out_features=120, bias=True)\n",
      "  (fc2): Linear(in_features=120, out_features=84, bias=True)\n",
      "  (fc3): Linear(in_features=84, out_features=10, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "class Net(nn.Module):\n",
    "  def __init__(self):\n",
    "    super(Net, self).__init__()\n",
    "    self.conv1 = nn.Conv2d(3, 6, 5)\n",
    "    self.pool = nn.MaxPool2d(2, 2)\n",
    "    self.conv2 = nn.Conv2d(6, 16, 5)\n",
    "    self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
    "    self.fc2 = nn.Linear(120, 84)\n",
    "    self.fc3 = nn.Linear(84, 10)\n",
    "\n",
    "  def forward(self, x):\n",
    "    x = self.pool(F.relu(self.conv1(x)))\n",
    "    x = self.pool(F.relu(self.conv2(x)))\n",
    "    x = x.view(-1, 16 * 5 * 5)\n",
    "    x = F.relu(self.fc1(x))\n",
    "    x = F.relu(self.fc2(x))\n",
    "    x = self.fc3(x)\n",
    "    return x\n",
    "\n",
    "net = Net()\n",
    "print(net)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialize the optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save the general checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCH = 5\n",
    "PATH = \"model.pt\"\n",
    "LOSS = 0.4\n",
    "\n",
    "torch.save({'epoch': EPOCH,\n",
    "            'model_state_dict': net.state_dict(),\n",
    "            'optimizer_state_dict': optimizer.state_dict(),\n",
    "            'loss': LOSS,\n",
    "            }, PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load the general checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Intialize model and optimizer first\n",
    "model     = Net()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)\n",
    "\n",
    "# Load checkpoint\n",
    "checkpoint = torch.load(PATH)\n",
    "\n",
    "# Load model state dict, optimizer, epoch, and loss\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "epoch = checkpoint['epoch']\n",
    "loss = checkpoint['loss']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Net(\n",
       "  (conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))\n",
       "  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))\n",
       "  (fc1): Linear(in_features=400, out_features=120, bias=True)\n",
       "  (fc2): Linear(in_features=120, out_features=84, bias=True)\n",
       "  (fc3): Linear(in_features=84, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ___Play with my own model___"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the general checkpoint\n",
    "\n",
    "Three steps skipped\n",
    "- Define and initialize the neural network\n",
    "- Initialize the optimizer\n",
    "- Save the general checkpoint\n",
    "\n",
    "- [how to continue training from a checkpoint with Trainer?](https://github.com/huggingface/transformers/issues/7198)\n",
    "  - __Ok, none of the following is necessary.__\n",
    "  - Just do: `trainer.train(checkpoint-dir)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "work_dir  = Path(\"/home/shius/projects/plantbert/\")\n",
    "model_dir = work_dir / \"models\"\n",
    "ckpt_dir  = model_dir / \"checkpoint-11500\" \n",
    "\n",
    "config_file = \"./config.yaml\"\n",
    "\n",
    "with open(config_file, 'r') as f:\n",
    "  config = yaml.safe_load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cannot just load the model by pointing out what the directory is \n",
    "#checkpoint = torch.load(ckpt_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize the model with the config\n",
    "vocab_size   = config['tokenize']['vocab_size']\n",
    "max_length   = config['tokenize']['max_length']\n",
    "model_config = BertConfig(vocab_size=vocab_size, \n",
    "                          max_position_embeddings=max_length)\n",
    "\n",
    "model = BertForMaskedLM(config=model_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['state', 'param_groups'])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optimizer = torch.load(ckpt_dir / \"optimizer.pt\")\n",
    "optimizer.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Intialize model and optimizer first\n",
    "model     = Net()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)\n",
    "\n",
    "# Load checkpoint\n",
    "checkpoint = torch.load(PATH)\n",
    "\n",
    "# Load model state dict, optimizer, epoch, and loss\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "epoch = checkpoint['epoch']\n",
    "loss = checkpoint['loss']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Net(\n",
       "  (conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))\n",
       "  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))\n",
       "  (fc1): Linear(in_features=400, out_features=120, bias=True)\n",
       "  (fc2): Linear(in_features=120, out_features=84, bias=True)\n",
       "  (fc3): Linear(in_features=84, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bert",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}