Implementation:Sktime Pytorch forecasting BaseModel V2

Knowledge Sources	Sktime_Pytorch_forecasting
Domains	Time_Series, Forecasting, Deep_Learning
Last Updated	2026-02-08 08:00 GMT

Overview

The V2 base model class for time series forecasting that extends LightningModule with a complete training, validation, testing, and prediction loop.

Description

BaseModel is an experimental V2 base class designed as the foundation for all forecasting models in pytorch-forecasting. It extends LightningModule and provides a standardized training pipeline including training_step, validation_step, test_step, and predict_step methods. The class manages loss computation, optimizer configuration (supporting "adam" and "sgd" strings or custom Optimizer instances), and learning rate scheduling (supporting "reduce_lr_on_plateau" and "step_lr").

The model also provides predict as a high-level prediction interface that uses a PredictCallback and Lightning Trainer to generate forecasts in "prediction", "quantiles", or "raw" modes. Conversion from raw outputs to point or quantile forecasts is delegated to the configured loss metric's to_prediction and to_quantiles methods.

Usage

Use this class as the base for implementing new forecasting models in the V2 architecture. Subclasses must implement the forward method. This class is experimental and intended for beta testing ahead of the v2.0.0 release.

Code Reference

Source Location

• Repository: Sktime_Pytorch_forecasting
• File: pytorch_forecasting/models/base/_base_model_v2.py
• Lines: 1-370

Signature

class BaseModel(LightningModule):
    def __init__(
        self,
        loss: Metric,
        logging_metrics: list[nn.Module] | None = None,
        optimizer: Optimizer | str | None = "adam",
        optimizer_params: dict | None = None,
        lr_scheduler: str | None = None,
        lr_scheduler_params: dict | None = None,
    )
    def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]
    def predict(
        self,
        dataloader: DataLoader,
        mode: str = "prediction",
        return_info: list[str] | None = None,
        mode_kwargs: dict[str, Any] = None,
        trainer_kwargs: dict[str, Any] = None,
    ) -> dict[str, torch.Tensor]
    def to_prediction(self, out: dict[str, Any], **kwargs) -> torch.Tensor
    def to_quantiles(self, out: dict[str, Any], **kwargs) -> torch.Tensor
    def training_step(self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int) -> STEP_OUTPUT
    def validation_step(self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int) -> STEP_OUTPUT
    def test_step(self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int) -> STEP_OUTPUT
    def predict_step(
        self,
        batch: tuple[dict[str, torch.Tensor]],
        batch_idx: int,
        dataloader_idx: int = 0,
    ) -> torch.Tensor
    def configure_optimizers(self) -> dict
    def log_metrics(self, y_hat: torch.Tensor, y: torch.Tensor, prefix: str = "val") -> None

Import

from pytorch_forecasting.models.base._base_model_v2 import BaseModel

I/O Contract

Inputs

Name	Type	Required	Description
loss	Metric	Yes	Loss function for training, must be a descendant of pytorch_forecasting.metrics.Metric
logging_metrics	list[nn.Module] or None	No	Additional metrics to log during training/validation/testing (default: empty list)
optimizer	Optimizer or str or None	No	Optimizer instance or name string ("adam", "sgd"). Defaults to "adam"
optimizer_params	dict or None	No	Parameters passed to the optimizer constructor
lr_scheduler	str or None	No	Learning rate scheduler name ("reduce_lr_on_plateau", "step_lr"). Defaults to None
lr_scheduler_params	dict or None	No	Parameters passed to the learning rate scheduler constructor
x	dict[str, torch.Tensor]	Yes	Input dictionary of tensors for forward pass
dataloader	DataLoader	Yes	DataLoader for predict method
mode	str	No	Prediction mode: "prediction", "quantiles", or "raw". Defaults to "prediction"

Outputs

Name	Type	Description
forward output	dict[str, torch.Tensor]	Dictionary containing output tensors including "prediction" key
training_step output	dict	Dictionary containing "loss" key with scalar loss tensor
validation_step output	dict	Dictionary containing "val_loss" key with scalar loss tensor
test_step output	dict	Dictionary containing "test_loss" key with scalar loss tensor
predict output	dict[str, torch.Tensor]	Dictionary of prediction results from PredictCallback

Usage Examples

from pytorch_forecasting.models.base._base_model_v2 import BaseModel
from pytorch_forecasting.metrics.point import MAE
import torch
import torch.nn as nn

class SimpleForecaster(BaseModel):
    def __init__(self, input_size, output_size, **kwargs):
        super().__init__(**kwargs)
        self.linear = nn.Linear(input_size, output_size)

    def forward(self, x):
        prediction = self.linear(x["encoder_cont"])
        return {"prediction": prediction}

# Create model with MAE loss and Adam optimizer
model = SimpleForecaster(
    input_size=10,
    output_size=5,
    loss=MAE(),
    optimizer="adam",
    optimizer_params={"lr": 1e-3},
    lr_scheduler="reduce_lr_on_plateau",
    lr_scheduler_params={"patience": 5},
)

# Generate predictions
results = model.predict(
    dataloader=test_dataloader,
    mode="prediction",
    trainer_kwargs={"accelerator": "auto"},
)

Related Pages

• Principle:Sktime_Pytorch_forecasting_V2_Model_Base