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. Open In Colab

​
1. Installing NeuralForecast

!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_iddsy
01.01949-01-31112.0
11.01949-02-28118.0
21.01949-03-31132.0
31.01949-04-30129.0
41.01949-05-31121.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

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_iddsNBEATSNHITSAutoMLP
01.01961-01-31446.882172447.219238454.914154
11.01961-02-28465.145813464.558014430.188446
21.01961-03-31469.978424474.637238458.478577
31.01961-04-30493.650665502.670349477.244507
41.01961-05-31537.569275559.405212522.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, two files are created and stored:
  • [model_name]_[suffix].ckpt: Pytorch Lightning checkpoint file with the model parameters and hyperparameters.
  • [model_name]_[suffix].pkl: Dictionary with configuration attributes.
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.

​
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_iddsNHITSNBEATSAutoMLP
01.01961-01-31447.219238446.882172454.914154
11.01961-02-28464.558014465.145813430.188446
21.01961-03-31474.637238469.978424458.478577
31.01961-04-30502.670349493.650665477.244507
41.01961-05-31559.405212537.569275522.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.