From 3280392da4d6924b1b22b7ada686d5bb8ea3b0a1 Mon Sep 17 00:00:00 2001 From: Faakhir30 Date: Thu, 6 Aug 2026 20:02:43 +0500 Subject: [PATCH 1/5] implmement adapter Signed-off-by: Faakhir30 --- .../models/nbeats/_nbeats_adapter_v2.py | 160 ++++++++++++++++++ 1 file changed, 160 insertions(+) create mode 100644 pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py diff --git a/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py b/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py new file mode 100644 index 000000000..b21fac234 --- /dev/null +++ b/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py @@ -0,0 +1,160 @@ +"""Shared N-Beats adapter for pytorch-forecasting v2.""" + +from typing import Any, Optional, Union + +import torch +from torch import nn +from torch.optim import Optimizer + +from pytorch_forecasting.layers._nbeats._blocks import ( + NBEATSSeasonalBlock, + NBEATSTrendBlock, +) +from pytorch_forecasting.metrics import Metric +from pytorch_forecasting.models.base._tslib_base_model_v2 import TslibBaseModel + + +class NBeatsAdapterV2(TslibBaseModel): + """Shared forward / training helpers for NBeats and NBeatsKAN (v2).""" + + 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, + metadata: dict | None = None, + backcast_loss_ratio: float = 0.0, + **kwargs: Any, + ): + super().__init__( + loss=loss, + logging_metrics=logging_metrics, + optimizer=optimizer, + optimizer_params=optimizer_params, + lr_scheduler=lr_scheduler, + lr_scheduler_params=lr_scheduler_params, + metadata=metadata, + ) + self.backcast_loss_ratio = backcast_loss_ratio + + def _target_from_batch(self, x: dict[str, torch.Tensor]) -> torch.Tensor: + """Extract univariate target history. + + v1 used ``x["encoder_cont"][..., 0]``. v2 tslib batches keep the target + in ``history_target``. + """ + target = x["history_target"] + if target.ndim == 3: + target = target[..., 0] + return target + + def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Pass forward of network. + + Network steps match v1 ``NBeatsAdapter.forward``; only input assembly + and output packaging differ for the v2 API. + """ + # --- v2 batch adapter (v1: target = x["encoder_cont"][..., 0]) --- + target = self._target_from_batch(x) + + # --- same as v1 from here --- + timesteps = self.context_length + self.prediction_length + generic_forecast = [ + torch.zeros( + (target.size(0), timesteps), dtype=torch.float32, device=self.device + ) + ] + trend_forecast = [ + torch.zeros( + (target.size(0), timesteps), dtype=torch.float32, device=self.device + ) + ] + seasonal_forecast = [ + torch.zeros( + (target.size(0), timesteps), dtype=torch.float32, device=self.device + ) + ] + forecast = torch.zeros( + (target.size(0), self.prediction_length), + dtype=torch.float32, + device=self.device, + ) + + backcast = target # initialize backcast + for i, block in enumerate(self.net_blocks): + # evaluate block + backcast_block, forecast_block = block(backcast) + + # add for interpretation + full = torch.cat([backcast_block.detach(), forecast_block.detach()], dim=1) + if isinstance(block, NBEATSTrendBlock): + trend_forecast.append(full) + elif isinstance(block, NBEATSSeasonalBlock): + seasonal_forecast.append(full) + else: + generic_forecast.append(full) + + # update backcast and forecast + backcast = ( + backcast - backcast_block + ) # do not use backcast -= backcast_block as this signifies an inline operation # noqa: E501 + forecast = forecast + forecast_block + + prediction = forecast.unsqueeze(-1) + backcast_out = (target - backcast).unsqueeze(-1) + trend = torch.stack(trend_forecast, dim=0).sum(0).unsqueeze(-1) + seasonality = torch.stack(seasonal_forecast, dim=0).sum(0).unsqueeze(-1) + generic = torch.stack(generic_forecast, dim=0).sum(0).unsqueeze(-1) + + # v1 applied transform_output via BaseModel; v2 tslib does so when scales exist + if "target_scale" in x: + prediction = self.transform_output(prediction, x["target_scale"]) + backcast_out = self.transform_output(backcast_out, x["target_scale"]) + trend = self.transform_output(trend, x["target_scale"]) + seasonality = self.transform_output(seasonality, x["target_scale"]) + generic = self.transform_output(generic, x["target_scale"]) + + # v1: to_network_output(...); v2: plain dict + return { + "prediction": prediction, + "backcast": backcast_out, + "trend": trend, + "seasonality": seasonality, + "generic": generic, + } + + def training_step( + self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int + ) -> dict[str, torch.Tensor]: + """Training step with optional backcast loss (v1 ``step`` parity).""" + x, y = batch + out = self(x) + y_hat = out["prediction"] + loss = self.loss(y_hat, y) + + if self.backcast_loss_ratio > 0: + backcast = out["backcast"].squeeze(-1) + encoder_target = self._target_from_batch(x) + + backcast_weight = ( + self.backcast_loss_ratio + * self.prediction_length + / max(self.context_length, 1) + ) + backcast_weight = backcast_weight / (backcast_weight + 1) + forecast_weight = 1 - backcast_weight + + # Compute backcast term directly (avoid Metric.update state / shape quirks). + # v1 used self.loss(backcast, encoder_target); v2 BaseModel losses are + # wired for forecast horizon shapes only. + backcast_loss = (backcast - encoder_target).abs().mean() * backcast_weight + loss = loss * forecast_weight + backcast_loss + + self.log( + "train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True + ) + self.log_metrics(y_hat, y, prefix="train") + return {"loss": loss} From 0fdcee247136ea744d36bf35aedb0c13424f3058 Mon Sep 17 00:00:00 2001 From: Faakhir30 Date: Fri, 7 Aug 2026 16:56:36 +0500 Subject: [PATCH 2/5] migrate nbeats Signed-off-by: Faakhir30 --- .../models/nbeats/_nbeats_pkg_v2.py | 98 +++++++++++++ .../models/nbeats/_nbeats_v2.py | 130 ++++++++++++++++++ 2 files changed, 228 insertions(+) create mode 100644 pytorch_forecasting/models/nbeats/_nbeats_pkg_v2.py create mode 100644 pytorch_forecasting/models/nbeats/_nbeats_v2.py diff --git a/pytorch_forecasting/models/nbeats/_nbeats_pkg_v2.py b/pytorch_forecasting/models/nbeats/_nbeats_pkg_v2.py new file mode 100644 index 000000000..d93a17aba --- /dev/null +++ b/pytorch_forecasting/models/nbeats/_nbeats_pkg_v2.py @@ -0,0 +1,98 @@ +"""NBeats v2 package container.""" + +from pytorch_forecasting.base._base_pkg import Base_pkg + + +class NBeats_pkg_v2(Base_pkg): + """NBeats v2 package container.""" + + _tags = { + "info:name": "NBeats", + "info:compute": 1, + "info:y_type": ["numeric"], + "authors": [ + "dmitri-carpov" # paper author + "jdb78", # for v1 + "Faakhir30", + ], + "capability:exogenous": False, + "capability:multivariate": False, + "capability:pred_int": False, + "capability:flexible_history_length": False, + "capability:cold_start": False, + } + + @classmethod + def get_cls(cls): + """Get model class.""" + from pytorch_forecasting.models.nbeats._nbeats_v2 import NBeats + + return NBeats + + @classmethod + def get_datamodule_cls(cls): + """Get the underlying DataModule class.""" + from pytorch_forecasting.data.data_module import TslibDataModule + + return TslibDataModule + + @classmethod + def get_test_train_params(cls): + """Return testing parameter settings for the trainer.""" + from pytorch_forecasting.metrics import MAE, MAPE, SMAPE + + params = [ + { + "widths": [16, 32], + "num_blocks": [1, 1], + "num_block_layers": [2, 2], + }, + { + "backcast_loss_ratio": 1.0, + "widths": [16, 32], + "num_blocks": [1, 1], + "num_block_layers": [2, 2], + }, + { + "stack_types": ["generic"], + "num_blocks": [1], + "num_block_layers": [2], + "widths": [16], + "expansion_coefficient_lengths": [8], + "sharing": [False], + }, + { + "loss": MAE(), + "widths": [16, 32], + "num_blocks": [1, 1], + "num_block_layers": [2, 2], + }, + { + "loss": MAPE(), + "logging_metrics": [SMAPE()], + "widths": [16, 32], + "num_blocks": [1, 1], + "num_block_layers": [2, 2], + }, + { + "optimizer": "adamw", + "lr_scheduler": "cosine_annealing", + "lr_scheduler_params": {"T_max": 5}, + "widths": [16, 32], + "num_blocks": [1, 1], + "num_block_layers": [2, 2], + }, + ] + + default_dm_cfg = { + "context_length": 8, + "prediction_length": 3, + "add_relative_time_idx": False, + } + + for param in params: + current_dm_cfg = param.get("datamodule_cfg", {}) + default_dm_cfg.update(current_dm_cfg) + param["datamodule_cfg"] = default_dm_cfg.copy() + + return params diff --git a/pytorch_forecasting/models/nbeats/_nbeats_v2.py b/pytorch_forecasting/models/nbeats/_nbeats_v2.py new file mode 100644 index 000000000..bfb59a6ef --- /dev/null +++ b/pytorch_forecasting/models/nbeats/_nbeats_v2.py @@ -0,0 +1,130 @@ +""" +N-Beats model for pytorch-forecasting v2 (no covariates). +""" + +from typing import Any, Optional, Union + +from torch import nn +from torch.optim import Optimizer + +from pytorch_forecasting.layers._nbeats._blocks import ( + NBEATSGenericBlock, + NBEATSSeasonalBlock, + NBEATSTrendBlock, +) +from pytorch_forecasting.metrics import MAE, MAPE, RMSE, SMAPE, Metric +from pytorch_forecasting.models.nbeats._nbeats_adapter_v2 import NBeatsAdapterV2 + + +class NBeats(NBeatsAdapterV2): + """ + N-BEATS for pytorch-forecasting v2. + + Based on + `N-BEATS: Neural basis expansion analysis for interpretable time series + forecasting `_. + + Network construction matches the v1 ``NBeats`` class; ``context_length`` / + ``prediction_length`` come from datamodule ``metadata`` instead of + ``from_dataset``. + """ + + @classmethod + def _pkg(cls): + """Package for the model.""" + from pytorch_forecasting.models.nbeats._nbeats_pkg_v2 import NBeats_pkg_v2 + + return NBeats_pkg_v2 + + def __init__( + self, + loss: Metric, + stack_types: list[str] | None = None, + num_blocks: list[int] | None = None, + num_block_layers: list[int] | None = None, + widths: list[int] | None = None, + sharing: list[bool] | None = None, + expansion_coefficient_lengths: list[int] | None = None, + dropout: float = 0.1, + backcast_loss_ratio: float = 0.0, + 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, + metadata: dict | None = None, + **kwargs: Any, + ): + if expansion_coefficient_lengths is None: + expansion_coefficient_lengths = [3, 7] + if sharing is None: + sharing = [True, True] + if widths is None: + widths = [32, 512] + if num_block_layers is None: + num_block_layers = [3, 3] + if num_blocks is None: + num_blocks = [3, 3] + if stack_types is None: + stack_types = ["trend", "seasonality"] + if logging_metrics is None: + logging_metrics = [SMAPE(), MAE(), RMSE(), MAPE()] + + super().__init__( + loss=loss, + logging_metrics=logging_metrics, + optimizer=optimizer, + optimizer_params=optimizer_params, + lr_scheduler=lr_scheduler, + lr_scheduler_params=lr_scheduler_params, + metadata=metadata, + backcast_loss_ratio=backcast_loss_ratio, + ) + self.save_hyperparameters(ignore=["loss", "logging_metrics", "metadata"]) + + self.stack_types = stack_types + self.num_blocks = num_blocks + self.num_block_layers = num_block_layers + self.widths = widths + self.sharing = sharing + self.expansion_coefficient_lengths = expansion_coefficient_lengths + self.dropout = dropout + + self._init_network() + + def _init_network(self): + """Build N-BEATS stacks (same block wiring as v1).""" + self.net_blocks = nn.ModuleList() + for stack_id, stack_type in enumerate(self.stack_types): + for _ in range(self.num_blocks[stack_id]): + if stack_type == "generic": + net_block = NBEATSGenericBlock( + units=self.widths[stack_id], + thetas_dim=self.expansion_coefficient_lengths[stack_id], + num_block_layers=self.num_block_layers[stack_id], + backcast_length=self.context_length, + forecast_length=self.prediction_length, + dropout=self.dropout, + ) + elif stack_type == "seasonality": + net_block = NBEATSSeasonalBlock( + units=self.widths[stack_id], + num_block_layers=self.num_block_layers[stack_id], + backcast_length=self.context_length, + forecast_length=self.prediction_length, + min_period=self.expansion_coefficient_lengths[stack_id], + dropout=self.dropout, + ) + elif stack_type == "trend": + net_block = NBEATSTrendBlock( + units=self.widths[stack_id], + thetas_dim=self.expansion_coefficient_lengths[stack_id], + num_block_layers=self.num_block_layers[stack_id], + backcast_length=self.context_length, + forecast_length=self.prediction_length, + dropout=self.dropout, + ) + else: + raise ValueError(f"Unknown stack type {stack_type}") + + self.net_blocks.append(net_block) From 1b827d9c4d7da2fa48bef9ad72cfd364f73f7ff5 Mon Sep 17 00:00:00 2001 From: Faakhir30 Date: Sat, 22 Aug 2026 21:40:03 +0500 Subject: [PATCH 3/5] update validatoin and test steps Signed-off-by: Faakhir30 --- .../models/nbeats/_nbeats_adapter_v2.py | 97 +++++++++++++++++-- 1 file changed, 88 insertions(+), 9 deletions(-) diff --git a/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py b/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py index b21fac234..834e354b6 100644 --- a/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py +++ b/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py @@ -126,12 +126,17 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: "generic": generic, } - def training_step( - self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int - ) -> dict[str, torch.Tensor]: - """Training step with optional backcast loss (v1 ``step`` parity).""" - x, y = batch - out = self(x) + def _compute_loss( + self, + x: dict[str, torch.Tensor], + y: torch.Tensor, + out: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Forecast loss plus optional backcast term (v1 ``step`` parity). + + Applied for train / val / test (not predict), matching v1's + ``not self.predicting`` guard on the shared ``step()``. + """ y_hat = out["prediction"] loss = self.loss(y_hat, y) @@ -147,14 +152,88 @@ def training_step( backcast_weight = backcast_weight / (backcast_weight + 1) forecast_weight = 1 - backcast_weight - # Compute backcast term directly (avoid Metric.update state / shape quirks). - # v1 used self.loss(backcast, encoder_target); v2 BaseModel losses are - # wired for forecast horizon shapes only. backcast_loss = (backcast - encoder_target).abs().mean() * backcast_weight loss = loss * forecast_weight + backcast_loss + return loss, y_hat + + def training_step( + self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int + ) -> dict[str, torch.Tensor]: + """ + Training step for the model with optional backcast loss. + + Parameters + ---------- + batch : Tuple[Dict[str, torch.Tensor]] + Batch of data containing input and target tensors. + batch_idx : int + Index of the batch. + + Returns + ------- + STEP_OUTPUT + Dictionary containing the loss and other metrics. + """ + x, y = batch + out = self(x) + loss, y_hat = self._compute_loss(x, y, out) self.log( "train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True ) self.log_metrics(y_hat, y, prefix="train") return {"loss": loss} + + def validation_step( + self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int + ) -> dict[str, torch.Tensor]: + """ + Validation step for the model with optional backcast loss. + + Parameters + ---------- + batch : Tuple[Dict[str, torch.Tensor]] + Batch of data containing input and target tensors. + batch_idx : int + Index of the batch. + + Returns + ------- + STEP_OUTPUT + Dictionary containing the loss and other metrics. + """ + x, y = batch + out = self(x) + loss, y_hat = self._compute_loss(x, y, out) + self.log( + "val_loss", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True + ) + self.log_metrics(y_hat, y, prefix="val") + return {"val_loss": loss} + + def test_step( + self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int + ) -> dict[str, torch.Tensor]: + """ + Test step for the model with optional backcast loss. + + Parameters + ---------- + batch : Tuple[Dict[str, torch.Tensor]] + Batch of data containing input and target tensors. + batch_idx : int + Index of the batch. + + Returns + ------- + STEP_OUTPUT + Dictionary containing the loss and other metrics. + """ + x, y = batch + out = self(x) + loss, y_hat = self._compute_loss(x, y, out) + self.log( + "test_loss", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True + ) + self.log_metrics(y_hat, y, prefix="test") + return {"test_loss": loss} From 4a34751129cb4bf70b4f6aeeb20e6f60276536ec Mon Sep 17 00:00:00 2001 From: Faakhir30 Date: Wed, 2 Sep 2026 06:43:00 +0500 Subject: [PATCH 4/5] rework adapter towards BaseModel --- .../models/nbeats/_nbeats_adapter_v2.py | 192 ++++++++++++++---- 1 file changed, 150 insertions(+), 42 deletions(-) diff --git a/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py b/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py index 834e354b6..b40f75955 100644 --- a/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py +++ b/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py @@ -1,6 +1,8 @@ -"""Shared N-Beats adapter for pytorch-forecasting v2.""" +""" +N-Beats model adapter for timeseries forecasting (v2). +""" -from typing import Any, Optional, Union +from typing import Any import torch from torch import nn @@ -11,11 +13,44 @@ NBEATSTrendBlock, ) from pytorch_forecasting.metrics import Metric -from pytorch_forecasting.models.base._tslib_base_model_v2 import TslibBaseModel - - -class NBeatsAdapterV2(TslibBaseModel): - """Shared forward / training helpers for NBeats and NBeatsKAN (v2).""" +from pytorch_forecasting.models.base._base_model_v2 import BaseModel + + +class NBeatsAdapterV2(BaseModel): + """ + Shared forward and training logic for the N-Beats model family (v2). + + Subclasses define stack construction in ``_init_network``; this + adapter implements the iterative backcast/forecast loop and optional + backcast loss. + + Univariate models use ``target_past``; exogenous variants (e.g. NBEATx) can + extend ``forward`` to consume ``encoder_cont`` / ``decoder_cont``. + + Parameters + ---------- + loss : Metric + Loss function for the forecast horizon. + logging_metrics : list[nn.Module], optional + Metrics to log during training, validation, and testing. + optimizer : Optimizer or str, optional + Optimizer for training. Default is ``"adam"``. + optimizer_params : dict, optional + Keyword arguments passed to the optimizer constructor. + lr_scheduler : str, optional + Learning rate scheduler name. + lr_scheduler_params : dict, optional + Keyword arguments passed to the scheduler constructor. + metadata : dict, optional + Metadata from ``EncoderDecoderTimeSeriesDataModule`` (``max_encoder_length``, + ``max_prediction_length``, ``encoder_cont``, etc.). + backcast_loss_ratio : float, default=0.0 + Weight of the backcast reconstruction term relative to forecast loss. + When ``0``, only forecast loss is used. When positive, train, validation, + and test steps combine forecast and backcast losses. + **kwargs + Ignored; reserved for subclass hyperparameters. + """ def __init__( self, @@ -36,31 +71,95 @@ def __init__( optimizer_params=optimizer_params, lr_scheduler=lr_scheduler, lr_scheduler_params=lr_scheduler_params, - metadata=metadata, ) + self.metadata = metadata or {} + self.context_length = self.metadata.get("max_encoder_length", 0) + self.prediction_length = self.metadata.get("max_prediction_length", 0) + self.encoder_cont_dim = self.metadata.get("encoder_cont", 0) + self.decoder_cont_dim = self.metadata.get("decoder_cont", 0) self.backcast_loss_ratio = backcast_loss_ratio def _target_from_batch(self, x: dict[str, torch.Tensor]) -> torch.Tensor: - """Extract univariate target history. + """ + Extract univariate target history from an encoder-decoder batch. + + Parameters + ---------- + x : dict[str, torch.Tensor] + Input batch. Uses ``target_past`` (v2 encoder-decoder) or, as a + fallback, ``history_target`` (tslib batches). - v1 used ``x["encoder_cont"][..., 0]``. v2 tslib batches keep the target - in ``history_target``. + Returns + ------- + torch.Tensor + Target history of shape ``(batch_size, context_length)``. """ - target = x["history_target"] + if "target_past" in x: + target = x["target_past"] + elif "history_target" in x: + target = x["history_target"] + else: + raise KeyError("Batch must contain 'target_past' or 'history_target'.") + if target.ndim == 3: target = target[..., 0] return target + def transform_output( + self, + y_hat: torch.Tensor, + target_scale: torch.Tensor | dict[str, torch.Tensor] | list[torch.Tensor], + ) -> torch.Tensor: + """ + Rescale model outputs to the original target scale. + + Parameters + ---------- + y_hat : torch.Tensor + Normalized model output. + target_scale : torch.Tensor or dict + Scale information from the batch. Encoder-decoder batches provide a + tensor; tslib batches may provide a dict with ``scale`` and ``center``. + + Returns + ------- + torch.Tensor + Output rescaled to the original target scale. + """ + if isinstance(target_scale, dict): + scale = target_scale["scale"] + center = target_scale.get("center", 0) + while scale.dim() < y_hat.dim(): + scale = scale.unsqueeze(-1) + if torch.is_tensor(center): + center = center.unsqueeze(-1) + return y_hat * scale + center + + scale = ( + target_scale[0] if isinstance(target_scale, (list, tuple)) else target_scale + ) + while scale.dim() < y_hat.dim(): + scale = scale.unsqueeze(-1) + return y_hat * scale + def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Pass forward of network. + """ + Forward pass through the N-Beats block stack. + + Parameters + ---------- + x : dict[str, torch.Tensor] + Input batch from the datamodule. Must contain ``target_past`` (or + ``history_target``). May contain ``target_scale`` for inverse scaling. - Network steps match v1 ``NBeatsAdapter.forward``; only input assembly - and output packaging differ for the v2 API. + Returns + ------- + dict[str, torch.Tensor] + Model outputs with keys ``prediction``, ``backcast``, ``trend``, + ``seasonality``, and ``generic``. """ - # --- v2 batch adapter (v1: target = x["encoder_cont"][..., 0]) --- target = self._target_from_batch(x) - # --- same as v1 from here --- timesteps = self.context_length + self.prediction_length generic_forecast = [ torch.zeros( @@ -83,12 +182,10 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: device=self.device, ) - backcast = target # initialize backcast - for i, block in enumerate(self.net_blocks): - # evaluate block + backcast = target + for block in self.net_blocks: backcast_block, forecast_block = block(backcast) - # add for interpretation full = torch.cat([backcast_block.detach(), forecast_block.detach()], dim=1) if isinstance(block, NBEATSTrendBlock): trend_forecast.append(full) @@ -97,7 +194,6 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: else: generic_forecast.append(full) - # update backcast and forecast backcast = ( backcast - backcast_block ) # do not use backcast -= backcast_block as this signifies an inline operation # noqa: E501 @@ -109,7 +205,6 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: seasonality = torch.stack(seasonal_forecast, dim=0).sum(0).unsqueeze(-1) generic = torch.stack(generic_forecast, dim=0).sum(0).unsqueeze(-1) - # v1 applied transform_output via BaseModel; v2 tslib does so when scales exist if "target_scale" in x: prediction = self.transform_output(prediction, x["target_scale"]) backcast_out = self.transform_output(backcast_out, x["target_scale"]) @@ -117,7 +212,6 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: seasonality = self.transform_output(seasonality, x["target_scale"]) generic = self.transform_output(generic, x["target_scale"]) - # v1: to_network_output(...); v2: plain dict return { "prediction": prediction, "backcast": backcast_out, @@ -132,10 +226,24 @@ def _compute_loss( y: torch.Tensor, out: dict[str, torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor]: - """Forecast loss plus optional backcast term (v1 ``step`` parity). + """ + Compute forecast loss, optionally combined with backcast loss. + + Parameters + ---------- + x : dict[str, torch.Tensor] + Input batch (used for encoder target when backcast loss is enabled). + y : torch.Tensor + Forecast horizon target. + out : dict[str, torch.Tensor] + Forward pass output. - Applied for train / val / test (not predict), matching v1's - ``not self.predicting`` guard on the shared ``step()``. + Returns + ------- + loss : torch.Tensor + Scalar loss for logging and optimization. + y_hat : torch.Tensor + Forecast predictions from ``out["prediction"]``. """ y_hat = out["prediction"] loss = self.loss(y_hat, y) @@ -161,19 +269,19 @@ def training_step( self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int ) -> dict[str, torch.Tensor]: """ - Training step for the model with optional backcast loss. + Training step with optional backcast loss. Parameters ---------- - batch : Tuple[Dict[str, torch.Tensor]] - Batch of data containing input and target tensors. + batch : tuple[dict[str, torch.Tensor]] + ``(x, y)`` from the dataloader. batch_idx : int Index of the batch. Returns ------- - STEP_OUTPUT - Dictionary containing the loss and other metrics. + dict[str, torch.Tensor] + Dictionary with key ``loss``. """ x, y = batch out = self(x) @@ -188,19 +296,19 @@ def validation_step( self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int ) -> dict[str, torch.Tensor]: """ - Validation step for the model with optional backcast loss. + Validation step with optional backcast loss. Parameters ---------- - batch : Tuple[Dict[str, torch.Tensor]] - Batch of data containing input and target tensors. + batch : tuple[dict[str, torch.Tensor]] + ``(x, y)`` from the dataloader. batch_idx : int Index of the batch. Returns ------- - STEP_OUTPUT - Dictionary containing the loss and other metrics. + dict[str, torch.Tensor] + Dictionary with key ``val_loss``. """ x, y = batch out = self(x) @@ -215,19 +323,19 @@ def test_step( self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int ) -> dict[str, torch.Tensor]: """ - Test step for the model with optional backcast loss. + Test step with optional backcast loss. Parameters ---------- - batch : Tuple[Dict[str, torch.Tensor]] - Batch of data containing input and target tensors. + batch : tuple[dict[str, torch.Tensor]] + ``(x, y)`` from the dataloader. batch_idx : int Index of the batch. Returns ------- - STEP_OUTPUT - Dictionary containing the loss and other metrics. + dict[str, torch.Tensor] + Dictionary with key ``test_loss``. """ x, y = batch out = self(x) From b41988ca93db4c691bed517e1f0ef000f2b7d54d Mon Sep 17 00:00:00 2001 From: Faakhir30 Date: Wed, 2 Sep 2026 07:23:00 +0500 Subject: [PATCH 5/5] rework NBEATS towards baseModel --- .../models/nbeats/_nbeats_pkg_v2.py | 12 ++-- .../models/nbeats/_nbeats_v2.py | 62 ++++++++++++++++--- 2 files changed, 63 insertions(+), 11 deletions(-) diff --git a/pytorch_forecasting/models/nbeats/_nbeats_pkg_v2.py b/pytorch_forecasting/models/nbeats/_nbeats_pkg_v2.py index d93a17aba..3c89c493b 100644 --- a/pytorch_forecasting/models/nbeats/_nbeats_pkg_v2.py +++ b/pytorch_forecasting/models/nbeats/_nbeats_pkg_v2.py @@ -32,9 +32,11 @@ def get_cls(cls): @classmethod def get_datamodule_cls(cls): """Get the underlying DataModule class.""" - from pytorch_forecasting.data.data_module import TslibDataModule + from pytorch_forecasting.data.data_module import ( + EncoderDecoderTimeSeriesDataModule, + ) - return TslibDataModule + return EncoderDecoderTimeSeriesDataModule @classmethod def get_test_train_params(cls): @@ -85,8 +87,10 @@ def get_test_train_params(cls): ] default_dm_cfg = { - "context_length": 8, - "prediction_length": 3, + "max_encoder_length": 8, + "max_prediction_length": 3, + "min_encoder_length": 8, + "min_prediction_length": 3, "add_relative_time_idx": False, } diff --git a/pytorch_forecasting/models/nbeats/_nbeats_v2.py b/pytorch_forecasting/models/nbeats/_nbeats_v2.py index bfb59a6ef..348ceeb6c 100644 --- a/pytorch_forecasting/models/nbeats/_nbeats_v2.py +++ b/pytorch_forecasting/models/nbeats/_nbeats_v2.py @@ -18,15 +18,63 @@ class NBeats(NBeatsAdapterV2): """ - N-BEATS for pytorch-forecasting v2. + N-BEATS for pytorch-forecasting v2 (univariate time series forecasting). - Based on - `N-BEATS: Neural basis expansion analysis for interpretable time series - forecasting `_. + Based on the article `N-BEATS: Neural basis expansion analysis for + interpretable time series forecasting + `_. The network has (if used as + ensemble) outperformed all other methods including ensembles of + traditional statistical methods in the M4 competition. The M4 + competition is arguably the most important benchmark for univariate + time series forecasting. - Network construction matches the v1 ``NBeats`` class; ``context_length`` / - ``prediction_length`` come from datamodule ``metadata`` instead of - ``from_dataset``. + Parameters + ---------- + loss : Metric + Loss metric to optimize during training. + stack_types : list of str, optional + One of "generic", "seasonality", or "trend". A list of strings + of length equal to the number of stacks. Default is ["trend", + "seasonality"] for interpretable mode. + num_blocks : list of int, optional + The number of blocks per stack. List length equal to number of + stacks. Default is [3, 3]. + num_block_layers : list of int, optional + Number of fully connected layers with ReLU activation per block. + List length equal to number of stacks. Default is [3, 3]. + widths : list of int, optional + Widths of fully connected layers with ReLU activation. List + length equal to number of stacks. Default is [32, 512]. + sharing : list of bool, optional + Whether weights are shared across blocks within a stack. List + length equal to number of stacks. Default is [True, True]. + expansion_coefficient_lengths : list of int, optional + If type is "generic", length of expansion coefficients; if + "trend", degree of polynomial; if "seasonality", minimum period. + List length equal to number of stacks. Default is [3, 7]. + dropout : float, optional + Dropout probability applied in the network. Helps prevent + overfitting. Default is 0.1. + backcast_loss_ratio : float, optional + Weight of backcast loss relative to forecast loss. 1.0 gives + equal weight; 0.0 means no backcast loss. Default is 0.0. + logging_metrics : list of nn.Module, optional + List of metrics logged during training. Defaults to + [SMAPE(), MAE(), RMSE(), MAPE()]. + optimizer : Optimizer or str, optional + Optimizer to use for training. Can be a torch.optim.Optimizer + class or string like "adam". Default is "adam". + optimizer_params : dict, optional + Additional parameters for the optimizer. Default is None. + lr_scheduler : str, optional + Learning rate scheduler type. Default is None. + lr_scheduler_params : dict, optional + Additional parameters for the learning rate scheduler. Default + is None. + metadata : dict, optional + Additional metadata for the model. Default is None. + **kwargs + Additional arguments forwarded to :py:class:`~NBeatsAdapterV2`. """ @classmethod