Package (P) Layer v2#

Warning

Please note that the v2 modules are currently in active-development and is in beta right now, so please use this API with caution. See complete documentation for v2 API here and stable v1 documentation here.

The Package Layer is the defining feature of the V2 user experience. While the underlying D1, D2, and Model layers are strictly decoupled to maintain separation of concerns, the Package layer acts as the Orchestrator.

It wraps the PyTorch Lightning boilerplate and provides a streamlined, scikit-learn-like interface. By passing standard Python dictionaries for configuration, the Package layer automatically handles the instantiation of the DataModule, extracts the necessary metadata, initializes the underlying forecasting model, and manages the training loop.

The Configuration Driven Workflow#

To use a V2 Package, you do not need to manually instantiate the DataModule or the Model. Instead, you provide three configuration dictionaries to the Package wrapper:

  1. datamodule_cfg: Parameters for the D2 layer (e.g., batch size, scalers, encoders, and maximum sequence lengths).

  2. model_cfg: Parameters for the core neural network (e.g., hidden sizes, dropout, loss metrics, and optimizers).

  3. trainer_cfg: Parameters for the PyTorch Lightning Trainer (e.g., max epochs, accelerator type, and logging configurations).

sklearn like Methods#

The Package layer exposes two primary methods for interacting with the model lifecycle:

  • fit(): Initiates the training process. You can pass a raw D1 TimeSeries dataset (which the package will internally wrap in a D2 DataModule) or a pre-configured D2 LightningDataModule.

  • predict(): Generates forecasts. Similar to fit(), this accepts a D1 dataset, a D2 DataModule, or even a standard PyTorch DataLoader. It also accepts a return_info parameter to easily append identifying columns (like time indices and series IDs) alongside the predictions.

All package classes inherit from Base_pkg

class pytorch_forecasting.base._base_pkg.Base_pkg(model_cfg: dict[str, Any] | str | Path | None = None, trainer_cfg: dict[str, Any] | str | Path | None = None, datamodule_cfg: dict[str, Any] | str | Path | None = None, ckpt_path: str | Path | None = None)[source]

Base model package class acting as a high-level wrapper for the Lightning workflow.

This class simplifies the user experience by managing model, datamodule, and trainer configurations, and providing streamlined fit and predict methods.

Parameters:

  • model_cfg (dict, optional) – Model configs for the initialisation of the model. Required if not loading from a checkpoint. Defaults to {}.

  • trainer_cfg (dict, optional) – Configs to initialise lightning.Trainer. Defaults to {}.

  • datamodule_cfg (Union[dict, str, Path], optional) –

    Configs to initialise a LightningDataModule.

    • If dict, the keys and values are used as configuration parameters.

    • If str or Path, it should be a path to a .pkl file containing the serialized configuration dictionary. Required for reproducibility when loading a model for inference. Defaults to {}.

  • ckpt_path (Union[str, Path], optional) – Path to the checkpoint from which to load the model. If provided, model_cfg is ignored. Defaults to None.

Construct BaseObject.

Code Example#

Here is how the configuration dictionaries and lifecycle methods come together using the Temporal Fusion Transformer package (TFT_pkg_v2):

from pytorch_forecasting.models.temporal_fusion_transformer._tft_pkg_v2 import TFT_pkg_v2
from pytorch_forecasting.metrics import MAE, SMAPE
from pytorch_forecasting.data.encoders import NaNLabelEncoder, TorchNormalizer
from sklearn.preprocessing import StandardScaler

# Define Configurations
datamodule_cfg = dict(
    max_encoder_length=30,
    max_prediction_length=1,
    batch_size=32,
)

model_cfg = dict(
    loss=MAE(),
    logging_metrics=[MAE(), SMAPE()],
    optimizer="adam",
    optimizer_params={"lr": 1e-3},
    hidden_size=64,
    num_layers=2,
)

trainer_cfg = dict(
    max_epochs=5,
    accelerator="auto",
    devices=1,
)

# Initialize the Package
model_pkg = TFT_pkg_v2(
    model_cfg=model_cfg,
    trainer_cfg=trainer_cfg,
    datamodule_cfg=datamodule_cfg,
)

# Train the model
# (Assuming `dataset` is a previously defined D1 TimeSeries object)
model_pkg.fit(dataset)

# Generate predictions
preds = model_pkg.predict(dataset, return_info=["index", "x", "y"])

API Reference#

See the detailed API documentation for the available V2 Package classes below:

models.temporal_fusion_transformer._tft_pkg_v2.TFT_pkg_v2([...])

TFT package container.

models.dlinear._dlinear_pkg_v2.DLinear_pkg_v2([...])

DLinear package container.

models.samformer._samformer_v2_pkg.Samformer_pkg_v2([...])

Samformer package container.

models.tide._tide_dsipts._tide_v2_pkg.TIDE_pkg_v2([...])

TIDE package container.

models.timexer._timexer_pkg_v2.TimeXer_pkg_v2([...])

TimeXer metadata container.