{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prototyping Models"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* The purpose of prototyping a model is to quickly build and test it, later enabling parallel training of different trials with minimal code change for hyperparameter optimization.\n",
    "* This chapter covers prototyping a model using Ablator with a popular **Fashion-mnist** dataset. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Running Experiments using Ablator"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Running an experiment involves loading configurations, training the model, and producing metrics. Ablator utilizes configurations, a model wrapper, and a trainer class to run an experiment for the given prototype. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Setting up Ablator"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install ablator using the command: ````pip install ablator````\n",
    "\n",
    "Import the **Configs**, **ModelWrapper**, and **ProtoTrainer** from the ablator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ablator import ModelConfig, OptimizerConfig, TrainConfig, RunConfig\n",
    "from ablator import ModelWrapper, ProtoTrainer, Stateless, Derived, configclass\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\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 os\n",
    "import shutil"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Configurations"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each config class has its arguments and serves a specific purpose in defining the configuration for the experiment."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Defining Configs:\n",
    "\n",
    "- **Optimizer Config**: adam (lr = 0.001).\n",
    "- **Train Config**: batch_size = 32, epochs = 20, random weights initialization is set as true.\n",
    "- **Run Config**: device details, experiment directory, and a random seed for the experiment.\n",
    "- **Model Config**: dimensions for the layers of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "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, \n",
    "    hidden_size = 256, \n",
    "    num_classes = 10\n",
    "    )\n",
    "\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",
    "\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=\"cpu\",\n",
    "    amp=False,\n",
    "    random_seed = 42\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Importing 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",
    "Shape of the training data tensor: [60000, 1, 28, 28]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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": [
    "#### 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",
    "````MyModel```` defines a model class that extends an existing model, ````FashionMNISTModel````. It adds a loss function, performs forward computation, and returns the predicted labels and loss during model training and evaluation. \n",
    "\n",
    "It is required to return the outputs and loss in the forward method of ````MyModel````. The outputs must be in a dictionary format. Example: ````{\"y_pred\": out, \"y_true\": labels}````."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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",
    "\n",
    "        out = out.argmax(dim=-1)\n",
    "\n",
    "        return {\"y_pred\": out, \"y_true\": labels}, loss"
   ]
  },
  {
   "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": "code",
   "execution_count": 5,
   "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')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Model Wrapper"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- This class serves as a comprehensive wrapper for PyTorch models, providing a high-level interface for handling various tasks involved in model training.\n",
    "\n",
    "- It takes care of importing parameters from configuration files into the model, setting up optimizers and schedulers, and checkpoints, logging metrics, handling interruptions, creating and utilizing data loaders, evaluating models, and much more.\n",
    "\n",
    "- Encapsulating these functionalities; significantly reduces development efforts and minimizes the need for writing complex code, ultimately improving efficiency and productivity."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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",
    "The data loaders and evaluation functions are passed to the ````ModelWrapper````"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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",
    "        }"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ProtoTrainer"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- This class is responsible to start training the model in the ````ModelWrapper```` and preparing resources for the model to avoid stalling during training or conflicts between other trainers.\n",
    "\n",
    "- Provides logging and syncing facilities to the provided directory or external remote servers like google cloud etc. It also does evaluation and syncing metrics to the directories.\n",
    "\n",
    "- Therefore, to achieve this, it requires ````ModelWrapper```` and ````run_config```` as inputs. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we wrap the model ````MyModel```` in a ModelWrapper ````MyModelWrapper````.\n",
    "Then, we create an instance of ````Prototrainer````, passing the **run_config** and **wrapper** as arguments, and then calling the ````launch()```` method to start the training.\n",
    "The ````launch()```` method returns an object of Class ````TrainMetrics````. It is used for calculating metrics for custom evaluation functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(run_config.experiment_dir):\n",
    "    shutil.os.mkdir(run_config.experiment_dir)\n",
    "\n",
    "shutil.rmtree(run_config.experiment_dir)\n",
    "\n",
    "wrapper = MyModelWrapper(\n",
    "    model_class=MyModel,\n",
    ")\n",
    "\n",
    "ablator = ProtoTrainer(\n",
    "    wrapper=wrapper,\n",
    "    run_config=run_config,\n",
    ")\n",
    "metrics = ablator.launch()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Results"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ````TrainMetrics```` object stores and manages predictions and calculates metrics using evaluation functions. We can access all the metrics from the ````TrainMetrics```` object using its ````to_dict()```` method.\n",
    "\n",
    "A more detailed exploration of interpreting results will be undertaken in a later chapter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_loss        : 2.3299765326192716\n",
      "val_loss          : 7.457184716395309\n",
      "train_accuracy    : 0.849365234375\n",
      "train_f1          : 0.8493114627262905\n",
      "val_accuracy      : 0.816905\n",
      "val_f1            : 0.816258432667835\n",
      "best_iteration    : 35625\n",
      "best_loss         : 7.65961674023207\n",
      "current_epoch     : 20\n",
      "current_iteration : 37500\n",
      "epochs            : 20\n",
      "learning_rate     : 0.001\n",
      "total_steps       : 37500\n"
     ]
    }
   ],
   "source": [
    "metrics_dict = metrics.to_dict()\n",
    "max_key_length = max(len(str(k)) for k in metrics_dict.keys())\n",
    "\n",
    "for k, v in metrics_dict.items():\n",
    "    print(f\"{k:{max_key_length}} : {v}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Conclusion\n",
    "\n",
    "Thus, we have successfully built and tested a prototype model using the ablator. In the later chapters, we will explore deeper into 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."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How to visualize metrics\n",
    "\n",
    "- We can also visualize metrics on TensorBoard with respect to every epoch.\n",
    "- Just install ````tensorboard````. Load using ````%load_ext tensorboard```` if using a notebook.\n",
    "- Run the command ````%tensorboard --logdir /tmp/experiments/[experiment_dir_name]/dashboard/tensorboard --port [port]````"
   ]
  }
 ],
 "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
}