{ "cells": [ { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "# Prototyping Models" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/fostiropoulos/ablator/blob/v0.0.1-mp/docs/source/notebooks/Prototyping-models.ipynb)\n", "\n", "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 and behave as expected, you can scale it to a parallel ablation study of multiple trials with minimal code change.\n", "\n", "This chapter covers prototyping a model using Ablator. We will train a simple neural network model on the popular **Fashion-mnist** dataset.\n", "\n", "There are three 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 prototype trainer and launch the experiment." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Let us first import all necessary dependencies:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "try:\n", " import ablator\n", "except:\n", " !pip install ablator\n", " print(\"Stopping RUNTIME! Please run again\") # This script automatically restart runtime (if ablator is not found and installing is needed) so changes are applied\n", " import os\n", "\n", " os.kill(os.getpid(), 9)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from ablator import (ModelConfig, OptimizerConfig, TrainConfig, RunConfig,\n", " ModelWrapper, ProtoTrainer, configclass)\n", "\n", "import torch\n", "import torch.nn as nn\n", "import torchvision\n", "import torchvision.transforms as transforms\n", "\n", "from sklearn.metrics import f1_score, accuracy_score\n", "\n", "import os" ] }, { "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 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.\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 define hyperparameters `input_size`, `hidden_size`, and `num_classes` as [stateful](./Configuration-Basics.ipynb#Ablator-custom-data-types-for-stateful-experiment-design) integer config attributes." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "CustomModelConfig(input_size=784, hidden_size=256, num_classes=10)" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "@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, # 28x28 image flattened\n", " hidden_size = 256, \n", " num_classes = 10\n", " )\n", "model_config" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Since the hyperparameters are defined as stateful, we must provide concrete values when initializing the `model_config` object.\n", "\n", "
\n", "\n", "Note\n", "\n", "In this tutorial, the model config is just used to construct the main model. In fact, if you don't plan to scale your prototype to an ablation study, you can skip defining a model config and directly construct the model. However, it is a good practice to define a model config, because ablator is mostly used for scaling up a prototype to an ablation study. Once shifted to scaling up, you will see that model config plays a critical role, specifically, it lets you create search spaces for hyperparameters and in turn run ablation study on them.\n", "\n", "
" ] }, { "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` is 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": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "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) # reshape as (batch, labels)\n", "\n", " out = out.argmax(dim=-1)\n", " out = out.reshape(-1, 1) # reshape as (batch, output)\n", "\n", " return {\"y_pred\": out, \"y_true\": labels}, loss" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "\n", "Note\n", "\n", "- Ablator requires the model's forward function to return two objects: one dictionary of model's batched output (e.g. labels, predictions, logits, probabilities, etc.), and the other is the loss value. Notice that these values must be tensors. You also have the choice to return `None` for either of the values, depending on the use case.\n", "\n", "- Depending on the evaluation metrics that you want to use, you can include in the model's dictionary output logits, probabilities, predicted labels, ground truth labels, etc. In this example, we return the predicted labels and the ground truth labels in the model's dictionary output, and these will be used later on to compute the accuracy and F1 score.\n", "\n", "
" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Configure the training process" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "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", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Configure the running configuration" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@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=\"cpu\",\n", " amp=False,\n", " random_seed = 42\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, or at least does not contain any prior experiment results.\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 prepare the datasets first. Then, we write some evaluation functions that will be used to evaluate our model. Finally, we will create the model wrapper, pass it and the configuration to the trainer and launch the experiment." ] }, { "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": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "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", ")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "#### Defining Custom Evaluation Metrics\n", "\n", "Evaluation metrics are a way for you to evaluate the model using different methods. . Defining evaluation functions for classification problems. Using average as \"weighted\" for multiclass evaluation." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "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')" ] }, { "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. These metrics are later used for logging and plotting (i.e. tensorboard tracking and analysis artifacts)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "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", " }" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now create the model wrapper object, passing the model class as its argument:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "wrapper = MyModelWrapper(\n", " model_class=MyModel,\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, call the `launch()` method, passing to `working_directory` the path to the main directory that you're working at (which stores codes, modules that will be pushed to ray). It's recommended that this directory be tracked by git for keeping track of any code changes." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ablator = ProtoTrainer(\n", " wrapper=wrapper,\n", " run_config=run_config,\n", ")\n", "\n", "metrics = ablator.launch(working_directory=os.getcwd()) # assuming the current directory is tracked by git" ] }, { "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." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "max_key_length = max(len(str(k)) for k in metrics)\n", "\n", "for k, v in metrics.items():\n", " print(f\"{k:{max_key_length}} : {v}\")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "```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": [ "Ablator automatically records metrics so that you can visualize them in TensorBoard and observe how they change every epoch:\n", "\n", "- Just install `pip install tensorboard` 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 ablation studies 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 }