tutorials/registering_your_model_with_the_benchmark.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# registering_your_model_with_the_benchmark\n",
    "\n",
    "A step-by-step tutorial for how to register a new model with the benchmark.\n",
    "\n",
    "This tutorial assumes you have followed the [installation instructions](https://nrel.github.io/BuildingsBench/#installation) for BuildingsBench."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from buildings_bench.models.base_model import BaseModel\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Model Implementation\n",
    "\n",
    "Let's create an example PyTorch model that is an MLP performing direct regression on the 24 hour forecast."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyModel(BaseModel):\n",
    "\n",
    "    def __init__(self, \n",
    "                 hidden_size,\n",
    "                 context_len=168,\n",
    "                 pred_len=24,\n",
    "                 continuous_loads=True):\n",
    "        \"\"\"Init method for MyModel.\n",
    "        \n",
    "        Args:\n",
    "            hidden_size (int): size of hidden layer\n",
    "            context_len (int): length of context window\n",
    "            pred_len (int): length of prediction window\n",
    "            continuous_loads (bool): whether this model uses continuous load values\n",
    "        \"\"\"\n",
    "        super().__init__(context_len, pred_len, continuous_loads)\n",
    "\n",
    "        # Our model will be a simple MLP with two hidden layers\n",
    "        # and will output two values (mean and std dev) for each hour.\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(context_len, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_size, pred_len*2)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        `x` is a dictionary with the following keys:\n",
    "\n",
    "        ```\n",
    "        'load': torch.Tensor,               # (batch_size, seq_len, 1)\n",
    "        'building_type': torch.LongTensor,  # (batch_size, seq_len, 1)\n",
    "        'day_of_year': torch.FloatTensor,   # (batch_size, seq_len, 1)\n",
    "        'hour_of_day': torch.FloatTensor,   # (batch_size, seq_len, 1)\n",
    "        'day_of_week': torch.FloatTensor,   # (batch_size, seq_len, 1)\n",
    "        'latitude': torch.FloatTensor,      # (batch_size, seq_len, 1)\n",
    "        'longitude': torch.FloatTensor,     # (batch_size, seq_len, 1)\n",
    "        ```\n",
    "\n",
    "        This model only uses the 'load'.\n",
    "        \"\"\"\n",
    "        # (batch_size, self.context_len)\n",
    "        x = x['load'][:, :self.context_len, 0]\n",
    "        out = self.mlp(x)  # (batch_size, self.pred_len*2)\n",
    "        return out.view(-1, self.pred_len, 2) # (batch_size, self.pred_len, 2)\n",
    "    \n",
    "    def loss(self, x, y):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            x (torch.Tensor): preds of shape (batch_size, seq_len, 2)\n",
    "            y (torch.Tensor): targets of shape (batch_size, seq_len, 1)\n",
    "        Returns:\n",
    "            loss (torch.Tensor): scalar loss\n",
    "        \"\"\"\n",
    "        return nn.functional.gaussian_nll_loss(x[:, :, 0].unsqueeze(2), y,\n",
    "                                       nn.functional.softplus(x[:, :, 1].unsqueeze(2))**2)\n",
    "    \n",
    "    def predict(self, x):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            x (Dict): dictionary of input tensors\n",
    "        Returns:\n",
    "            predictions (torch.Tensor): of shape (batch_size, pred_len, 1)\n",
    "            distribution_parameters (torch.Tensor): of shape (batch_size, pred_len, -1)\n",
    "        \"\"\"\n",
    "        out = self.forward(x)\n",
    "        means = out[:, :, 0].unsqueeze(2)\n",
    "        stds = nn.functional.softplus(out[:, :, 1].unsqueeze(2))\n",
    "        return means, torch.cat([means, stds], dim=2)\n",
    "    \n",
    "    def unfreeze_and_get_parameters_for_finetuning(self):\n",
    "        \"\"\"For transfer learning.\"\"\"\n",
    "        return self.parameters()\n",
    "    \n",
    "    def load_from_checkpoint(self, checkpoint_path):\n",
    "        \"\"\"Describe how to load model from a checkpoint.\"\"\"\n",
    "        self.load_state_dict(torch.load(checkpoint_path)['model'])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Adding your model to the model registry\n",
    "\n",
    "Store this class definition in a file called my_model.py under `./buildings_bench/models/my_model.py`.\n",
    "At the top of `./buildings_bench/models/__init__.py`, add the following import statement:\n",
    "\n",
    "```python\n",
    "from buildings_bench.models.my_model import MyModel\n",
    "```\n",
    "\n",
    "Then, add your model to the `model_registry` dictionary:\n",
    "\n",
    "```python\n",
    "model_registry = {\n",
    "    ...\n",
    "    'MyModel': MyModel,\n",
    "    ...\n",
    "}\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Creating a TOML config file\n",
    "\n",
    "Create a TOML config file under `./configs/MyModel.toml` with each keyword argument your model expects in its constructor (i.e., the hyperparameters for your model) and any argparse args you wish to override for the script you want to run.\n",
    "\n",
    "The example model's config file `MyModel.toml` should look something like this:\n",
    "\n",
    "```toml\n",
    "[model]\n",
    "hidden_size = 1024\n",
    "context_len = 168\n",
    "pred_len = 24\n",
    "continuous_loads = true\n",
    "\n",
    "[pretrain]\n",
    "batch_size = 256\n",
    "init_scale = 0.02\n",
    "warmup_steps = 10000\n",
    "lr = 0.00006\n",
    "scheduler_steps = 162760\n",
    "apply_scaler_transform = 'boxcox'\n",
    "\n",
    "[transfer_learning]\n",
    "apply_scaler_transform = 'boxcox'\n",
    "```\n",
    "\n",
    "We specified the model's hidden size, context length, prediction length, and whether the model should be trained on continuous loads or not. \n",
    "\n",
    "For pretraining, we specified the batch size, learning rate, and number of warmup steps. We also specified the initial scale for the weight initialization and the number of steps for the learning rate scheduler. Finally, we specified to apply a Box-Cox transform to the loads.\n",
    "\n",
    "For transfer learning, we will use the default values for the batch size and learning rate, but we specify we want to apply the same Box-Cox transform to the loads."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Running the scripts\n",
    "\n",
    "### Pretraining\n",
    "\n",
    "`python3 scripts/pretrain.py --model MyModel`\n",
    "\n",
    "This script is implemented with PyTorch `DistributedDataParallel`, so it can be launched with `torchrun`. See `./scripts/pretrain.sh` for an example.\n",
    "\n",
    "### Zero-shot STLF\n",
    "\n",
    "`python3 scripts/zero_shot.py --model MyModel --checkpoint /path/to/checkpoint.pt`\n",
    "\n",
    "### Transfer Learning for STLF\n",
    "\n",
    "`python3 scripts/transfer_learning_torch.py --model MyModel --checkpoint /path/to/checkpoint.pt`"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}