Skip to content

[ENH] Add forecaster support and remove _pkg class for new API #2434

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
phoeenniixx wants to merge 6 commits into
base: dm
Choose a base branch
from
Open

[ENH] Add forecaster support and remove _pkg class for new API #2434

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
6f21533
add forecaster refactor
phoeenniixx Sep 20, 2026
1b934cc
updates
phoeenniixx Sep 28, 2026
7fbc169
solve checkpointing failures
phoeenniixx Sep 28, 2026
e7624d8
update workflow target branch
phoeenniixx Sep 28, 2026
5c5cec7
rename BasePtforecaster to BasePtObject, reimplement check_is_fitted …
phoeenniixx Oct 4, 2026
bbedfa8
rename pkg files and bring all v2 pkg to forecaster
phoeenniixx Oct 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .github/workflows/test.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,9 @@ name: Test

on:
push:
branches: [v2-dev]
branches: [dm]

@phoeenniixx phoeenniixx Sep 28, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

need to change this once #2423 is merged!

pull_request:
branches: [v2-dev]
branches: [dm]

concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
Expand Down
2 changes: 2 additions & 0 deletions pytorch_forecasting/models/base/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
"""Base classes for pytorch-foercasting models."""

from pytorch_forecasting.models.base._base_forecaster import BaseForecaster
from pytorch_forecasting.models.base._base_model import (
AutoRegressiveBaseModel,
AutoRegressiveBaseModelWithCovariates,
Expand All @@ -19,6 +20,7 @@
"_BasePtForecasterV2",
"AutoRegressiveBaseModel",
"AutoRegressiveBaseModelWithCovariates",
"BaseForecaster",
"BaseModel",
"BaseModelWithCovariates",
"Prediction",
Expand Down
193 changes: 95 additions & 98 deletions pytorch_forecasting/base/_base_pkg.py → ...recasting/models/base/_base_forecaster.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from lightning import Trainer
from lightning.pytorch.callbacks import ModelCheckpoint
from lightning.pytorch.core.datamodule import LightningDataModule
from skbase.base import BaseEstimator
import torch
from torch.utils.data import DataLoader
import yaml
Expand All @@ -13,64 +14,53 @@
from pytorch_forecasting.models.base._base_object import _BasePtForecasterV2


class Base_pkg(_BasePtForecasterV2):
class BaseForecaster(_BasePtForecasterV2, BaseEstimator):
"""
Base model package class acting as a high-level wrapper for the Lightning workflow.
Base forecaster 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.
Model parameters are flattened into the concrete subclass's ``__init__``; the
model is built inside ``fit`` from the datamodule's ``metadata``. The class
manages model, datamodule and trainer, and provides 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 {}.

trainer : lightning.Trainer, optional
Trainer used by ``fit`` when none is passed there. Defaults to ``Trainer()``.
datamodule : LightningDataModule, optional
A configured but data-less datamodule; ``fit`` attaches the data with
``with_data``. Defaults to ``get_datamodule_cls()()``, or to the
``datamodule_cfg.pkl`` saved next to ``ckpt_path`` if loading.
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.
Path to the checkpoint from which to load the model. Defaults to None.
"""

def __init__(
self,
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,
trainer: Trainer | None = None,
datamodule: Any = None,
ckpt_path: str | Path | None = None,
):
self.ckpt_path = Path(ckpt_path) if ckpt_path else None
self.model_cfg = self._load_config(
model_cfg, ckpt_path=self.ckpt_path, auto_file_name="model_cfg.pkl"
)
print(self.model_cfg)
self.trainer = trainer
self.datamodule = datamodule
self.ckpt_path = ckpt_path
super().__init__()

self.datamodule_cfg = self._load_config(
datamodule_cfg,
ckpt_path=self.ckpt_path,
auto_file_name="datamodule_cfg.pkl",
self._loaded_datamodule_cfg = self._load_config(
None, ckpt_path=self.ckpt_path, auto_file_name="datamodule_cfg.pkl"
)
self.trainer_cfg = self._load_config(trainer_cfg)
self.metadata = self._load_config(
None, ckpt_path=self.ckpt_path, auto_file_name="metadata.pkl"
)
self._loaded_model_cfg = self._load_config(
None, ckpt_path=self.ckpt_path, auto_file_name="model_cfg.pkl"
)

self.model = None
self.trainer = None
self.datamodule = None
self.trainer_ = None
self.datamodule_ = None
if self.ckpt_path:
print(self.metadata)
self._build_model(metadata=self.metadata, **self.model_cfg)
else:
self.model = None
self._build_model(metadata=self.metadata)

@staticmethod
def _load_config(
Expand All @@ -97,7 +87,6 @@ def _load_config(
raise FileNotFoundError(f"Configuration file not found: {path}")

suffix = path.suffix.lower()
print(suffix)

if suffix in [".yaml", ".yml"]:
with open(path) as f:
Expand All @@ -121,6 +110,25 @@ def get_datamodule_cls(cls):
"""Get the underlying DataModule class."""
raise NotImplementedError("Subclasses must implement `get_datamodule_cls`.")

def get_model_params(self) -> dict[str, Any]:
"""Return the kwargs for ``get_cls()``, excluding ``metadata``."""
raise NotImplementedError("Subclasses must implement `get_model_params`.")

@property
def model_cfg(self) -> dict[str, Any]:
"""Model kwargs, as passed to ``get_cls()``."""
return self.get_model_params()

@property
def datamodule_cfg(self) -> dict[str, Any]:
"""Constructor kwargs of the datamodule, fitted one preferred."""
dm = self.datamodule_ if self.datamodule_ is not None else self.datamodule
if dm is None:
return self._loaded_datamodule_cfg
if hasattr(dm, "_init_kwargs"):
return dm._init_kwargs
return dict(getattr(dm, "hparams", {}))

@classmethod
def get_test_dataset_from(cls, **kwargs):
"""
Expand All @@ -140,43 +148,39 @@ def get_test_dataset_from(cls, **kwargs):
"predict": datasets_info["validation_dataset"],
}

def _build_model(self, metadata: dict, **kwargs):
def _build_model(self, metadata: dict):
"""Instantiates the model, either from a checkpoint or from config."""
model_cls = self.get_cls()
if self.ckpt_path:

@CloseChoice CloseChoice Oct 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here is a problem, this only updates self.model but other properites are missed when reloading from checkpoint.

E.g.

"""Existing checkpoint limitation: reload loses fitted preprocessing state."""

import warnings
from pathlib import Path
from tempfile import TemporaryDirectory

from lightning.pytorch import Trainer
import numpy as np
import pandas as pd
import torch

from pytorch_forecasting.data import TimeSeries
from pytorch_forecasting.data.data_module import EncoderDecoderTimeSeriesDataModule
from pytorch_forecasting.models.mlp import DecoderMLPForecaster

warnings.filterwarnings("ignore", message=".*experimental.*")
torch.manual_seed(0)

frame = pd.DataFrame({
    "group": np.repeat(np.arange(10), 20),
    "time": np.tile(np.arange(20), 10),
    "target": 100.0 + np.tile(np.arange(20), 10),
    "feature": np.tile(np.arange(20), 10).astype(float),
})
data = TimeSeries(
    frame, time="time", target="target", group=["group"],
    num=["feature"], known=["feature"],
)

TRAINER_KWARGS = dict(
    accelerator="cpu", max_epochs=1, limit_train_batches=1,
    limit_val_batches=1, logger=False, enable_progress_bar=False,
    enable_model_summary=False, num_sanity_val_steps=0,
)

from sklearn.preprocessing import StandardScaler

with TemporaryDirectory(prefix="forecaster-review-") as directory:
    original = DecoderMLPForecaster(
        hidden_size=4, n_hidden_layers=1,
        datamodule=EncoderDecoderTimeSeriesDataModule(
            max_encoder_length=4, max_prediction_length=3,
            target_normalizer="auto", scalers={"feature": StandardScaler()},
        ),
    )
    checkpoint = original.fit(
        data, trainer=Trainer(**TRAINER_KWARGS, default_root_dir=directory),
        ckpt_dir=directory,
    )
    assert checkpoint.is_file()
    before_loader = original._load_dataloader(data)
    loaded = DecoderMLPForecaster(ckpt_path=checkpoint)
    after_loader = loaded._load_dataloader(data)
    breakpoint()
    before_dm = before_loader.dataset.data_module
    after_dm = after_loader.dataset.data_module
    before_x = next(iter(before_loader))[0]
    after_x = next(iter(after_loader))[0]
    print("Target fitted before / after reload:",
          before_dm._target_normalizer_fitted, after_dm._target_normalizer_fitted)
    print("Features fitted before / after reload:",
          before_dm._feature_scalers_fitted, after_dm._feature_scalers_fitted)
    for key in ("target_past", "decoder_cont"):
        print(f"{key} mean before / after reload:",
              before_x[key].mean().item(), after_x[key].mean().item())
        assert not torch.allclose(before_x[key], after_x[key])
    assert before_dm._target_normalizer_fitted and before_dm._feature_scalers_fitted
    assert not after_dm._target_normalizer_fitted and not after_dm._feature_scalers_fitted
    print("BUG reproduced: checkpoint reload changes preprocessing of identical input data.")

self.model = model_cls.load_from_checkpoint(

@CloseChoice CloseChoice Oct 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I actually think there is quite something wrong with the serialization of this.

"""PR #2434: refitting a loaded forecaster saves incompatible model config."""

import warnings
from pathlib import Path
from tempfile import TemporaryDirectory

from lightning.pytorch import Trainer
import numpy as np
import pandas as pd
import torch

from pytorch_forecasting.data import TimeSeries
from pytorch_forecasting.data.data_module import EncoderDecoderTimeSeriesDataModule
from pytorch_forecasting.models.mlp import DecoderMLPForecaster

warnings.filterwarnings("ignore", message=".*experimental.*")
torch.manual_seed(0)

frame = pd.DataFrame({
    "group": np.repeat(np.arange(10), 20),
    "time": np.tile(np.arange(20), 10),
    "target": 100.0 + np.tile(np.arange(20), 10),
    "feature": np.tile(np.arange(20), 10).astype(float),
})
data = TimeSeries(
    frame, time="time", target="target", group=["group"],
    num=["feature"], known=["feature"],
)

TRAINER_KWARGS = dict(
    accelerator="cpu", max_epochs=1, limit_train_batches=1,
    limit_val_batches=1, logger=False, enable_progress_bar=False,
    enable_model_summary=False, num_sanity_val_steps=0,
)

with TemporaryDirectory(prefix="forecaster-review-") as directory:
    root = Path(directory)
    original = DecoderMLPForecaster(
        hidden_size=4, n_hidden_layers=1,
        datamodule=EncoderDecoderTimeSeriesDataModule(
            max_encoder_length=4, max_prediction_length=3,
        ),
    )
    first_path = original.fit(
        data, trainer=Trainer(**TRAINER_KWARGS, default_root_dir=root),
        ckpt_dir=root / "first",
    )
    assert first_path.is_file()
    loaded = DecoderMLPForecaster(ckpt_path=first_path)
    print("Loaded wrapper hidden_size:", loaded.hidden_size)
    print("Loaded model hidden_size:", loaded.model.hidden_size)
    assert loaded.hidden_size != loaded.model.hidden_size
    second_path = loaded.fit(
        data, trainer=Trainer(**TRAINER_KWARGS, default_root_dir=root),
        ckpt_dir=root / "second",
    )
    assert second_path.is_file()
    try:
        DecoderMLPForecaster(ckpt_path=second_path)
    except RuntimeError as error:
        assert "size mismatch" in str(error)
        print("BUG reproduced: reloading the second checkpoint fails:")
        print(error)
    else:
        raise AssertionError("Reload succeeded; the reported bug was not reproduced.")

Couldn't we simply store all init parameters as well and then set them again when loading as a first try to tackle serialization?

self.ckpt_path, metadata=metadata, **kwargs
self.ckpt_path, metadata=metadata, **self._loaded_model_cfg
)
elif self.model_cfg:
self.model = model_cls(**self.model_cfg, metadata=metadata)
else:
self.model = None
self.model = model_cls(**self.model_cfg, metadata=metadata)

def _build_datamodule(self, data: TimeSeries) -> LightningDataModule:
"""Constructs a DataModule from a D1 layer object."""
if not self.datamodule_cfg:
raise ValueError("`datamodule_cfg` must be provided to build a datamodule.")
datamodule_cls = self.get_datamodule_cls()
return datamodule_cls(data, **self.datamodule_cfg)

def _load_dataloader(
self, data: TimeSeries | LightningDataModule | DataLoader
) -> DataLoader:
"""Converts various data input types into a DataLoader for prediction."""
if isinstance(data, TimeSeries): # D1 Layer
dm = self._build_datamodule(data)
dm.setup(stage="predict")
return dm.predict_dataloader()
elif isinstance(data, LightningDataModule): # D2 Layer
data.setup(stage="predict")
return data.predict_dataloader()
elif isinstance(data, DataLoader):
return data
"""Attach ``data`` to the configured datamodule, with no fitted transforms."""
dm = self.datamodule
if dm is None:
dm = self.get_datamodule_cls()(**self._loaded_datamodule_cfg)
return dm.with_data(data)

def _resolve_trainer(self, trainer: Trainer | None) -> Trainer:
"""``fit``'s trainer > ``__init__``'s trainer > default ``Trainer()``."""
if trainer is not None:
return trainer
if self.trainer is not None:
return self.trainer
return Trainer()

def _load_dataloader(self, data: TimeSeries) -> DataLoader:
"""Build the prediction dataloader, with the transforms fitted in ``fit``."""
if self.is_fitted:
dm = self.datamodule_.with_data(data)
else:
raise TypeError(
f"Unsupported data type for prediction: {type(data).__name__}. "
"Expected TimeSeriesDataSet, LightningDataModule, or DataLoader."
)
dm = self._build_datamodule(data)
dm.setup(stage="predict")
return dm.predict_dataloader()

def _save_artifact(self, output_dir: Path):
"""Save all configuration artifacts."""
Expand All @@ -188,14 +192,14 @@ def _save_artifact(self, output_dir: Path):
with open(output_dir / "model_cfg.pkl", "wb") as f:
pickle.dump(self.model_cfg, f)

if self.datamodule is not None and hasattr(self.datamodule, "metadata"):
if self.datamodule_ is not None and hasattr(self.datamodule_, "metadata"):
with open(output_dir / "metadata.pkl", "wb") as f:
pickle.dump(self.datamodule.metadata, f)
pickle.dump(self.datamodule_.metadata, f)

def fit(
self,
data: TimeSeries | LightningDataModule,
# todo: we should create a base data_module for different data_modules
data: TimeSeries,
trainer: Trainer | None = None,
save_ckpt: bool = True,
ckpt_dir: str | Path = "checkpoints",
ckpt_kwargs: dict[str, Any] | None = None,
Expand All @@ -206,9 +210,11 @@ def fit(

Parameters
----------
data : Union[TimeSeries, LightningDataModule]
The data to fit on (D1 or D2 layer). This object is responsible
for providing both training and validation data.
data : TimeSeries
The data to fit on. The datamodule splits it into training and
validation data.
trainer : lightning.Trainer, optional
Overrides the trainer given to ``__init__`` for this call only.
save_ckpt : bool, default=True
If True, save the best model checkpoint and the `datamodule_cfg`.
ckpt_dir : Union[str, Path], default="checkpoints"
Expand All @@ -223,21 +229,14 @@ def fit(
Optional[Path]
The path to the best model checkpoint if `save_ckpt=True`, else None.
"""
if isinstance(data, TimeSeries):
self.datamodule = self._build_datamodule(data)
else:
self.datamodule = data
self.datamodule.setup(stage="fit")
self.datamodule_ = self._build_datamodule(data)

if self.model is None:
if not self.model_cfg:
raise RuntimeError(
"`model_cfg` must be provided to train from scratch."
)
metadata = self.datamodule.metadata
self._build_model(metadata)

callbacks = self.trainer_cfg.get("callbacks", []).copy()
# the model is built only here, because only now are its input
# shapes known - they come from the metadata of the data module
metadata = self.datamodule_.metadata
self._build_model(metadata)

self.trainer_ = self._resolve_trainer(trainer)
checkpoint_cb = None
if save_ckpt:
ckpt_dir = Path(ckpt_dir)
Expand All @@ -252,13 +251,10 @@ def fit(
if ckpt_kwargs:
default_ckpt_kwargs.update(ckpt_kwargs)
checkpoint_cb = ModelCheckpoint(**default_ckpt_kwargs)
callbacks.append(checkpoint_cb)
trainer_init_cfg = self.trainer_cfg.copy()
trainer_init_cfg.pop("callbacks", None)

self.trainer = Trainer(**trainer_init_cfg, callbacks=callbacks)
self.trainer_.callbacks.append(checkpoint_cb)

self.trainer.fit(self.model, datamodule=self.datamodule, **trainer_fit_kwargs)
self.trainer_.fit(self.model, datamodule=self.datamodule_, **trainer_fit_kwargs)
self._is_fitted = True
if save_ckpt and checkpoint_cb:
best_model_path = Path(checkpoint_cb.best_model_path)
self._save_artifact(best_model_path.parent)
Expand All @@ -268,7 +264,7 @@ def fit(

def predict(
self,
data: TimeSeries | LightningDataModule | DataLoader,
data: TimeSeries,
output_dir: str | Path | None = None,
**kwargs,
) -> dict[str, torch.Tensor] | None:
Expand All @@ -280,8 +276,9 @@ def predict(

Parameters
----------
data : Union[TimeSeries, LightningDataModule, DataLoader]
The data to predict on (D1, D2, or DataLoader).
data : TimeSeries
The data to predict on. If the datamodule has a target normalizer
or scalers, the data is scaled with the ones fitted in ``fit``.
**kwargs :
Additional keyword arguments passed directly to the model's ``.predict()``
method. This includes `mode`, `return_info`, `output_dir`, and any
Expand All @@ -295,7 +292,7 @@ def predict(
"""
if self.model is None:
raise RuntimeError(
"Model is not initialized. Provide `model_cfg` or `ckpt_path`."
"Model is not initialized. Call `fit` or pass `ckpt_path`."
)

dataloader = self._load_dataloader(data)
Expand Down
4 changes: 2 additions & 2 deletions pytorch_forecasting/models/dlinear/_dlinear_pkg_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,10 @@
Packages container for DLinear model.
"""

from pytorch_forecasting.base._base_pkg import Base_pkg
from pytorch_forecasting.models.base._base_forecaster import BaseForecaster


class DLinear_pkg_v2(Base_pkg):
class DLinear_pkg_v2(BaseForecaster):
"""DLinear package container."""

_tags = {
Expand Down
4 changes: 2 additions & 2 deletions pytorch_forecasting/models/frets/_frets_pkg_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
"""FreTS v2 package container."""

from pytorch_forecasting.base._base_pkg import Base_pkg
from pytorch_forecasting.models.base._base_forecaster import BaseForecaster


class FreTS_pkg_v2(Base_pkg):
class FreTS_pkg_v2(BaseForecaster):
"""FreTS v2 package container."""

_tags = {
Expand Down
4 changes: 2 additions & 2 deletions pytorch_forecasting/models/mlp/_decodermlp_pkg_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
"""DecoderMLP v2 package container."""

from pytorch_forecasting.base._base_pkg import Base_pkg
from pytorch_forecasting.models.base._base_forecaster import BaseForecaster


class DecoderMLP_pkg_v2(Base_pkg):
class DecoderMLP_pkg_v2(BaseForecaster):
"""DecoderMLP v2 package container."""

_tags = {
Expand Down
4 changes: 2 additions & 2 deletions pytorch_forecasting/models/patch_tst/_patch_tst_pkg_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
"""PatchTST package container for V2."""

from pytorch_forecasting.base._base_pkg import Base_pkg
from pytorch_forecasting.models.base._base_forecaster import BaseForecaster


class PatchTST_pkg_v2(Base_pkg):
class PatchTST_pkg_v2(BaseForecaster):
"""PatchTST package container for V2."""

_tags = {
Expand Down
6 changes: 4 additions & 2 deletions pytorch_forecasting/models/samformer/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
"""

from pytorch_forecasting.models.samformer._samformer_v2 import Samformer
from pytorch_forecasting.models.samformer._samformer_v2_pkg import Samformer_pkg_v2
from pytorch_forecasting.models.samformer._samformer_v2_pkg import (
SamformerForecaster,
)

__all__ = ["Samformer", "Samformer_pkg_v2"]
__all__ = ["Samformer", "SamformerForecaster"]
Loading
Loading