{ "cells": [ { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "FxRTm4zWpwoo" }, "source": [ "# Ablation experiment" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "uHitmULK0W1v" }, "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/HPO-tutorial.ipynb)\n", "\n", "After your prototype has been verified and runs smoothly with `ProtoTrainer`, you can scale it to an ablation study/ HPO and analyze the results.\n", "\n", "In this chapter, we will learn how to set up and launch a parallel experiment for an ablation study with Ablator.\n", "\n", "Similarly to launching a prototype experiment, there are also 3 main steps to run an ablation experiment in ablator:\n", "\n", "- Configure the parallel 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.\n", "\n", "Recall from the [Introduction tutorial](./Introduction.ipynb), Ablator combines Optuna for running multiple trials from defined search spaces and Ray back-end for parallelizing the trials. So, an extra step is to start a ray cluster before launching the experiment (but if you don't want to do this, abator will automatically start a ray cluster for you)." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "iw2ie4ik0fiP" }, "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": { "id": "bPChwlLe_BUq" }, "outputs": [], "source": [ "from ablator import ModelConfig, OptimizerConfig, TrainConfig, ParallelConfig\n", "from ablator import ModelWrapper, ParallelTrainer, configclass\n", "from ablator.config.hpo import SearchSpace\n", "\n", "import torch\n", "import torch.nn as nn\n", "import torchvision\n", "import torchvision.transforms as transforms\n", "\n", "import shutil\n", "from sklearn.metrics import accuracy_score" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Launch the parallel experiment" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Configure the experiment" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "2u5Ovi0-YrtJ" }, "source": [ "We will follow the same steps as in the tutorial on [Prototyping models](./Prototyping-models.ipynb) to configure the experiment. Here's a summary of how we will configure it:\n", "\n", "- **Model Configuration**: defines hyperparameters for the number of filters and activation function.\n", "\n", "- **Optimizer Configuration**: adam (lr = 0.001).\n", "\n", "- **Train Configuration**: `batch_size = 32`, `epochs = 10`.\n", "\n", "- **Runing Configuration**:\n", " - GPU as hardware, a random seed for the experiment.\n", " - We let the experiment runs HPO for `total_trials = 20` trials, allowing `concurrent_trials = 2` trials to run in parallel.\n", " - We also use a search space for the model and the optimizer.\n", " - Use validation loss as the metric to optimize, in specific, we want to minimize this (`{\"val_loss\": \"min\"}`).\n", "\n", "
\n", "\n", "Note\n", "\n", "By default, any parallel experiment will be an ablation study (`search_algo` is set to `SearchAlgo.random`). However, if the experiment is HPO, you can set `search_algo` to `SearchAlgo.tpe` to use TPE algorithm for HPO.\n", "\n", "
" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Configure the model" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "##### Model configuration\n", "\n", "For the model configuration, we defines the following hyperparameters:\n", "\n", "- `num_filter1`, `num_filter2` (integer): number of filters at each convolutional layer\n", "\n", "- `activation` (string): activation function to use in layers." ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "id": "JUcqbxNe_3P-" }, "outputs": [ { "data": { "text/plain": [ "CustomModelConfig(num_filter1=32, num_filter2=64, activation='relu')" ] }, "execution_count": 2, "metadata": {}, "output_type": "execute_result" } ], "source": [ "@configclass\n", "class CustomModelConfig(ModelConfig):\n", " num_filter1: int\n", " num_filter2: int\n", " activation: str\n", "\n", "model_config = CustomModelConfig(\n", " num_filter1 =32,\n", " num_filter2 = 64,\n", " activation = \"relu\"\n", ")\n", "\n", "model_config" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Since the hyperparameters are defined as [stateful](./Configuration-Basics.ipynb#Ablator-custom-data-types-for-stateful-experiment-design), we must provide concrete values when initializing the `model_config` object. " ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "PKPFcebsYvJY" }, "source": [ "##### Creating Pytorch CNN Model\n", "\n", "We define a custom CNN model `FashionCNN` with the following architecture:\n", "\n", "- The first convolutional layer: takes a single channel and applies `num_filter1` filters to it. Then, applies an activation function and a max pooling layer.\n", "\n", "- The second convolutional layer: takes `num_filter1` channels and applies `num_filter2` filters to them. It also utilizes an activation function and a pooling layer.\n", "\n", "- The third convolutional layer: This is an additional layer that applies `num_filter2` filters.\n", "\n", "- A flattening layer: converts the convolutional layers into a linear format and subsequently produces a 10-dimensional output for labeling.\n", "\n", "`FashionCNN` is then included in `MyModel` as a sub-module. `MyModel`'s forward function performs forward computation, add a loss function, and returns the predicted labels and loss during model training and evaluation. " ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "id": "U3TMcBbNAFGV" }, "outputs": [], "source": [ "# Define the model\n", "class FashionCNN(nn.Module):\n", " def __init__(self, config: CustomModelConfig):\n", " super(FashionCNN, self).__init__()\n", "\n", " activation_list = {\"relu\": nn.ReLU(), \"elu\": nn.ELU(), \"leakyRelu\": nn.LeakyReLU()}\n", "\n", " num_filter1 = config.num_filter1\n", " num_filter2 = config.num_filter2\n", " activation = activation_list[config.activation]\n", "\n", " self.conv1 = nn.Conv2d(1, num_filter1, kernel_size=3, stride=1, padding=1)\n", " self.act1 = activation\n", " self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n", "\n", " self.conv2 = nn.Conv2d(num_filter1, num_filter2, kernel_size=3, stride=1, padding=1)\n", " self.act2 = activation\n", " self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)\n", "\n", " self.conv3 = nn.Conv2d(num_filter2, num_filter2, kernel_size=3, stride=1, padding=1)\n", " self.act3 = activation\n", "\n", " self.flatten = nn.Flatten()\n", " self.fc1 = nn.Linear(num_filter2 * 7 * 7, 10)\n", " \n", "\n", " def forward(self, x):\n", " x = self.conv1(x)\n", " x = self.act1(x)\n", " x = self.maxpool1(x)\n", " x = self.conv2(x)\n", " x = self.act2(x)\n", " x = self.maxpool2(x)\n", " x = self.conv3(x)\n", " x = self.act3(x)\n", " x = self.flatten(x)\n", " x = self.fc1(x)\n", " \n", " return x\n", "\n", "class MyModel(nn.Module):\n", " def __init__(self, config: CustomModelConfig) -> None:\n", " super().__init__()\n", "\n", " self.model = FashionCNN(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" ] }, { "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 score.\n", "\n", "
" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Configure the training process" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "TrainConfig(dataset='Fashion-mnist', batch_size=32, epochs=10, optimizer_config={'name': 'adam', 'arguments': {'betas': (0.9, 0.999), 'weight_decay': 0.0, 'lr': 0.001}}, scheduler_config=None)" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "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=10,\n", " optimizer_config=optimizer_config,\n", " scheduler_config=None\n", ")\n", "\n", "train_config" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Configure the running configuration\n", "\n", "To run an ablation study, we need to specify a search space for the hyperparameters of interest. This search space will then be used to configure the running configuration." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "r847W0MDY4Ai" }, "source": [ "##### Search Space\n", "\n", "For this tutorial, we have defined `search_space` object for four different hyperparameters:\n", "\n", "- Number of filters in the first and second convolutional layers: range between 32 and 64, and 64 and 128, respectively.\n", "\n", "- The activation function to use: any of `relu`, `elu`, and `leakyRelu`.\n", "\n", "- Learning rate value: ranges between 1e-3 and 1e-2." ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "id": "AJfHDNVUBNBB" }, "outputs": [], "source": [ "search_space = {\n", " \"model_config.num_filter1\": SearchSpace(value_range = [32, 64], value_type = 'int'),\n", " \"model_config.num_filter2\": SearchSpace(value_range = [64, 128], value_type = 'int'),\n", " \"train_config.optimizer_config.arguments.lr\": SearchSpace(\n", " value_range = [0.001, 0.01],\n", " value_type = 'float'\n", " ),\n", " \"model_config.activation\": SearchSpace(categorical_values = [\"relu\", \"elu\", \"leakyRelu\"])\n", "}" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "LWBSub7uZM6W" }, "source": [ "#### Parallel Configuration\n", "\n", "As the last step to configure the experiment, we pass `search_space`, `train_config`, and `model_config` to the `ParallelConfig`. Other parameters are also set (refer to this [Configuration Basics section](./Configuration-Basics.ipynb#Configuration-categories) or this [Config module documentation](../config.train.parallel.experiment.rst) for more details on the list of possible parameters to pass and their meanings, as well as the default values):" ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "id": "aAxNLP2QZNyI" }, "outputs": [], "source": [ "@configclass\n", "class CustomParallelConfig(ParallelConfig):\n", " model_config: CustomModelConfig\n", "\n", "parallel_config = CustomParallelConfig(\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=True,\n", " random_seed = 42,\n", " total_trials = 20,\n", " concurrent_trials = 2,\n", " search_space = search_space,\n", " optim_metrics = {\"val_loss\": \"min\"},\n", " optim_metric_name = \"val_loss\",\n", " gpu_mb_per_experiment = 1024,\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "\n", "Note\n", "\n", "- We recommend that the experiment directory `ParallelConfig.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, 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 preparing the datasets first. Then, we will create the model wrapper, pass it and the configuration to the trainer and launch the experiment." ] }, { "cell_type": "markdown", "metadata": { "id": "DtrEqbVGZCZX" }, "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": 7, "metadata": { "id": "qciJdRxT4v1M" }, "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": { "id": "bIQiuoj2Y6VJ" }, "source": [ "#### Create the Model Wrapper\n", "We will now create a model wrapper class and overwrite the following functions, similar to [Prototyping models](./Prototyping-models.ipynb) tutorial.\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": "code", "execution_count": 8, "metadata": { "id": "mE38Vg5HAKPK" }, "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: CustomParallelConfig):\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: CustomParallelConfig):\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\": lambda y_true, y_pred: accuracy_score(y_true.flatten(), y_pred.flatten()),\n", " }" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now create the model wrapper object, passing the model class as its argument:" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [], "source": [ "wrapper = MyModelWrapper(\n", " model_class=MyModel,\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "aHywLWRfZb5e" }, "source": [ "### Create the trainer and launch the experiment with `ParallelTrainer`\n", "\n", "`ParallelTrainer`, an extention from `ProtoTrainer`, is responsible for creating and pushing trials to the Ray cluster for parallelization of the ablation study.\n", "\n", "- We first initialize the trainer, providing it with the model wrapper and the running configuration.\n", "\n", "- Next, 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). This directory should also be a tracked by git for keeping track of any code changes.\n", "\n", "- Here we did not specify the ray address, this means ablator will automatically set up a ray cluster in the local machine, and all trials will be populated to this local cluster." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "ZE_wIo66AMdn" }, "outputs": [], "source": [ "shutil.rmtree(parallel_config.experiment_dir, ignore_errors=True) # Assure empty experiment directory\n", "\n", "ablator = ParallelTrainer(\n", " wrapper=wrapper,\n", " run_config=parallel_config,\n", ")\n", "\n", "ablator.launch(working_directory = os.getcwd()) # assuming the current directory is tracked by git, ray_head_address=None as default" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "\n", "Note\n", "\n", "- By default, `ablator.launch(working_directory = os.getcwd())` will initialize a ray cluster on your machine, and this cluster will be used for the experiment.\n", "\n", "- You have the option to scale the experiment to a cluster that's running somewhere else (e.g. on a cloud service like AWS). Given a ray cluster, you can use `ablator.launch(working_directory = os.getcwd(), ray_address =
)` to launch the experiment on that cluster.\n", "\n", "- To learn about running ablation experiments on cloud ray clusters, refer to [Launch-in-cloud-cluster](./Searchspace-for-diff-optimizers.ipynb) tutorial.\n", "\n", "
\n" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "We can provide `resume = True` to the `launch()` method to resume training the model from existing checkpoints and existing experiment state. Refer to the [Resume experiments](./Searchspace-for-diff-optimizers.ipynb) tutorial for more details." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "tRK0EQBuZ9vB" }, "source": [ "## Visualizing experiment results in TensorBoard\n", "\n", "Since ablator automatically stores TensorBoard event files for each training process, we can perform a short visualization with TensorBoard to compare how trials perform:\n", "\n", "- Install `tensorboard` and load using `%load_ext tensorboard` if using a notebook.\n", "\n", "- Run the command `%tensorboard --logdir /experiments_ --port [port]`, where `` is the experiment directory that we passed to the parallel config (`parallel_config.experiment_dir = \"/tmp/experiments/\"`), and `experiments_` is generated by ablator." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%load_ext tensorboard\n", "%tensorboard --logdir /tmp/experiments/ --port 6008" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "![TensorBoard-Output](./Images/tensorboard-output.jpg)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "More detailed analysis for ablation studies will be explored in later tutorials." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "mU9xvrfgZ1FJ" }, "source": [ "## Conclusion\n", "\n", "Finally, after completing all the trials, metrics obtained in each trial will be stored in the `experiment_dir`. This directory contains subdirectories representing the trials, as well as the experiment's state.\n", "\n", "Components stored in each trial directory are: best_checkpoints, checkpoints, results, training log, configurations, and metadata.\n", "\n", "To learn more, you can read the [Experiment output directory](./Experiment-dir.ipynb) tutorial, which explains the content of the experiment directory in detail.\n", "\n", "In the next tutorial, we will learn how to analyze the results from the trained trials." ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "T4", "provenance": [] }, "kernelspec": { "display_name": "Python 3", "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" } }, "nbformat": 4, "nbformat_minor": 0 }