{ "cells": [ { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "# Prototyping Models" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Let's say you have a novel idea for a model architecture and you want to run ablation study on it with `ablator`. Ablator simplifies the process of prototyping your model, allowing you to swiftly construct and evaluate your innovative concept. Once a prototype runs smoothly, you can switch to parallel ablation study, which trains and runs HPO of different trials, with minimal code change for hyperparameter optimization.\n", "\n", "This chapter covers prototyping a model using Ablator, training the model on the popular **Fashion-mnist** dataset.\n", "\n", "There are 3 main steps to run a prototype experiment in ablator:\n", "\n", "- Configure the prototype experiment.\n", "\n", "- Create model wrapper that defines boiler-plate code for training and evaluating models.\n", "\n", "- Create the trainer and launch the experiment." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Let us first import all necessary dependencies:" ] }, { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "```python\n", "from ablator import ModelConfig, OptimizerConfig, TrainConfig, RunConfig\n", "from ablator import ModelWrapper, ProtoTrainer, configclass\n", "\n", "import torch\n", "import torch.nn as nn\n", "from torch.utils.data import DataLoader, Dataset\n", "import torchvision\n", "import torchvision.transforms as transforms\n", "\n", "from sklearn.metrics import f1_score, accuracy_score\n", "\n", "import shutil\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Launch the prototype experiment" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "### Configure the experiment" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "We will follow exactly the same steps as in the previous tutorial on [Configuration Basics](./Configuration-Basics.ipynb) to configure the experiment:\n", "\n", "Here's a summary of how we will configure it:\n", "\n", "- **Model Configuration**: dimensions for the layers of the model.\n", "\n", "- **Optimizer Configuration**: adam (lr = 0.001).\n", "\n", "- **Train Configuration**: `batch_size` = 32, `epochs` = 20, random weights initialization is set as true.\n", "\n", "- **Running Configuration**: CPU as hardware and a random seed for the experiment." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Configure the model" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "##### Model configuration\n", "\n", "For the model configuration, we defines hyperparameters `input_size`, `hidden_size`, and `num_classes` as integer config attributes." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "```python\n", "@configclass\n", "class CustomModelConfig(ModelConfig):\n", " input_size :int\n", " hidden_size :int \n", " num_classes :int\n", "\n", "model_config = CustomModelConfig(\n", " input_size = 28*28, \n", " hidden_size = 256, \n", " num_classes = 10\n", " )\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Since the hyperparameters are defined using primitive data type integer (aka Stateful), we must provide concrete values when initializing the `model_config` object. " ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "##### Creating Pytorch Model " ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Model Architecture (Simple Neural Network with Linear Layers):\n", "\n", "Linear_1_(28*28, 256) -> ReLU -> Linear_2_(256, 256) -> ReLU -> Linear_3_(256, 10). (where; ReLU is an Activation function) \n", "\n", "Note that here we depart from the Configuration Basics tutorial, we construct our model as a 2-level module:\n", "\n", "- `FashionMNISTModel` defines the model architecture (your novel idea), this is where we use the model config attributes to construct the model.\n", "\n", "- `MyModel` includes the main model architecture as a sub-module, adds a loss function, performs forward computation, and returns the predicted labels and loss during model training and evaluation. " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "```python\n", "class FashionMNISTModel(nn.Module):\n", " def __init__(self, config: CustomModelConfig):\n", " super(FashionMNISTModel, self).__init__()\n", "\n", " input_size = config.input_size \n", " hidden_size = config.hidden_size\n", " num_classes = config.num_classes\n", "\n", " self.fc1 = nn.Linear(input_size, hidden_size)\n", " self.relu1 = nn.ReLU()\n", " self.fc2 = nn.Linear(hidden_size, hidden_size)\n", " self.relu2 = nn.ReLU()\n", " self.fc3 = nn.Linear(hidden_size, num_classes)\n", " \n", " def forward(self, x):\n", " x = x.view(x.size(0), -1) \n", " x = self.fc1(x)\n", " x = self.relu1(x)\n", " x = self.fc2(x)\n", " x = self.relu2(x)\n", " x = self.fc3(x)\n", " return x\n", "\n", "class MyModel(nn.Module):\n", " def __init__(self, config: CustomModelConfig) -> None:\n", " super().__init__()\n", " \n", " self.model = FashionMNISTModel(config)\n", " self.loss = nn.CrossEntropyLoss()\n", "\n", " def forward(self, x, labels=None):\n", " out = self.model(x)\n", " loss = None\n", "\n", " if labels is not None:\n", " loss = self.loss(out, labels)\n", " labels = labels.reshape(-1, 1)\n", "\n", " out = out.argmax(dim=-1)\n", " out = out.reshape(-1, 1)\n", "\n", " return {\"y_pred\": out, \"y_true\": labels}, loss\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Configure the training process" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "```python\n", "optimizer_config = OptimizerConfig(\n", " name=\"adam\", \n", " arguments={\"lr\": 0.001}\n", ")\n", "\n", "train_config = TrainConfig(\n", " dataset=\"Fashion-mnist\",\n", " batch_size=32,\n", " epochs=20,\n", " optimizer_config=optimizer_config,\n", " scheduler_config=None,\n", " rand_weights_init = True\n", ")\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Configure the running configuration" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "```python\n", "@configclass\n", "class CustomRunConfig(RunConfig):\n", " model_config: CustomModelConfig\n", "\n", "run_config = CustomRunConfig(\n", " train_config=train_config,\n", " model_config=model_config,\n", " metrics_n_batches = 800,\n", " experiment_dir = \"/tmp/experiments\",\n", " device=\"cuda\",\n", " amp=False,\n", " random_seed = 42\n", ")\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "\n", "Note\n", "\n", "- We recommend that the experiment directory `RunConfig.experiment_dir` should be an empty directory.\n", "- Make sure to redefine the running configuration class to update its `model_config` attribute from `ModelConfig` (by default) to `CustomModelConfig` before creating the config object.\n", "\n", "
" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Create the model wrapper\n", "\n", "The model wrapper class `ModelWrapper` serves as a comprehensive wrapper for PyTorch models, providing a high-level interface for handling various tasks involved in model training. It defines boiler-plate code for training and evaluating models, which significantly reduces development efforts and minimizes the need for writing complex code, ultimately improving efficiency and productivity:\n", "\n", "- It takes care of creating and utilizing data loaders, evaluating models, importing parameters from configuration files into the model, setting up optimizers and schedulers, and checkpoints, logging metrics, handling interruptions, and much more.\n", "\n", "- Its functions are over-writable to support for custom use-cases (read more about these functions in [this documentation of Model Wrapper](../training.interface.rst)).\n", "\n", "An important function of the `ModelWrapper` is `make_dataloader_train`, which is used to create a data loader for training the model. In fact, you must provide a train dataloader to `make_dataloader_train` before launching the experiment.\n", "\n", "Therefore, we will start prepare the datasets first. Then, we write some eluation functions to be used to evaluate our model. Finally, we will create the model wrapper and train the model." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "#### Prepare the dataset\n", "\n", "**Fashion MNIST** is a dataset consisting of 60,000 grayscale images of fashion items. The images are categorized into ten classes, which include clothing items. \n", "\n", "- Image dimensions: 28 pixels x 28 pixels (grayscale)\n", "\n", "- Shape of the training data tensor: [60000, 1, 28, 28]\n", "\n", "Here we will create two datasets: one for training and one for validation." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "```python\n", "transform = transforms.ToTensor()\n", "\n", "train_dataset = torchvision.datasets.FashionMNIST(\n", " root='./data',\n", " train=True,\n", " download=True,\n", " transform=transform\n", ")\n", "\n", "test_dataset = torchvision.datasets.FashionMNIST(\n", " root='./data',\n", " train=False,\n", " download=True,\n", " transform=transform\n", ")\n", "```" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "#### Defining Custom Evaluation Metrics\n", "\n", "Defining evaluation functions for classification problems. Using average as \"weighted\" for multiclass evaluation." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "```python\n", "def my_accuracy(y_true, y_pred):\n", " return accuracy_score(y_true.flatten(), y_pred.flatten())\n", "\n", "def my_f1_score(y_true, y_pred):\n", " return f1_score(y_true.flatten(), y_pred.flatten(), average='weighted')\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "\n", "Note\n", "\n", "Make sure that parameters to the evaluation function match the model's forward dictionary output. Since MyModel's returned dictionary has keys `\"y_true\"` and `\"y_pred\"`, the evaluation function must have parameters `\"y_true\"` and `\"y_pred\"`.\n", "\n", "
" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "#### Create the Model Wrapper" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "We will now create a model wrapper class and overwrite the following functions:\n", "\n", "- `make_dataloader_train` and `make_dataloader_val`: to provide the training dataset and validation dataset as dataloaders (In PyTorch, a **DataLoader** is a utility class that provides an iterable over a dataset. It is commonly used for handling data loading and batching in machine learning and deep learning tasks).\n", "\n", "- `evaluation_functions`: to provide the evaluation functions that will evaluate the model on the datasets. In this function, you must return a dictionary of callables, where the keys are the names of the evaluation metrics and the values are the functions that compute the metrics." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "```python\n", "class MyModelWrapper(ModelWrapper):\n", " def __init__(self, *args, **kwargs):\n", " super().__init__(*args, **kwargs)\n", "\n", " def make_dataloader_train(self, run_config: CustomRunConfig):\n", " return torch.utils.data.DataLoader(\n", " train_dataset,\n", " batch_size=32,\n", " shuffle=True\n", " )\n", "\n", " def make_dataloader_val(self, run_config: CustomRunConfig):\n", " return torch.utils.data.DataLoader(\n", " test_dataset,\n", " batch_size=32,\n", " shuffle=False\n", " )\n", "\n", " def evaluation_functions(self):\n", " return {\n", " \"accuracy\": my_accuracy,\n", " \"f1\": my_f1_score\n", " }\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now create the model wrapper object, passing the model class as its argument:\n", "```python\n", "wrapper = MyModelWrapper(\n", " model_class=MyModel,\n", ")\n", "```" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "### Create the trainer and launch the experiment" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "For a prototype experiment, we will use the prototype trainer `ProtoTrainer` to launch the experiment.\n", "\n", "Initialize the trainer, providing it with the model wrapper and the running configuration. After that, calling the `launch()` method will start the training process." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "```python\n", "ablator = ProtoTrainer(\n", " wrapper=wrapper,\n", " run_config=run_config,\n", ")\n", "metrics = ablator.launch()\n", "```" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Experiment results" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "The `ProtoTrainer.launch()` method returns a dictionary which stores metrics of the experiment\n", "\n", "A more detailed exploration of interpreting results will be undertaken in a later chapter." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "```python\n", "max_key_length = max(len(str(k)) for k in metrics.keys())\n", "\n", "for k, v in metrics.items():\n", " print(f\"{k:{max_key_length}} : {v}\")\n", "```\n", "```shell\n", "val_loss : 0.5586626408626636\n", "val_accuracy : 0.8687149999999999\n", "val_f1 : 0.8684085851245271\n", "train_loss : 0.2816645764191945\n", "train_accuracy : 0.8915705128205127\n", "train_f1 : 0.891141942313593\n", "best_iteration : 3750\n", "best_loss : 0.4098668480262208\n", "current_epoch : 20\n", "current_iteration : 37500\n", "epochs : 20\n", "learning_rate : 0.001\n", "total_steps : 37500\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "How to visualize metrics\n", "\n", "ablator automatically records metrics so that you can visualize them in TensorBoard and observe how they change every epoch:\n", "\n", "- Just install `tensorboard`, import it, and load using `%load_ext tensorboard` if using a notebook.\n", "\n", "- Run the command `%tensorboard --logdir /dashboard/tensorboard --port [port]`, where `` is the experiment directory that we passed to the parallel config (`run_config.experiment_dir = \"/tmp/experiments/\"`)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Conclusion\n", "\n", "That's it! We have successfully built and tested a prototype model using ablator. In the later chapters, we will learn how to scale a prototype to a cluster of parallel processes to explore hyperparameter optimization with more complex models." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "#### Additional Info" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Why train with ProtoTrainer?\n", "\n", "- It provides a robust way to handle errors during training.\n", "- Ideal for prototyping experiments in a local environment.\n", "- Easily adaptable for hyperparameter optimization with larger configurations and horizontal scaling.\n", "- Quick transition to ````ParallelConfig```` and ````ParallelTrainer```` for parallel execution of trials using Ray." ] } ], "metadata": { "kernelspec": { "display_name": "env", "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.10.6" }, "orig_nbformat": 4 }, "nbformat": 4, "nbformat_minor": 2 }