Adding Models to NeuralForecast
Tutorial on how to add new models to NeuralForecast
Prerequisites
This Guide assumes advanced familiarity with NeuralForecast.
We highly recommend reading first the Getting Started and the NeuralForecast Map tutorials!
Additionally, refer to the CONTRIBUTING guide for the basics how to contribute to NeuralForecast.
Introduction
This tutorial is aimed at contributors who want to add a new model to
the NeuralForecast library. The library’s existing modules handle
optimization, training, selection, and evaluation of deep learning
models. The core
class simplifies building entire pipelines, both for
industry and academia, on any dataset, with user-friendly methods such
as fit
and predict
.
Adding a new model to NeuralForecast is simpler than building a new PyTorch model from scratch. You only need to write the forward method.
It has the following additional advantages:
- Existing modules in NeuralForecast already implement the essential training and evaluating aspects for deep learning models.
- Integrated with PyTorch-Lightning and Tune libraries for efficient optimization and distributed computation.
- The
BaseModel
classes provide common optimization components, such as early stopping and learning rate schedulers. - Automatic performance tests are scheduled on Github to ensure quality standards.
- Users can easily compare the performance and computation of the new model with existing models.
- Opportunity for exposure to a large community of users and contributors.
Example: simplified MLP model
We will present the tutorial following an example on how to add a
simplified version of the current
MLP
model, which does not include exogenous covariates.
At a given timestamp , the
MLP
model will forecast the next values of the univariate target time,
, using as inputs the last historical values, given by
. The following figure presents a diagram of the model.
0. Preliminaries
Follow our tutorial on contributing here to set up your development environment.
Here is a short list of the most important steps:
- Create a fork of the
neuralforecast
library. - Clone the fork to your computer.
- Set an environment with the
neuralforecast
library, core dependencies, andnbdev
package to code your model in an interactive notebook.
1. Inherit the Base Class (BaseWindows
)
The library contains three types of base models: BaseWindows
,
BaseRecurrent
, and BaseMultivariate
. In this tutorial, we will focus
on the BaseWindows
class, which is the most common type of model in
the library, with models such as
NBEATS
,
NHITS
,
TFT
,
and
PatchTST
.
The main difference between the three types is the sampling procedure
and input batch for the forward
method, which determines the type of
model.
Important
If you want to add a
BaseRecurrent
orBaseMultivariate
model, please add an issue in our github.
a. Sampling process
During training, the base class receives a sample of time series of the
dataset from the
TimeSeriesLoader
module. The BaseWindows
models will sample individual windows of size
input_size+h
, starting from random timestamps.
b. BaseWindows
’ hyperparameters
Get familiar with the hyperparameters specified in the base class,
including h
(horizon), input_size
, and optimization hyperparameters
such as learning_rate
, max_steps
, among others. The following list
presents the hyperparameters related to the sampling of windows:
h
(h): number of future values to predict.input_size
(L): number of historic values to use as input for the model.batch_size
(bs): number of time series sampled by the loader during training.valid_batch_size
(v_bs): number of time series sampled by the loader during inference (validation and test).windows_batch_size
(w_bs): number of individual windows sampled during training (from the previous time series) to form the batch.inference_windows_batch_size
(i_bs): number of individual windows sampled during inference to form each batch. Used to control the GPU memory.
c. Input and Output batch shapes
The forward
method receives a batch of data in a dictionary with the
following keys:
insample_y
: historic values of the time series.insample_mask
: mask indicating the available values of the time series (1 if available, 0 if missing).futr_exog
: future exogenous covariates (if any).hist_exog
: historic exogenous covariates (if any).stat_exog
: static exogenous covariates (if any).
The following table presents the shape for each tensor:
tensor | BaseWindows |
---|---|
insample_y | (w_bs , L ) |
insample_mask | (w_bs , L ) |
futr_exog | (w_bs , L +h , n_f ) |
hist_exog | (w_bs , L , n_h ) |
stat_exog | (w_bs ,n_s ) |
The forward
function should return a single tensor with the forecasts
of the next h
timestamps for each window. Use the attributes of the
loss
class to automatically parse the output to the correct shape (see
the example below).
Tip
Since we are using
nbdev
, you can easily add prints to the code and see the shapes of the tensors during training.
d. BaseWindows
’ methods
The BaseWindows
class contains several common methods for all
windows-based models, simplifying the development of new models by
preventing code duplication. The most important methods of the class
are:
_create_windows
: parses the time series from theTimeSeriesLoader
into individual windows of sizeinput_size+h
._normalization
: normalizes each window based on thescaler
type._inv_normalization
: inverse normalization of the forecasts.training_step
: training step of the model, called by PyTorch-Lightning’sTrainer
class during training (fit
method).validation_step
: validation step of the model, called by PyTorch-Lightning’sTrainer
class during validation.predict_step
: prediction step of the model, called by PyTorch-Lightning’sTrainer
class during inference (predict
method).
2. Create the model file and class
Once familiar with the basics of the BaseWindows
class, the next step
is creating your particular model.
The main steps are:
- Create the file in the
nbs
folder (https://github.com/Nixtla/neuralforecast/tree/main/nbs). It should be namedmodels.YOUR_MODEL_NAME.ipynb
. - Add the header of the
nbdev
file. - Import libraries in the file.
- Define the
__init__
method with the model’s inherited and particular hyperparameters and instantiate the architecture. - Define the
forward
method, which recieves the input batch dictionary and returns the forecast.
a. Model class
First, add the following two cells on top of the nbdev
file.
#| default_exp models.mlp
Important
Change
mlp
to your model’s name, using lowercase and underscores. When you later runnbdev_export
, it will create aYOUR_MODEL.py
script in theneuralforecast/models/
directory.
#| hide
%load_ext autoreload
%autoreload 2
Next, add the dependencies of the model.
#| export
from typing import Optional
import torch
import torch.nn as nn
from neuralforecast.losses.pytorch import MAE
from neuralforecast.common._base_windows import BaseWindows
Tip
Don’t forget to add the
#| export
tag on this cell.
Next, create the class with the init
and forward
methods. The
following example shows the example for the simplified
MLP
model. We explain important details after the code.
#| export
class MLP(BaseWindows): # <<---- Inherits from BaseWindows
def __init__(self,
# Inhereted hyperparameters with no defaults
h,
input_size,
# Model specific hyperparameters
num_layers = 2,
hidden_size = 1024,
# Inhereted hyperparameters with defaults
exclude_insample_y = False,
loss = MAE(),
valid_loss = None,
max_steps: int = 1000,
learning_rate: float = 1e-3,
num_lr_decays: int = -1,
early_stop_patience_steps: int =-1,
val_check_steps: int = 100,
batch_size: int = 32,
valid_batch_size: Optional[int] = None,
windows_batch_size = 1024,
inference_windows_batch_size = -1,
step_size: int = 1,
scaler_type: str = 'identity',
random_seed: int = 1,
num_workers_loader: int = 0,
drop_last_loader: bool = False,
**trainer_kwargs):
# Inherit BaseWindows class
super(MLP, self).__init__(h=h,
input_size=input_size,
..., # <<--- Add all inhereted hyperparameters
random_seed=random_seed,
**trainer_kwargs)
# Architecture
self.num_layers = num_layers
self.hidden_size = hidden_size
# MultiLayer Perceptron
layers = [nn.Linear(in_features=input_size, out_features=hidden_size)]
layers += [nn.ReLU()]
for i in range(num_layers - 1):
layers += [nn.Linear(in_features=hidden_size, out_features=hidden_size)]
layers += [nn.ReLU()]
self.mlp = nn.ModuleList(layers)
# Adapter with Loss dependent dimensions
self.out = nn.Linear(in_features=hidden_size,
out_features=h * self.loss.outputsize_multiplier) ## <<--- Use outputsize_multiplier to adjust output size
def forward(self, windows_batch): # <<--- Receives windows_batch dictionary
# Parse windows_batch
insample_y = windows_batch['insample_y'].clone()
# MLP
y_pred = self.mlp(insample_y)
# Reshape and map to loss domain
y_pred = y_pred.reshape(batch_size, self.h, self.loss.outputsize_multiplier)
y_pred = self.loss.domain_map(y_pred)
return y_pred
Tip
- Don’t forget to add the
#| export
tag on each cell.- Larger architectures, such as Transformers, might require splitting the
forward
by using intermediate functions.
Important notes
The base class has many hyperparameters, and models must have default
values for all of them (except h
and input_size
). If you are unsure
of what default value to use, we recommend copying the default values
from existing models for most optimization and sampling hyperparameters.
You can change the default values later at any time.
The reshape
method at the end of the forward
step is used to adjust
the output shape. The loss
class contains an outputsize_multiplier
attribute to automatically adjust the output size of the forecast
depending on the loss
. For example, for the Multi-quantile loss
(MQLoss
),
the model needs to output each quantile for each horizon.
Finally, always include y_pred = self.loss.domain_map(y_pred)
at the
end of the forward
. This is necessary to map the output to the
domain (and shape) of the loss function. For example, if the loss
function is the
MAE
,
it maps the output of shape (batch_size, h, 1)
to (batch_size, h)
.
b. Tests and documentation
nbdev
allows for testing and documenting the model during the
development process. It allows users to iterate the development within
the notebook, testing the code in the same environment. Refer to
existing models, such as the complete MLP model
here.
These files already contain the tests, documentation, and usage examples
that were used during the development process.
c. Export the new model to the library with nbdev
Following the CONTRIBUTING guide, the next step is to export the new
model from the development notebook to the neuralforecast
folder with
the actual scripts.
To export the model, run nbdev_export
in your terminal. You should see
a new file with your model in the neuralforecast/models/
folder.
3. Core class and additional files
Finally, add the model to the core
class and additional files:
-
Manually add the model in the following init file.
-
Add the model to the
core
class, using thenbdev
file here:- Add the model to the initial model list:
from neuralforecast.models import ( GRU, LSTM, RNN, TCN, DilatedRNN, MLP, NHITS, NBEATS, NBEATSx, TFT, VanillaTransformer, Informer, Autoformer, FEDformer, StemGNN, PatchTST )
- Add the model to the
MODEL_FILENAME_DICT
dictionary (used for thesave
andload
functions).
4. Add the model to the documentation
It’s important to add the model to the necessary documentation pages so that everyone can find the documentation:
- Add the model to the model overview table.
- Add the model to the sidebar for the API reference.
- Add the model to mint.json.
5. Upload to GitHub
Congratulations! The model is ready to be used in the library following the steps above.
Follow our contributing guide’s final steps to upload the model to GitHub: here.
One of the maintainers will review the PR, request changes if necessary, and merge it into the library.
Quick Checklist
- Get familiar with the
BaseWindows
class hyperparameters and input/output shapes of theforward
method. - Create the notebook with your model class in the
nbs
folder:models.YOUR_MODEL_NAME.ipynb
- Add the header and import libraries.
- Implement
init
andforward
methods. - Export model with
nbdev_export
. - Add model to this init file.
- Add the model to the
core
class here. - Follow the CONTRIBUTING guide to create the PR to upload the model.