[BUG] Torch model with checkpoints is not saved independent of checkpoint

[tharwan]

Describe the bug
When fitting a torch based model with checkpoints and then using save_model and load_model to persist it will result in an error during loading because the model will contain a reference to the checkpoints and tries to load the last checkpoint during prediction.

To Reproduce

from darts.dataprocessing.transformers import Scaler
from darts.models import RNNModel
from darts.datasets import AirPassengersDataset
import pandas as pd

# Read data:
series = AirPassengersDataset().load()

# Create training and validation sets:
train, val = series.split_after(pd.Timestamp("19590101"))

# Normalize the time series (note: we avoid fitting the transformer on the validation set)
transformer = Scaler()
train_transformed = transformer.fit_transform(train)
val_transformed = transformer.transform(val)
series_transformed = transformer.transform(series)

model_name = "Air_RNN"
my_model = RNNModel(
    n_epochs=20,
    model_name=model_name,
    random_state=42,
    input_chunk_length=14,
    training_length=20,
    save_checkpoints=True,
)

my_model.fit(
    train_transformed,
    val_series=val_transformed,
)

best_model = RNNModel.load_from_checkpoint(model_name, best=True)
best_model.save_model("best_model.pth.tar")

# delete logs
!rm -rf darts_logs

# reload model
best_model = RNNModel.load_model("best_model.pth.tar")

# this will fail because best_model.load_ckpt_path != None but the checkpoints are deleted
best_model.predict(1, val_transformed)
# Raises FileNotFoundError: Checkpoint at ....ckpt not found. Aborting training.

Expected behavior
loaded model predicts without needing the checkpoint.

System (please complete the following information):

• Python version: 3.7
• darts version 0.17.1

Additional context
For us this is somewhat critical that the model can not be loaded outside its training context. However the easy workaround seems to be to just set best_model.load_ckpt_path = None

From my limited perspective there is also no reason why the checkpoint should be needed for prediction in the first place. I can prepare a PR that just skips loading it if that helps.

[dennisbader]

Hi @tharwan.
You are right that this workflow fails:

• load from checkpoint -> manually save -> load manually saved model -> predict.

Why does it fail? When you load from checkpoint, the checkpointed model state is actually only loaded either when calling fit() or predict():

# behind the scenes when fitting
pl.Trainer.fit(model, ..., ckpt_path=my_checkpoint_path)

# behind the scenes when predicting directly from checkpointed model state
pl.Trainer.predict(model, ..., ckpt_path=my_checkpoint_path)

We don't fully load the model from checkpoint when calling load_from_checkpoint to cover two scenarios:

• user wants to predict from a given checkpoint
• run crashed during training and user wants to resume training from checkpoint. This requires loading the checkpoint at training time for current training state.

So when you save your model manually before calling fit or predict, the model state has not yet been loaded.
That's why in your case prediction fails when removing the checkpoints.

The intended way to use the models is:

• load from checkpoint -> resume training with fit() or predict directly -> (Optionally: manually save model) -> (Optionally: load manually saved model)

May I ask why you directly want to load the best model and then save it again?
The best model is already saved as a checkpoint which you can load outside the training context.

[tharwan]

regarding our workflow:
We integrate darts with mlflow. So we have a training job and then log a model to mlflow. Of course we would like to log the best model.

I did not see any difference on the predictions from the saved model with or without the checkpoint being loaded, so I am not sure why it is needed.

Also the API of darts seems highly confusing about this issue. Consider for example that calling predict a second time without the checkpoint being available will not produce an error because then the load path is set to None.

I fully understand that the checkpoint might be needed to continue training.

[dennisbader]

Hmm I think you're right. Now that I look at it again, we actually do load the model state from checkpoint already when calling load_from_checkpoint(). Which is why you get the same prediction results and why we don't need the checkpoint for prediction.

I will fix this.

[tharwan]

Great thank you!