{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Model Checkpointing" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Model Checkpointing\n", "\n", "So say you've trained a model, it trained fine, has a nice learning curve (train vs validation) and now you want to save it. Or, you want to save checkpoints of the model during specific stages of the training process and later, use the best checkpoints for inference. Technically, all flax modules are dataclasses and params (part of model state in flax) are what store the model, so what we need to do for checkpointing is to persist the params. \n", "\n", "However, that poses on small problem. You see, params from flax modules are not regular python data types. They're tree maps. In other libraries, e.g. Pytorch, you can save a state dict as a regular python dictionary. Such isn't compatible with flax. Instead, we use a package called [orbax](https://github.com/google/orbax) for checkpointing. \n", "\n", "Let's take the cnn from the previous notebook and add checkpointing to it. You can skip the notebook cells until the `Checkpointing` section as they have nothing to do with checkpointing. " ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/home/shawon/Code/jax_examples/venv/lib64/python3.12/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n", " _torch_pytree._register_pytree_node(\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Files already downloaded and verified\n", "Files already downloaded and verified\n" ] } ], "source": [ "import torch\n", "import jax_dataloader as jdl\n", "import torchvision\n", "import torchvision.transforms as transforms\n", "\n", "\n", "class ToNumpy:\n", " def __call__(self, x: torch.Tensor):\n", " return x.numpy()\n", " \n", " \n", "transform = transforms.Compose(\n", " [transforms.ToTensor(),\n", " transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ToNumpy()])\n", "\n", "\n", "trainset = torchvision.datasets.CIFAR10(root='./data', train=True,\n", " download=True, transform=transform)\n", "\n", "testset = torchvision.datasets.CIFAR10(root='./data', train=False,\n", " download=True, transform=transform)\n", "\n", "\n", "batch_size = 128\n", "\n", "trainloader = jdl.DataLoader(trainset, backend=\"pytorch\", batch_size=batch_size,\n", " shuffle=True)\n", "testloader = jdl.DataLoader(testset, backend=\"pytorch\", batch_size=batch_size,\n", " shuffle=False)\n", "\n", "# classes in cifar10\n", "classes = ('plane', 'car', 'bird', 'cat',\n", " 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "(10,)\n" ] } ], "source": [ "import jax\n", "import jax.numpy as jnp\n", "import numpy as np\n", "import flax\n", "import flax.linen as nn\n", "\n", "\n", "from einops import rearrange\n", "\n", "\n", "class ConvNet(nn.Module):\n", " @nn.compact\n", " def __call__(self, x):\n", " # convs\n", " out = nn.Conv(features=6, kernel_size=(5, 5))(x)\n", " out = nn.max_pool(out, window_shape=(2, 2))\n", " out = nn.Conv(features=16, kernel_size=(5, 5))(out)\n", " out = nn.max_pool(out, window_shape=(2, 2))\n", "\n", " # flatten into a vector\n", " # skip the batch dim\n", " if len(x.shape) > 3:\n", " out = rearrange(x, \"batch c h w -> batch (c h w)\")\n", " else:\n", " out = out.flatten()\n", "\n", " # dense\n", " out = nn.Dense(features=120)(out)\n", " out = nn.Dense(features=84)(out)\n", " out = nn.Dense(features=10)(out)\n", "\n", " return out\n", " \n", "\n", "# ======================\n", "model = ConvNet()\n", "rng = jax.random.key(0)\n", "params = model.init(rng, jnp.empty((3, 32, 32)))\n", "\n", "# run a sample forward pass\n", "logits = model.apply(params, jnp.empty((3, 32, 32)))\n", "print(logits.shape)" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "import optax\n", "from tqdm.notebook import trange, tqdm\n", "from flax.training import train_state\n", "\n", "\n", "@jax.jit\n", "def calculate_loss(params, x, y):\n", " logits = model.apply(params, x)\n", " loss = optax.softmax_cross_entropy_with_integer_labels(logits, y)\n", " return loss\n", "\n", "\n", "@jax.jit\n", "def batched_loss(params, xs, ys):\n", " batch_loss = jax.vmap(calculate_loss, in_axes=(None, 0, 0))(params, xs, ys)\n", " return batch_loss.mean(axis=-1)\n", "\n", "\n", "optimiser = optax.adam(learning_rate=0.001)\n", "criterion = jax.value_and_grad(batched_loss)\n", "state = train_state.TrainState.create(\n", " apply_fn=model.apply,\n", " params=params,\n", " tx=optimiser\n", ")\n", "\n", "\n", "@jax.jit\n", "def train_step(state, batch):\n", " loss_value, grads = criterion(state.params, *batch)\n", " updated_state = state.apply_gradients(grads=grads)\n", " return loss_value, updated_state" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "from sklearn.metrics import classification_report\n", "from sklearn.metrics import f1_score\n", "\n", "\n", "@jax.jit\n", "def test_step(state, xs):\n", " def infer(params, x):\n", " logits = model.apply(params, x)\n", " return jax.nn.softmax(logits, axis=-1)\n", "\n", " preds = jax.vmap(jax.jit(infer), in_axes=(None, 0))(state.params, xs)\n", " return preds\n", "\n", "\n", "def evaluate(state, test_loader):\n", " scores = list()\n", " for batch in tqdm(test_loader):\n", " xs, ys = batch\n", " preds = test_step(state, xs)\n", " preds = jnp.argmax(preds, axis=-1)\n", " f1 = f1_score(preds, ys, average=\"micro\")\n", " scores.append(f1)\n", "\n", " return np.array(scores).mean(axis=-1)\n", "\n", "\n", "def custom_classification_report(state, test_loader):\n", " preds = list()\n", " actual = list()\n", " for batch in tqdm(test_loader):\n", " xs, ys = batch\n", " pred = test_step(state, xs)\n", " pred = jnp.argmax(pred, axis=-1).tolist()\n", " preds.extend(pred)\n", " actual.extend(ys.tolist())\n", "\n", " clf = classification_report(preds, actual, target_names=classes)\n", " print(clf)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Checkpointing\n", "\n", "Let's first prepare orbax and then go through how checkpointing via it works. " ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "!rm -rf /tmp/cnn_cifar10_checkpointing_example\n", "# delete the existing checkpoints" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [ "import orbax\n", "from flax.training import orbax_utils\n", "\n", "# since everything in jax is a pytree\n", "# the checkpoints are basically the pytree versions of the params\n", "orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer()\n", "# checkpoint manager for managing how many checkpoints to keep\n", "# keep a max of 5 checkpoints\n", "options = orbax.checkpoint.CheckpointManagerOptions(max_to_keep=5, create=True)\n", "# structure for the checkpoint\n", "# will be used by orbax later to restore the saved model\n", "# you can also add other information regarding the model\n", "# but make sure to keep the structure consistent\n", "ckpt = {\n", " \"state\": state,\n", "}\n", "# tell orbax to use the structure\n", "save_args = orbax_utils.save_args_from_target(ckpt)\n", "# add the save path\n", "save_path = \"/tmp/cnn_cifar10_checkpointing_example\"\n", "# ckpt manager for versioning\n", "checkpoint_manager = orbax.checkpoint.CheckpointManager(save_path, orbax_checkpointer, options)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Breaking down the mess\n", "\n", "Since we're training using train_state from flax, we update the params with the `train_state`. So it would be wise to checkpoint the entire state instead of the params only. Now, there are two stages of checkpointing using Orbax. First, we have `PyTreeCheckpointer`, which can save a single checkpoint. It doesn't keep track of updates over training iterations so no matter how many iterations your training runs for, it'll only save a checkpoint on the first call. To track different checkpoints throughout training, we need `CheckpointManager`.\n", "\n", "Now, orbax will save the state as a pytree object. But if we want to restore a full state from it, the checkpoint manager provides no such direct API. So we have to improvise a bit (just follow through, comes right after the training function). Furthermore, orbax has no clue about the data structure of your checkpoints. All it cares is that it gets a pytree, as long as you supply the structure for it. The `ckpt` dict here provides the structure to orbax. It basically acts as a schema for the checkpoints.\n", "\n", "Sounds complicated? Kinda is. You see in Pytorch, you can just use the model class and map a saved state_dict to it. Could flax have made it simpler? May be! You can read more [here](https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html).\n", "\n", "*N.B: Checkpoint manager creates the checkpoint directory during init and maintains the directory and checkpoint metadata (paths etc.) as state. So if you want to rerun the training, delete the save_path directory first and re init the manager.*" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we can add the checkpointing code inside the `train` function." ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "def train(state, epochs, train_loader, test_loader, ckpt_manager=checkpoint_manager, save_args=save_args):\n", " steps = 0\n", " losses = []\n", " f1_scores = []\n", " \n", " lowest_loss = np.inf\n", "\n", " # =============\n", " for e in trange(epochs):\n", " for batch in tqdm(train_loader):\n", " loss, state = train_step(state, batch)\n", " steps += 1\n", "\n", " # log every 200 steps\n", " if steps % 200 == 0:\n", " losses.append(loss)\n", "\n", " # run evaluation\n", " print(\"Evaluating ... \")\n", " score = evaluate(state, test_loader)\n", "\n", " f1_scores.append(score)\n", "\n", " print(f\"Epoch : {e + 1} :: Step : {steps} :: Loss : {loss} :: F1 : {score}\")\n", " \n", " # checkpoint only if train loss decreases\n", " # ideally one would checkpoint on validation loss\n", " # but we don't have a validation split on the dataset\n", " # take this as an example\n", " \n", " if loss < lowest_loss:\n", " print(\"Checkpointing\")\n", " lowest_loss = loss\n", " # save model ckpt\n", " ckpt = {\n", " \"state\": state\n", " }\n", " ckpt_manager.save(steps, ckpt, save_kwargs={\"save_args\": save_args})\n", "\n", " # ============\n", " return state, losses, f1_scores" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "94543a2220f14ffd94bf618a9a7b4edb", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/5 [00:00