Skip to content

Running the Benchmark

Getting Started

We provide scripts in the ./scripts directory for pretraining and to run the benchmark tasks (zero-shot STLF and transfer learning), either with our provided baselines or your own model.

PyTorch checkpoint files for our trained models are available for download as a single tar file here or as individual files on S3 here.

Our benchmark assumes each model takes as input a dictionary of torch tensors with the following keys:

    'load': torch.Tensor,               # (batch_size, seq_len, 1)
    'building_type': torch.LongTensor,  # (batch_size, seq_len, 1)
    'day_of_year': torch.FloatTensor,   # (batch_size, seq_len, 1)
    'hour_of_day': torch.FloatTensor,   # (batch_size, seq_len, 1)
    'day_of_week': torch.FloatTensor,   # (batch_size, seq_len, 1)
    'latitude': torch.FloatTensor,      # (batch_size, seq_len, 1)
    'longitude': torch.FloatTensor,     # (batch_size, seq_len, 1)

To use these scripts with your model you'll need to register your model with our platform.

Registering your model

Please see this step-by-step tutorial for a Jupyter Notebook version of the following instructions.

Make sure to have installed the benchmark in editable mode: pip install -e .[benchmark]

  1. Create a file called with your model's implementation, and make your model a subclass of the base model in ./buildings_bench/models/ Make sure to implement the abstract methods: forward, loss, load_from_checkpoint, predict, unfreeze_and_get_parameters_for_finetuning.
  2. Place this file under ./buildings_bench/models/
  3. Import your model class and add your model's name to the model_registry dictionary in ./buildings_bench/models/
  4. Create a TOML config file under ./buildings_bench/configs/your_model.toml with each keyword argument your model expects in its constructor (i.e., the hyperparameters for your model) and any additional args for the script you want to run.

The TOML config file should look something like this:

# your model's keyword arguments

# override any of the default pretraining argparse args here

# override any of the default zero_shot argparse args here

# override any of the default transfer_learning argparse args here
See ./buildings_bench/configs/TransformerWithTokenizer-S.toml for an example.


Without SLURM

The script is implemented with PyTorch DistributedDataParallel so it must be launched with torchrun from the command line and the argument --disable_slurm must be passed. See ./scripts/ for an example.


export WORLD_SIZE=1

torchrun \
    --nnodes=1 \
    --nproc_per_node=$NUM_GPUS \
    --rdzv-backend=c10d \
    --rdzv-endpoint=localhost:0 \
    scripts/ --model TransformerWithGaussian-S --disable_slurm

The argument --disable_slurm is not needed if you are running this script on a Slurm cluster as a batch job.

This script will automatically log outputs to wandb if the environment variables WANDB_ENTITY and WANDB_PROJECT are set. Otherwise, pass the argument --disable_wandb to disable logging to wandb.


To launch pretraining as a SLURM batch job:

export MASTER_PORT=$(expr 10000 + $(echo -n $SLURM_JOBID | tail -c 4))

master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_ADDR=$master_addr

srun python3 scripts/ \
        --model TransformerWithGaussian-S

Zero-shot STLF

This script scripts/ and the script for transfer learning scripts/ do not use DistributedDataParallel so they can be run without torchrun.

python3 scripts/ --model TransformerWithGaussian-S --checkpoint /path/to/

Transfer Learning for STLF

python3 scripts/ --model TransformerWithGaussian-S --checkpoint /path/to/