Clarification on MetricsCallback in optimize_hyperparameters()

• PyTorch-Forecasting version: 0.9.0
• PyTorch version: 1.9.0
• Python version: 3.7.10
• Operating System: Win10

Expected behavior

I am tuning a TFT model with the function optimize_hyperparameters. Each trial returns my specified metric (in this case MAE) like so:

[I 2021-08-20 18:39:06,537] Trial 219 finished with value: 11.637566566467285 and parameters: {'gradient_clip_val': 0.9433828074072016, 'hidden_size': 11, 'dropout': 0.15422157867252945, 'hidden_continuous_size': 8, 'attention_head_size': 2, 'learning_rate': 0.08458313458926149}. Best is trial 219 with value: 11.637566566467285.

I expect this returned value 11.63 to be from the best epoch in that trial. The callback used in the function also saves only the best epoch for each trial

checkpoint_callback = pl.callbacks.ModelCheckpoint(
    dirpath=os.path.join(model_path, "trial_{}".format(trial.number)), filename="{epoch}", monitor="val_loss"
)

Actual behavior

However, the best trial value reported by optuna is actually the value from the last epoch of that trial and so depends a lot on max_epochs. For each trial this is actually a good thing, e.g. for the trial above the best epoch had an MAE of 11,512 so it was better.

But for the entire optimization it is bad, because optuna does not get the actual best value of each trial.

As I see it, it all depends on this function that provides optuna with the final value:

class MetricsCallback(Callback):
    """PyTorch Lightning metric callback."""

    def __init__(self):
        super().__init__()
        self.metrics = []

    def on_validation_end(self, trainer, pl_module):
        self.metrics.append(trainer.callback_metrics)

Maybe I do not understand these callbacks correctly. If so, please close this issue.

Also thanks a lot for this great package.

Start in pytorch_forecasting/models/temporal_fusion_transformer/tuning.py, reading MetricsCallback and optimize_hyperparameters together with the linked ModelCheckpoint setup. Reproduce the reported behavior and compare the value passed to Optuna with the best validation epoch. Done means the optimization result reflects the intended trial metric rather than only the final epoch, with the behavior verified against the issue’s example.