Optimize model hyperparameters using Optuna

# Import to be able to import python package from src
import sys
sys.path.insert(0, '../src')

import ontime as on
import pandas as pd

Prerequisite

Install Optuna within your project

!pip install optuna optuna-integration optuna-dashboard

Create the dataset

We will directly use the BenchmarkDataset class, so that we can quickly create splits and samples for evaluating the model.

from ontime.module.benchmarking import BenchmarkDataset
from sklearn.preprocessing import StandardScaler
from darts.dataprocessing.transformers import Scaler

# generate a random time series
ts = on.generators.random_walk().generate(start=pd.Timestamp('2019-01-01'), end=pd.Timestamp('2023-12-31'))

# create the dataset
dataset = BenchmarkDataset(
    ts,
    "Random series",
    input_length=120,
    target_length=48,
    gap=0,
    stride=48,
    scaler_type=StandardScaler # add normalization cause why not
)

# split dataset (we only need train and val)
train, val = dataset.get_train_val_split()
scaler = Scaler(dataset.scaler_type())
train = scaler.fit_transform(train)
val = scaler.transform(val)

Setup the model with its hyperparameters

The setup method defines the hyperparameters and their values that must be tested in during the optimization. This method is dependent of the model, and should therefore be changed for your use case.

from darts.models import TCNModel
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from optuna_integration import PyTorchLightningPruningCallback

# avoid issue with optuna https://github.com/optuna/optuna/issues/4689
import pytorch_lightning as pl

class OptunaPruning(PyTorchLightningPruningCallback, pl.Callback):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

def setup_model(trial):
    # define the model

    pl_trainer_kwargs = {
        "accelerator": "auto",
        "callbacks": [
            EarlyStopping(
                monitor="val_loss",
                patience=10,
                mode="min",
            ),
            OptunaPruning(
                trial,
                monitor="val_loss",
            ),
        ]
    }

    tcn_model = TCNModel(
        model_name="optuna_tcn_model",
        input_chunk_length=dataset.input_length,
        output_chunk_length=trial.suggest_int("output_chunk_length", 1, 48),
        kernel_size=trial.suggest_int("kernel_size", 2, 5),
        num_filters=trial.suggest_int("num_filters", 8, 64),
        num_layers=trial.suggest_int("num_layers", 1, 4),
        dropout=trial.suggest_float("dropout", 0.0, 0.5),
        weight_norm=trial.suggest_categorical("weight_norm", [True, False]),
        n_epochs=5,
        random_state=42,
        optimizer_kwargs={
            "lr": trial.suggest_float("lr", 5e-5, 1e-3, log=True)
        },
        pl_trainer_kwargs=pl_trainer_kwargs,
        batch_size=trial.suggest_categorical("batch_size", [16, 32, 64]),
        save_checkpoints=True,
        force_reset=True,
    )

    return on.Model(tcn_model)

Create objective method

The objective method defines the objective function to optimized. Therefore, it must return a metric. We use the BenchmarkEvaluator to compute the metric to optimized, as it is able to directly create the samples according to the defined dataset configuration.

from darts.metrics import mse
from ontime.module.benchmarking import BenchmarkEvaluator, BenchmarkMetric

def objective(trial, metric=mse):
    # setup the model
    tcn_model = setup_model(trial)

    # fit the model
    tcn_model.fit(train, val_series=val)

    # make predictions
    evaluator = BenchmarkEvaluator(dataset, [BenchmarkMetric("val_metric", metric)], on_val_ts=True)

    results = evaluator.evaluate(tcn_model, scaler=scaler, scaled_evaluation=True)

    return results["val_metric"]

Run the optimization

Once the optimization is run, you can access to the dashboard on your web browser using the command optuna-dashboard sqlite:///relative/path/to/file.db.

If you are using Visual Studio Code, you can install the Optuna Dashboard extension, right click on your db file and select “Open in Optuna Dashboard”. However, dashboard in VS Code is not as complete as the one opened in the browser.

import optuna

study = optuna.create_study(
    storage="sqlite:///db.sqlite3", # for dashboard
    study_name="tcn_optuna_tuto",
    load_if_exists=True,
    direction="minimize")
study.optimize(objective, n_trials=5)

print(f"Best trial: {study.best_trial.value}")
print(f"Best hyperparameters: {study.best_trial.params}")

Best trial: 0.16961007783457446
Best hyperparameters: {'output_chunk_length': 6, 'kernel_size': 4, 'num_filters': 61, 'num_layers': 2, 'dropout': 0.2106525223389889, 'weight_norm': True, 'lr': 0.0009837089623902606, 'batch_size': 32}