{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quick start with Ablator"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Welcome to the Ablator tutorial! In this chapter, you will learn how to run an experiment from scratch. We will provide a simple demo to see what it looks like to run Ablator. You are also welcome to download this demo @[Colab](https://colab.research.google.com/drive/127l02PicoLxAZ3b_JL9eVpMxQ_UQKgiG?usp=sharing) or [Github](https://github.com/SeanXiaoby/ablator-fork/tree/tutorial-demos/examples/demo-basics-usage-vscode)\n",
    "\n",
    "Let's get started!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installing"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We assume you have installed Python (Python >= 3.10) and pip on your local machine. Please use the following command to install Ablator:\n",
    "\n",
    "```bash\n",
    "pip install ablator\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparations\n",
    "\n",
    "Below is a summary of the steps to launch an experiment with Ablator:\n",
    "\n",
    "- Set up experiment configurations\n",
    "- Define your idea - a model and datasets\n",
    "- Wrap model with model wrapper and launch experiment\n",
    "\n",
    "To showcase that Ablator takes care of training and evaluation of your idea (so you don't need to write training and evaluation scripts), in this demo, we will define a simple LeNet-5 model and classic MNIST dataset, wrap them with Ablator, and run the experiment.\n",
    "\n",
    "Before we start, let's import the necessary packages:"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "import shutil\n",
    "import argparse\n",
    "from typing import Any, Callable, Dict\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from sklearn.metrics import accuracy_score\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up configurations\n",
    "\n",
    "Configuring an experiment requires defining several configuration objects: model configuration, training configuration, optimizer configuration, and run configuration.\n",
    "\n",
    "In this chapter, we set up the experiment configurations using traditional Python objects and classes. The following code shows how to set up configurations for Ablator:"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "from ablator import ModelConfig, TrainConfig, OptimizerConfig, RunConfig,\n",
    "                     configclass, Literal, ModelWrapper, ProtoTrainer\n",
    "\n",
    "@configclass\n",
    "class SimpleConfig(ModelConfig):\n",
    "    name: Literal[\"simplenet\"]\n",
    "\n",
    "@configclass\n",
    "class SimpleRunConfig(RunConfig):\n",
    "    model_config: SimpleConfig\n",
    "\n",
    "run_config = SimpleRunConfig(\n",
    "    experiment_dir = \"/tmp/dir\",\n",
    "    train_config = TrainConfig(\n",
    "        dataset = \"mnist\",\n",
    "        batch_size = 64,\n",
    "        epochs = 10,\n",
    "        scheduler_config = None,\n",
    "        rand_weights_init = False,\n",
    "        optimizer_config = OptimizerConfig(\n",
    "            name = \"sgd\",\n",
    "            arguments = {\n",
    "                \"lr\": 0.001,\n",
    "                \"momentum\": 0.1\n",
    "            }\n",
    "        )\n",
    "    ),\n",
    "    model_config = SimpleConfig(name = \"simplenet\"),\n",
    "    metrics_n_batches = 200,\n",
    "    device= \"cpu\",\n",
    "    amp=False\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that there are multiple ways to set up the configuration. To learn more, go to the [Configuration Basics](./Configuration-Basics.ipynb) tutorial."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define your idea - a model and datasets\n",
    "\n",
    "This is where you define your novel idea (e.g. an original model or dataset). The following code shows how to define a model and datasets (note that you can customize this step to your needs):\n",
    "\n",
    "- Model detail:\n",
    "  - Define a simple CNN module using components from PyTorch packages.\n",
    "  - We will include the simple CNN as a part of the main model, define the loss function, and define the forward pass."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "class SimpleCNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SimpleCNN, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 6, 5)\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.pool1 = nn.MaxPool2d(2, 2)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
    "        self.relu2 = nn.ReLU()\n",
    "        self.pool2 = nn.MaxPool2d(2, 2)\n",
    "        self.fc1 = nn.Linear(16 * 4 * 4, 120)\n",
    "        self.relu3 = nn.ReLU()\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.relu4 = nn.ReLU()\n",
    "        self.fc3 = nn.Linear(84, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.pool1(self.relu1(self.conv1(x)))\n",
    "        x = self.pool2(self.relu2(self.conv2(x)))\n",
    "        x = x.view(-1, 16 * 4 * 4)\n",
    "        x = self.relu3(self.fc1(x))\n",
    "        x = self.relu4(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class MyModel(nn.Module):\n",
    "    def __init__(self, config: SimpleConfig) -> None:\n",
    "        super().__init__()\n",
    "        self.model = SimpleCNN()\n",
    "        self.loss = nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, x, labels, custom_input=None):\n",
    "        # custom_input is for demo purposes only, defined in the dataset wrapper\n",
    "        out = self.model(x)\n",
    "        loss = self.loss(out, labels)\n",
    "        if labels is not None:\n",
    "            loss = self.loss(out, labels)\n",
    "\n",
    "        out = out.argmax(dim=-1)\n",
    "        out = out.reshape(-1,1)\n",
    "        labels = labels.reshape(-1,1)\n",
    "        \n",
    "        return {\"y_pred\": out, \"y_true\": labels}, loss\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Dataset detail: we create the training & validation data loaders from the MNIST dataset. Data preprocessing includes normalization and transformations to tensor:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5,), (0.5,))\n",
    "])\n",
    "\n",
    "trainset = torchvision.datasets.MNIST(root='./datasets', train=True, download=True, transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, num_workers=2)\n",
    "\n",
    "testset = torchvision.datasets.MNIST(root='./datasets', train=False, download=True, transform=transform)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False, num_workers=2)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- An evaluation function is definded here for Ablator to evaluate the model."
   ]
  },
  {
   "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",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Wrap model with model wrapper and launch experiment\n",
    "\n",
    "As a final step, use the model wrapper to wrap the main model (this will add boiler-plate codes on the training and evaluation of your model on your dataset; we also add accuracy as an evaluation metric to the evaluation step) and launch Ablator."
   ]
  },
  {
   "attachments": {},
   "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: SimpleRunConfig):\n",
    "        return trainloader\n",
    "\n",
    "    def make_dataloader_val(self, run_config: SimpleRunConfig):\n",
    "        return testloader\n",
    "\n",
    "    def evaluation_functions(self) -> Dict[str, Callable]:\n",
    "        return {\"accuracy_score\": my_accuracy}\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    wrapper = MyModelWrapper(model_class=MyModel)\n",
    "    shutil.rmtree(run_config.experiment_dir, ignore_errors=True)    # Remove previous experiment results if existed to avoid experiment existed error when launching the experiment\n",
    "    ablator = ProtoTrainer(\n",
    "        wrapper=wrapper,\n",
    "        run_config=run_config,\n",
    "    )\n",
    "    ablator.launch()\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For Jupyter notebooks, directly run the above codes in the notebook. You can so save the above code snippets in a Python script and run it with the following command:\n",
    "\n",
    "```shell\n",
    "python <your_script_name>.py\n",
    "```\n",
    "\n",
    "If Ablator is successfully launched, you should see information printed on the console!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Access the results\n",
    "\n",
    "The training process (experiment results: training, evaluation results) is recorded by Ablator and is saved in the experiment directory (`run_config.experiment_dir`). You can access the training results of the model by using the following codes:\n",
    "\n",
    "```shell\n",
    "cd /tmp/dir/\n",
    "cat results.json\n",
    "```\n",
    "\n",
    "You should see the training results from each epoch.\n",
    "\n",
    "You can also visualize the results by using Tensorboard in Jupyter Notebook:\n",
    "\n",
    "```python\n",
    "# Load the TensorBoard extension\n",
    "import tensorboard\n",
    "%load_ext tensorboard\n",
    "\n",
    "# Start TensorBoard\n",
    "%tensorboard --logdir /tmp/dir/dashboard/tensorboard\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "Ablator is far beyond what we show you in this tutorial. Please refer to the following chapters for more features and functionalies of Ablator!"
   ]
  }
 ],
 "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
}