Save and Load Models

Saving and loading trained Deep Learning models has multiple valuable uses. These models are often costly to train; storing a pre-trained model can help reduce costs as it can be loaded and reused to forecast multiple times. Moreover, it enables Transfer learning capabilities, consisting of pre-training a flexible model on a large dataset and using it later on other data with little to no training. It is one of the most outstanding 🚀 achievements in Machine Learning 🧠 and has many practical applications. In this notebook we show an example on how to save and load NeuralForecast models. The two methods to consider are:
1. NeuralForecast.save: Saves models into disk, allows save dataset and config.
2. NeuralForecast.load: Loads models from a given path.

Important This Guide assumes basic knowledge on the NeuralForecast library. For a minimal example visit the Getting Started guide.

You can run these experiments using GPU with Google Colab.

1. Installing NeuralForecast

%%capture
!pip install neuralforecast

2. Loading AirPassengers Data

For this example we will use the classical AirPassenger Data set. Import the pre-processed AirPassenger from utils.

from neuralforecast.utils import AirPassengersDF

Y_df = AirPassengersDF
Y_df.head()

	unique_id	ds	y
0	1.0	1949-01-31	112.0
1	1.0	1949-02-28	118.0
2	1.0	1949-03-31	132.0
3	1.0	1949-04-30	129.0
4	1.0	1949-05-31	121.0

3. Model Training

Next, we instantiate and train three models: NBEATS, NHITS, and AutoMLP. The models with their hyperparameters are defined in the models list.

import logging

from ray import tune

from neuralforecast.core import NeuralForecast
from neuralforecast.auto import AutoMLP
from neuralforecast.models import NBEATS, NHITS

logging.getLogger('pytorch_lightning').setLevel(logging.ERROR)

horizon = 12
models = [NBEATS(input_size=2 * horizon, h=horizon, max_steps=50),
          NHITS(input_size=2 * horizon, h=horizon, max_steps=50),
          AutoMLP(# Ray tune explore config
                  config=dict(max_steps=100, # Operates with steps not epochs
                              input_size=tune.choice([3*horizon]),
                              learning_rate=tune.choice([1e-3])),
                  h=horizon,
                  num_samples=1, cpus=1)]

Seed set to 1
Seed set to 1

%%capture
nf = NeuralForecast(models=models, freq='ME')
nf.fit(df=Y_df)

Produce the forecasts with the predict method.

Y_hat_df = nf.predict()
Y_hat_df.head()

Predicting: |                                                                                                 …

Predicting: |                                                                                                 …

Predicting: |                                                                                                 …

	unique_id	ds	NBEATS	NHITS	AutoMLP
0	1.0	1961-01-31	446.882172	447.219238	454.914154
1	1.0	1961-02-28	465.145813	464.558014	430.188446
2	1.0	1961-03-31	469.978424	474.637238	458.478577
3	1.0	1961-04-30	493.650665	502.670349	477.244507
4	1.0	1961-05-31	537.569275	559.405212	522.252991

We plot the forecasts for each model.

from utilsforecast.plotting import plot_series

plot_series(Y_df, Y_hat_df)

4. Save models

To save all the trained models use the save method. This method will save both the hyperparameters and the learnable weights (parameters). The save method has the following inputs:

path: directory where models will be saved.
model_index: optional list to specify which models to save. For example, to only save the NHITS model use model_index=[2].
overwrite: boolean to overwrite existing files in path. When True, the method will only overwrite models with conflicting names.
save_dataset: boolean to save Dataset object with the dataset.

nf.save(path='./checkpoints/test_run/',
        model_index=None, 
        overwrite=True,
        save_dataset=True)

For each model a single file is created:

[model_name]_[suffix].safetensors: the learnable weights, with the hyperparameters stored in the file’s metadata header.

Two files describe the rest of the state:

configuration.json: the NeuralForecast configuration, including which class each checkpoint holds.
dataset.json and dataset.safetensors: the stored dataset, when save_dataset=True.

Where model_name corresponds to the name of the model in lowercase (eg. nhits). We use a numerical suffix to distinguish multiple models of each class. In this example the names will be automlp_0, nbeats_0, and nhits_0.

Important The Auto models will be stored as their base model. For example, the AutoMLP trained above is stored as an MLP model, with the best hyparparameters found during tuning.

A note on the artifact format

Directories saved by version 3.3.0 and later hold safetensors weights and JSON metadata, and they are loaded without executing any code they contain. Earlier versions used pickle, which executes arbitrary code contained in the artifact when it is loaded. Reading one of those directories therefore requires you to say so:

nf = NeuralForecast.load(path='./old_checkpoints/', allow_pickle=True)

Only do that for artifacts you trust. To convert an old directory once, instead of passing the flag every time:

python -m neuralforecast.migrate ./old_checkpoints/

This writes ./old_checkpoints_v2/, leaves the original untouched, and checks that the result loads with no pickle consent at all. It also takes a single checkpoint: python -m neuralforecast.migrate ./NHITS_0.ckpt. Two cases to be aware of:

A directory saved with save_dataset=True by an older version must be migrated. Its dataset.pkl stores tensors inside a plain pickle, and there is no way to read that safely. If you would rather not migrate, re-save with save_dataset=False and pass df to predict().
TimeLLM can no longer be saved or loaded. It resolves its llm argument through from_pretrained while being constructed, so an artifact could direct that fetch. Train and predict with it in the same process.

Loading from remote storage

load refuses non-local paths unless you opt in:

nf = NeuralForecast.load(path='s3://my-bucket/models/', trust_remote=True)

Anyone who can write that location chooses what gets loaded on the machine doing the loading, so this is worth being deliberate about.

Custom losses, optimizers and schedulers

Class names inside an artifact are resolved through a closed registry, so anything the library does not ship has to be registered before it can be saved:

from neuralforecast import register_loss, register_optimizer, register_lr_scheduler

register_loss(MyLoss)
register_optimizer(MyOptimizer)
register_lr_scheduler(MyScheduler)

Two things to know:

Registration has to run in the process that loads, too, not only the one that saves. The artifact stores the registered name; the class itself comes from your code.
Any nn.Module can be registered as a loss. One that does not inherit a neuralforecast loss base has no recorded constructor arguments, so they are recovered from the object and checked by rebuilding it when you save. If the rebuild would differ, saving fails rather than writing an artifact that loads as a different loss.

Registration is what makes a class saveable. load additionally accepts the same four arguments, to replace what an artifact stored — useful when the stored class is no longer importable, or when you want to swap it:

model = NHITS.load(path, loss=MyLoss(), optimizer=MyOptimizer)

That works for both formats, but it cannot rescue a model that failed to save: those errors appear before any file is written. load also accepts map_location, and rejects anything else rather than ignoring it.

What is not carried in an artifact

Some settings describe the run rather than the model, and are dropped on save or on load with a warning naming them. Set them on the loaded object if you want them back.

Callbacks are dropped on save, as in earlier versions. Early stopping is rebuilt from early_stop_patience_steps, so that keeps working. Re-attach others with model.trainer_kwargs['callbacks'] = [...].
A logger instance is dropped; logger=False is kept, since it is configuration rather than an object.
Trainer and DataLoader settings that belong to the machine — default_root_dir, strategy, num_nodes, profiler, num_workers, prefetch_factor, multiprocessing_context — are not restored from an artifact, because a file should not choose where checkpoints are written or how many processes are spawned on the host loading it.
Anything that needs pickle to serialize, such as a worker_init_fn in dataloader_kwargs or an lr_lambda in lr_scheduler_kwargs, is dropped on save with a warning. A model argument that cannot be encoded is an error instead, so a model never silently reloads without its loss.

5. Load models

Load the saved models with the load method, specifying the path, and use the new nf2 object to produce forecasts.

nf2 = NeuralForecast.load(path='./checkpoints/test_run/')
Y_hat_df2 = nf2.predict()
Y_hat_df2.head()

Seed set to 1
Seed set to 1
Seed set to 1

Predicting: |                                                                                                 …

Predicting: |                                                                                                 …

Predicting: |                                                                                                 …

	unique_id	ds	NHITS	NBEATS	AutoMLP
0	1.0	1961-01-31	447.219238	446.882172	454.914154
1	1.0	1961-02-28	464.558014	465.145813	430.188446
2	1.0	1961-03-31	474.637238	469.978424	458.478577
3	1.0	1961-04-30	502.670349	493.650665	477.244507
4	1.0	1961-05-31	559.405212	537.569275	522.252991

Finally, plot the forecasts to confirm they are identical to the original forecasts.

plot_series(Y_df, Y_hat_df2)

References

https://pytorch-lightning.readthedocs.io/en/stable/common/checkpointing_basic.html Oreshkin, B. N., Carpov, D., Chapados, N., & Bengio, Y. (2019). N-BEATS: Neural basis expansion analysis for interpretable time series forecasting. ICLR 2020 Cristian Challu, Kin G. Olivares, Boris N. Oreshkin, Federico Garza, Max Mergenthaler-Canseco, Artur Dubrawski (2021). N-HiTS: Neural Hierarchical Interpolation for Time Series Forecasting. Accepted at AAAI 2023.

Getting Started

Capabilities

Tutorials

Use cases

API Reference

Save and Load Models

1. Installing NeuralForecast

2. Loading AirPassengers Data

3. Model Training

4. Save models

A note on the artifact format

Loading from remote storage

Custom losses, optimizers and schedulers

What is not carried in an artifact

5. Load models

References

​1. Installing NeuralForecast

​2. Loading AirPassengers Data

​3. Model Training

​4. Save models

​A note on the artifact format

​Loading from remote storage

​Custom losses, optimizers and schedulers

​What is not carried in an artifact

​5. Load models

​References

1. Installing NeuralForecast

2. Loading AirPassengers Data

3. Model Training

4. Save models

A note on the artifact format

Loading from remote storage

Custom losses, optimizers and schedulers

What is not carried in an artifact

5. Load models

References