From 3d704c6ac6f6cdf532a91c60a8e2c5975b424506 Mon Sep 17 00:00:00 2001 From: Lucas Date: Tue, 15 Sep 2026 23:37:59 +0200 Subject: [PATCH 1/7] feat: add T0Model foundation model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add classes to use t0-alpha for inference and fine-tuning. - Add t0-alpha fine-tuning example to the dedicated notebook. - Add tests for T0Model classes - Add tfc-t0 v0.4.0 to the optional dependencies Co-authored-by: Huikan Xiang Co-authored-by: Geoffrey NΓ©giar --- CHANGELOG.md | 2 + INSTALL.md | 2 + README.md | 1 + darts/models/__init__.py | 2 + darts/models/forecasting/__init__.py | 1 + darts/models/forecasting/t0_model.py | 425 ++++++++++++++++++ darts/tests/conftest.py | 1 + .../forecasting/foundation_test_utils.py | 66 +++ .../models/forecasting/test_foundation.py | 27 +- darts/tests/models/forecasting/test_t0.py | 214 +++++++++ docs/source/index.rst | 6 + ...oundation-Model-Fine-Tuning-examples.ipynb | 107 ++++- pyproject.toml | 1 + 13 files changed, 852 insertions(+), 3 deletions(-) create mode 100644 darts/models/forecasting/t0_model.py create mode 100644 darts/tests/models/forecasting/test_t0.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1e27074682..8b5ff690b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,8 @@ but cannot always guarantee backwards compatibility. Changes that may **break co **Improved** +- πŸš€ Added new forecasting model `T0Model` : [The Forecasting Company's open-weights ~100M-parameter foundation model](https://huggingface.co/theforecastingcompany/t0-alpha) for zero-shot forecasting. It supports univariate, multivariate, and multiple time series as well as past and future covariates, without training, and can output deterministic or probabilistic forecasts. It can also be fine-tuned (full or partial) with `enable_finetuning`. [#3142](https://github.com/unit8co/darts/pull/3142) by [Geoffrey NΓ©giar](https://github.com/GeoffNN), [Huikan Xiang](https://github.com/huikan-tfc) and [Lucas Meyer](https://github.com/LTMeyer). + **Fixed** **Dependencies** diff --git a/INSTALL.md b/INSTALL.md index ad775a2881..050750e6a1 100644 --- a/INSTALL.md +++ b/INSTALL.md @@ -22,6 +22,7 @@ Some models have additional dependencies that are not included in the `all` inst |-----------------------|-----------------------| | `NeuralForecastModel` | neuralforecast>=3.0.0 | | `TiRexModel` | tirex-ts>=1.4.0 | +| `T0Model` | tfc-t0>=0.4.0 | Some optional integrations also require additional dependencies: @@ -59,6 +60,7 @@ Some models have dependencies not available on conda-forge. To use them, you nee | Model | Dependencies | |-----------------------|-----------------------| | `TiRexModel` | tirex-ts>=1.4.0 | +| `T0Model` | tfc-t0>=0.4.0 | ## Other Information diff --git a/README.md b/README.md index 755e3d4745..5ffa8a8bf5 100644 --- a/README.md +++ b/README.md @@ -283,6 +283,7 @@ Here's a breakdown of the forecasting models currently implemented in Darts. Our | [TimesFM2p5Model](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.timesfm2p5_model.html#darts.models.forecasting.timesfm2p5_model.TimesFM2p5Model) | [TimesFM 1.0 paper](https://arxiv.org/abs/2310.10688), [Google blog post](https://research.google/blog/a-decoder-only-foundation-model-for-time-series-forecasting) | βœ… βœ… | πŸ”΄ πŸ”΄ πŸ”΄ | βœ… βœ… | βœ… | | [TiRexModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.tirex_model.html#darts.models.forecasting.tirex_model.TiRexModel) | [TiRex paper](https://arxiv.org/abs/2505.23719), [TiRex GitHub](https://github.com/NX-AI/tirex) | βœ… βœ… | πŸ”΄ πŸ”΄ πŸ”΄ | βœ… βœ… | βœ… | | [PatchTSTFMModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.patchtst_fm_model.html#darts.models.forecasting.patchtst_fm_model.PatchTSTFMModel) | [PatchTST-FM paper](https://arxiv.org/abs/2602.06909), [PatchTST-FM GitHub](https://github.com/ibm-granite/granite-tsfm) | βœ… βœ… | πŸ”΄ πŸ”΄ πŸ”΄ | βœ… βœ… | βœ… | +| [T0Model](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.t0_model.html#darts.models.forecasting.t0_model.T0Model) | [T0 model card](https://huggingface.co/theforecastingcompany/t0-alpha), [tfc-t0 GitHub](https://github.com/theforecastingcompany/tfc-t0) | βœ… βœ… | βœ… βœ… πŸ”΄ | βœ… βœ… | βœ… | | **Ensemble Models**
([GlobalForecastingModel](https://unit8co.github.io/darts/userguide/covariates.html#global-forecasting-models-gfms)): Model support is dependent on ensembled forecasting models and the ensemble model itself | | | | | | | [NaiveEnsembleModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.naive_ensemble_model.html#darts.models.forecasting.naive_ensemble_model.NaiveEnsembleModel) | | βœ… βœ… | βœ… βœ… βœ… | βœ… βœ… | βœ… | | [RegressionEnsembleModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.regression_ensemble_model.html#darts.models.forecasting.regression_ensemble_model.RegressionEnsembleModel) | | βœ… βœ… | βœ… βœ… βœ… | βœ… βœ… | βœ… | diff --git a/darts/models/__init__.py b/darts/models/__init__.py index 5fb832fd34..5821425d3a 100644 --- a/darts/models/__init__.py +++ b/darts/models/__init__.py @@ -106,6 +106,7 @@ SKLearnClassifierModel as SKLearnClassifierModel, ) from darts.models.forecasting.sklearn_model import SKLearnModel as SKLearnModel + from darts.models.forecasting.t0_model import T0Model as T0Model from darts.models.forecasting.tcn_model import TCNModel as TCNModel from darts.models.forecasting.tft_model import TFTModel as TFTModel from darts.models.forecasting.theta import FourTheta as FourTheta @@ -191,6 +192,7 @@ "PatchTSTFMModel": ("darts.models.forecasting.patchtst_fm_model", "(Py)Torch"), "TimesFM2p5Model": ("darts.models.forecasting.timesfm2p5_model", "(Py)Torch"), "TiRexModel": ("darts.models.forecasting.tirex_model", "(Py)Torch and/or TiRex-TS"), + "T0Model": ("darts.models.forecasting.t0_model", "(Py)Torch and/or tfc-t0"), # --- Forecasting: NeuralForecast --- "NeuralForecastModel": ("darts.models.forecasting.nf_model", "NeuralForecast"), # --- Forecasting: Prophet --- diff --git a/darts/models/forecasting/__init__.py b/darts/models/forecasting/__init__.py index 413c4e2826..aea964fe55 100644 --- a/darts/models/forecasting/__init__.py +++ b/darts/models/forecasting/__init__.py @@ -59,6 +59,7 @@ - :class:`~darts.models.forecasting.timesfm2p5_model.TimesFM2p5Model` - :class:`~darts.models.forecasting.tirex_model.TiRexModel` - :class:`~darts.models.forecasting.patchtst_fm_model.PatchTSTFMModel` + - :class:`~darts.models.forecasting.t0_model.T0Model` Ensemble Models (`GlobalForecastingModel `__) - :class:`~darts.models.forecasting.naive_ensemble_model.NaiveEnsembleModel` - :class:`~darts.models.forecasting.regression_ensemble_model.RegressionEnsembleModel` diff --git a/darts/models/forecasting/t0_model.py b/darts/models/forecasting/t0_model.py new file mode 100644 index 0000000000..fba6dc4ad1 --- /dev/null +++ b/darts/models/forecasting/t0_model.py @@ -0,0 +1,425 @@ +""" +T0: Zero-Shot Forecasting +------------------------- + +T0 can be used the same way as other foundation models (e.g. Chronos2). In addition to univariate and +multivariate series, it supports past and future covariates. + +The weights are published on the Hugging Face Hub at `theforecastingcompany/t0-alpha +`__. The repository is gated: request access on +the model page and authenticate (``huggingface-cli login``, or set ``HF_TOKEN``) before first use. + +For detailed examples and tutorials, see: + +* `Foundation Model Examples + `__ +* `Fine-Tuning Examples + `__ +""" + +import dataclasses +import os + +import torch +from t0 import T0Config, T0Forecaster +from t0 import TimeSeries as T0TimeSeries +from t0.data import VariateType + +from darts.logging import get_logger, raise_log +from darts.models.components.huggingface_connector import HuggingFaceConnector +from darts.models.forecasting.foundation_model import FoundationModel +from darts.models.forecasting.pl_forecasting_module import PLForecastingModule +from darts.utils.data.torch_datasets.utils import PLModuleInput +from darts.utils.likelihood_models.torch import QuantileRegression + +logger = get_logger(__name__) + + +class _T0Module(PLForecastingModule): + """PyTorch Lightning module wrapping a pre-loaded T0 forecaster. + + Targets and past covariates are forecast jointly, and past-covariate predictions are dropped. + Future covariates are mapped to T0's ``[batch, n_covariates, context + horizon]`` format. + + Fine-tuning runs T0's forward pass over all pre-trained quantiles; prediction returns only the + user-specified ones. + """ + + def __init__( + self, + hub_model_name: str, + hub_model_revision: str | None, + local_dir: str | os.PathLike | None, + all_quantiles: tuple[float, ...], + enable_finetuning: bool | dict = False, + **kwargs, + ): + super().__init__(**kwargs) + self._enable_finetuning = bool(enable_finetuning) + + # rebuild T0 from the Hub config, then load the weights into it + connector = HuggingFaceConnector( + model_name=hub_model_name, + model_revision=hub_model_revision, + local_dir=local_dir, + ) + config = connector.load_config() + config_kwargs = { + field.name: config[field.name] + for field in dataclasses.fields(T0Config) + if field.name in config + } + if "quantile_levels" in config_kwargs: + config_kwargs["quantile_levels"] = tuple(config_kwargs["quantile_levels"]) + if self._enable_finetuning: + # set dropout to 0, which would otherwise inject noise into a frozen backbone + config_kwargs["dropout"] = 0.0 + self.t0: T0Forecaster = T0Forecaster.from_config(T0Config(**config_kwargs)) + connector.load_model_weights(self.t0) + + # switch the model to train mode if needed. + self.t0.train(self._enable_finetuning) + self.future_len = (self.output_chunk_length or 0) + self.output_chunk_shift + + if self._enable_finetuning: + if self.future_len > self.t0.max_horizon: + raise_log( + ValueError( + f"T0 fine-tuning only supports up to {self.t0.max_horizon} steps got {self.future_len}" + ), + ) + # loss is computed over all pre-trained quantiles to preserve the distribution; + # user-specified quantiles are selected at prediction time + self._finetuning_likelihood = QuantileRegression(list(all_quantiles)) + else: + self._finetuning_likelihood = None + + def forward(self, x_in: PLModuleInput, *args, **kwargs): + """Forward pass returning quantile predictions shaped ``(batch, time, n_targets, n_quantiles)``. + + During training with fine-tuning enabled, all pre-trained quantiles are returned for the loss. + At prediction time, only user-specified quantiles are returned. + + Parameters + ---------- + x_in + ``(x_past, x_future, x_static, future_target)`` the past, future, and static features, as well as + the future target. + *args + Positional arguments passed to the forward method. + **kwargs + Optional keyword arguments. + """ + # Dimension notation in comments below: + # B: batch size + # L: input chunk length + # T: output chunk length + # S: output chunk shift + # H: future length = T + S + # C: target components + # P: past covariate components + # F: future covariate components + # V: context variates = C + P (target + past covariates, jointly forecast) + # Qp: pre-trained quantiles (returned during fine-tuning) + # N: likelihood quantiles (user-specified, 1 if deterministic) + + # `x_past`: (B, L, C + P + F) stack of [past_target, past_covariates, historic_future_covariates]; + # `x_future`: (B, T, F) future covariates, or None. + x_past, x_future, _, _ = x_in + batch_size = x_past.shape[0] + + # Past covariates are forecast jointly with the target (T0 is variate-agnostic) and dropped from the + # output. Future covariates are the trailing columns of `x_past` (their historic part) and are passed to + # T0's covariate branch instead. `x_future` width gives the number of future covariates. + n_future_covs = x_future.shape[-1] if x_future is not None else 0 + n_context = x_past.shape[-1] - n_future_covs + + # context: (B, V, L) + context = x_past[:, :, :n_context].transpose(1, 2) + + # T0 expects covariates over context + horizon. Re-assemble them from the historic part (in `x_past`) + # and the future chunk (`x_future`); the `output_chunk_shift` gap is left NaN (T0 treats NaN as missing). + future_covariates = None + if n_future_covs > 0: + historic = x_past[:, :, n_context:] # (B, L, F) + future = torch.full( + (batch_size, self.future_len, n_future_covs), + torch.nan, + device=x_past.device, + dtype=x_past.dtype, + ) + if x_future is not None: + future[:, -(self.output_chunk_length or 0) :, :] = x_future + # (B, L + H, F) -> (B, F, L + H) + future_covariates = torch.cat([historic, future], dim=1).transpose(1, 2) + + if self.training and self._enable_finetuning: + # scale the input, run one forward pass, and rescale the predictions + patch_size = self.t0.patch_size + context_length = context.shape[-1] + model_input = self.t0.patcher.pad( + T0TimeSeries.from_array( + context, future_covariates, horizon=self.future_len + ) + ) + scaled, loc_scale = self.t0.scaler.scale_input(model_input) + # (rows, patches, patch_size, Qp) in data space + patches = self.t0.scaler.rescale_predictions( + self.t0(scaled), loc_scale, patch_size + ) + # each position predicts the patch that follows it, so the forecast starts at the last + # context patch; ceil-divide to count whole patches + first = -(-context_length // patch_size) - 1 + n_patches = -(-self.future_len // patch_size) + is_target = model_input.variate_type[:, -1] == VariateType.TARGET + # -> (B * V, H, Qp) -> (B, V, H, Qp), matching `predict()`'s multivariate layout + quantiles = ( + patches[is_target][:, first : first + n_patches] + .flatten(1, 2)[:, : self.future_len] + .unflatten(0, (batch_size, n_context)) + ) + else: + user_q: list[float] = ( + self.likelihood.quantiles + if isinstance(self.likelihood, QuantileRegression) + else [0.5] + ) + # quantiles: (B, V, H, N) + quantiles = self.t0.predict( + context, + horizon=self.future_len, + quantiles=user_q, + future_covariates=future_covariates, + ).quantiles + # drop the past-covariate variates, keep targets: (B, V, H, N) -> (B, C, H, N) + quantiles = quantiles[:, : self.n_targets] + # (B, C, H, N) -> (B, H, C, N) -> slice output shift -> (B, T, C, N) + return quantiles.permute(0, 2, 1, 3)[:, self.output_chunk_shift :, :, :] + + def _compute_loss(self, output, target, criterion, sample_weight): + if self.training and self._enable_finetuning: + # compute loss on pre-trained quantiles + return self._finetuning_likelihood.compute_loss( + output, target, sample_weight + ) + return super()._compute_loss(output, target, criterion, sample_weight) + + +class T0Model(FoundationModel): + # Quantile levels T0 was trained on. Other levels are interpolated, so any quantiles in (0, 1) are accepted. + _PRETRAINED_QUANTILES: tuple[float, ...] = (0.1, 0.25, 0.5, 0.75, 0.9) + + def __init__( + self, + input_chunk_length: int, + output_chunk_length: int, + output_chunk_shift: int = 0, + likelihood: QuantileRegression | None = None, + hub_model_name: str = "theforecastingcompany/t0-alpha", + hub_model_revision: str | None = None, + local_dir: str | os.PathLike | None = None, + **kwargs, + ): + """ + T0 foundation model for zero-shot time series forecasting. + + This is a Darts wrapper around The Forecasting Company's open-weights T0 model. The implementation delegates + all forecasting logic to the optional `tfc-t0 `_ package and loads the config + and weights through Darts' :class:`HuggingFaceConnector`, while exposing a standard + :class:`TorchForecastingModel` interface. + + T0 is a ~100M-parameter pre-trained patch-transformer foundation model designed for zero-shot forecasting + across both short and long horizons. + + This model supports univariate and multivariate time series, as well as past and future covariates. + Because T0 is variate-agnostic, past covariates are forecast jointly with the target series (and dropped + from the output); future covariates are conditioned on but not forecast. + + By default, the model is deterministic (median forecast only). To enable probabilistic forecasts, pass a + :class:`~darts.utils.likelihood_models.torch.QuantileRegression` instance to the ``likelihood`` parameter. + It is recommended to call :func:`predict()` with ``predict_likelihood_parameters=True`` or ``num_samples >> 1`` + to get meaningful results. T0 was trained on quantile levels [0.1, 0.25, 0.5, 0.75, 0.9]; other levels are + interpolated, so any quantiles in the open interval (0, 1) may be requested. + + For more details on the T0 model, see the `model card `_ + and the `tfc-t0 repository `_. The weights are open but the + Hub repository is gated, so request access on the model card and authenticate before first use. + + .. note:: + The model can be fine-tuned (full or partial) via ``enable_finetuning``. Fine-tuning runs a + single forward pass, so ``output_chunk_length`` plus ``output_chunk_shift`` must not exceed + T0's ``max_horizon``. The training loss is computed on all pre-trained quantiles to preserve + the pre-trained distribution, and dropout is disabled so a frozen backbone stays + deterministic. + + Parameters + ---------- + input_chunk_length + Number of time steps in the past to take as a model input (per chunk). Applies to the target + series, and past and/or future covariates (if the model supports it). + output_chunk_length + Number of time steps predicted at once (per chunk) by the internal model. Also, the number of future values + from future covariates to use as a model input (if the model supports future covariates). It is not the same + as forecast horizon `n` used in `predict()`, which is the desired number of prediction points generated + using either a one-shot- or autoregressive forecast. Setting `n <= output_chunk_length` prevents + auto-regression. This is useful when the covariates don't extend far enough into the future, or to prohibit + the model from using future values of past and / or future covariates for prediction (depending on the + model's covariate support). + output_chunk_shift + Optionally, the number of steps to shift the start of the output chunk into the future (relative to the + input chunk end). This will create a gap between the input and output. If the model supports + `future_covariates`, the future values are extracted from the shifted output chunk. Predictions will start + `output_chunk_shift` steps after the end of the target `series`. If `output_chunk_shift` is set, the model + cannot generate autoregressive predictions (`n > output_chunk_length`). + likelihood + The likelihood model to be used for probabilistic forecasts. Must be ``None`` or an instance of + :class:`~darts.utils.likelihood_models.torch.QuantileRegression`. Any quantiles in the open interval + (0, 1) are supported (T0 interpolates levels it was not trained on). Default: ``None``, which will make + the model deterministic (median quantile only). + hub_model_name + The model ID on HuggingFace Hub. Default: ``"theforecastingcompany/t0-alpha"``. + hub_model_revision + The model version to use. This can be a branch name, tag name, or commit hash. Default: ``None``, which + will use the default branch from ``hub_model_name``. + local_dir + Optional local directory to load the pre-downloaded model. If specified and the directory is empty, the + model will be downloaded from HuggingFace Hub and saved to this directory. Default is ``None``, which will + use a cache directory managed by ``huggingface_hub`` instead. Note that this is different from the + ``work_dir`` parameter used for saving model checkpoints during fine-tuning. + **kwargs + Optional arguments to initialize the pytorch_lightning.Module, pytorch_lightning.Trainer, and + Darts' :class:`TorchForecastingModel`. + + torch_metrics + A torch metric or a ``MetricCollection`` used for evaluation. A full list of available metrics can be found + at https://torchmetrics.readthedocs.io/en/latest/. Default: ``None``. + batch_size + Number of time series (input and output sequences) used in each prediction pass. Default: ``32``. + model_name + Name of the model. Used for creating checkpoints and saving tensorboard data. If not specified, + defaults to the following string ``"YYYY-mm-dd_HH_MM_SS_torch_model_run_PID"``, where the initial part + of the name is formatted with the local date and time, while PID is the process ID (preventing models + spawned at the same time by different processes to share the same model_name). E.g., + ``"2021-06-14_09_53_32_torch_model_run_44607"``. + work_dir + Path of the working directory, where to save checkpoints and Tensorboard summaries. + Default: current working directory. + log_tensorboard + If set, use Tensorboard to log the different parameters. The logs will be located in: + ``"{work_dir}/darts_logs/{model_name}/logs/"``. Default: ``False``. + force_reset + If set to ``True``, any previously-existing model with the same name will be reset (all checkpoints will + be discarded). Default: ``False``. + save_checkpoints + Whether to automatically save the untrained model and checkpoints from training. + To load the model from checkpoint, call :func:`MyModelClass.load_from_checkpoint()`, where + :class:`MyModelClass` is the :class:`TorchForecastingModel` class that was used (such as :class:`TFTModel`, + :class:`NBEATSModel`, etc.). If set to ``False``, the model can still be manually saved using + :func:`save()` and loaded using :func:`load()`. Default: ``False``. + add_encoders + A large number of past and future covariates can be automatically generated with `add_encoders`. + This can be done by adding multiple pre-defined index encoders and/or custom user-made functions that + will be used as index encoders. Additionally, a transformer such as Darts' :class:`Scaler` can be added to + transform the generated covariates. This happens all under one hood and only needs to be specified at + model creation. + Read :meth:`SequentialEncoder ` to find out more about + ``add_encoders``. Default: ``None``. An example showing some of ``add_encoders`` features: + + .. highlight:: python + .. code-block:: python + + def encode_year(idx): + return (idx.year - 1950) / 50 + + add_encoders={ + 'cyclic': {'future': ['month']}, + 'datetime_attribute': {'future': ['hour', 'dayofweek']}, + 'position': {'past': ['relative'], 'future': ['relative']}, + 'custom': {'past': [encode_year]}, + 'transformer': Scaler(), + 'tz': 'CET' + } + .. + random_state + Controls the randomness of reproducible forecasting. + pl_trainer_kwargs + By default :class:`TorchForecastingModel` creates a PyTorch Lightning Trainer with several useful presets + that performs the training, validation and prediction processes. These presets include automatic + checkpointing, tensorboard logging, setting the torch device and more. + With ``pl_trainer_kwargs`` you can add additional kwargs to instantiate the PyTorch Lightning trainer + object. Check the `PL Trainer documentation + `__ for more information about the + supported kwargs. Default: ``None``. + Running on GPU(s) is also possible using ``pl_trainer_kwargs`` by specifying keys ``"accelerator", + "devices", and "auto_select_gpus"``. Some examples for setting the devices inside the ``pl_trainer_kwargs`` + dict: + + - ``{"accelerator": "cpu"}`` for CPU, + - ``{"accelerator": "gpu", "devices": [i]}`` to use only GPU ``i`` (``i`` must be an integer), + - ``{"accelerator": "gpu", "devices": -1, "auto_select_gpus": True}`` to use all available GPUs. + + For more info, see here: + https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html#trainer-flags , and + https://pytorch-lightning.readthedocs.io/en/stable/accelerators/gpu_basic.html#train-on-multiple-gpus + show_warnings + whether to show warnings raised from PyTorch Lightning. Useful to detect potential issues of + your forecasting use case. Default: ``False``. + + References + ---------- + .. [1] The Forecasting Company, "T0", https://huggingface.co/theforecastingcompany/t0-alpha. + + Examples + -------- + Point forecasting: + + >>> from darts.models import T0Model + >>> from darts.datasets import AirPassengersDataset + >>> series = AirPassengersDataset().load().astype("float32") + >>> model = T0Model(input_chunk_length=24, output_chunk_length=12) + >>> model.fit(series) + >>> pred = model.predict(n=12) + + Probabilistic forecasting: + + >>> from darts.utils.likelihood_models import QuantileRegression + >>> model = T0Model( + ... input_chunk_length=24, + ... output_chunk_length=12, + ... likelihood=QuantileRegression(quantiles=[0.1, 0.5, 0.9]), + ... ) + >>> model.fit(series) + >>> pred = model.predict(n=12, predict_likelihood_parameters=True) + """ + if likelihood is not None and not isinstance(likelihood, QuantileRegression): + raise_log( + ValueError( + f"Only QuantileRegression likelihood is supported for T0 in Darts. " + f"Got {type(likelihood)}." + ), + ) + + super().__init__(**kwargs) + + self.hub_model_name = hub_model_name + self.hub_model_revision = hub_model_revision + self.local_dir = local_dir + + @property + def supports_past_covariates(self) -> bool: + return True + + @property + def supports_future_covariates(self) -> bool: + return True + + def _create_model(self, train_sample) -> PLForecastingModule: + # enable_finetuning is injected into pl_module_params by the base class; + # _T0Module accepts it as an explicit parameter and converts dict form to bool + return _T0Module( + hub_model_name=self.hub_model_name, + hub_model_revision=self.hub_model_revision, + local_dir=self.local_dir, + all_quantiles=self._PRETRAINED_QUANTILES, + **(self.pl_module_params or {}), + ) diff --git a/darts/tests/conftest.py b/darts/tests/conftest.py index e0f48db067..613235c6b3 100644 --- a/darts/tests/conftest.py +++ b/darts/tests/conftest.py @@ -32,6 +32,7 @@ def _package_available(*names: str) -> bool: PLOTLY_AVAILABLE = _package_available("plotly") IPYTHON_AVAILABLE = _package_available("IPython") TIREX_AVAILABLE = _package_available("tirex") +T0_AVAILABLE = _package_available("t0") MLFLOW_AVAILABLE = _package_available("mlflow") tfm_kwargs: dict[str, Any] = { diff --git a/darts/tests/models/forecasting/foundation_test_utils.py b/darts/tests/models/forecasting/foundation_test_utils.py index 08262df0ed..4167842f64 100644 --- a/darts/tests/models/forecasting/foundation_test_utils.py +++ b/darts/tests/models/forecasting/foundation_test_utils.py @@ -8,6 +8,7 @@ import contextlib import functools import shutil +import tempfile from pathlib import Path from unittest.mock import patch @@ -27,6 +28,34 @@ TIMESFM2P5_TINY_MAX_CONTEXT_LENGTH = 64 TIMESFM2P5_TINY_MAX_PREDICTION_LENGTH = 8 + +# ── T0 tiny model ─────────────────────────────────────────────────────────── +# Create a local checkpoint for t0-alpha model. +@functools.lru_cache(maxsize=1) +def tiny_t0_dir() -> str: + """Build a tiny T0 model and save it to a temporary directory, returning the path. + + ``T0Model`` loads through Darts' ``HuggingFaceConnector``, so the path is passed as + ``local_dir`` and the gated t0-alpha weights are never downloaded.""" + from t0 import T0Config, T0Forecaster + + directory = tempfile.mkdtemp(prefix="darts_tiny_t0_") + T0Forecaster.from_config( + T0Config( + embed_dim=32, + num_layers=2, + num_heads=4, + mlp_hidden_dim=64, + patch_size=8, + group_every_n=2, + dropout=0.0, + quantile_levels=(0.1, 0.25, 0.5, 0.75, 0.9), + scaler_use_arcsinh=True, + ) + ).save_pretrained(directory) + return directory + + # ── HuggingFace mock download (Chronos-2 tiny artefact) ──────────────────── HF_HUB_DOWNLOAD_PATCH_TARGET = ( "darts.models.components.huggingface_connector.hf_hub_download" @@ -74,6 +103,43 @@ def _forecast_quantiles(self, context, prediction_length: int, **_kwargs): return quantiles, mean +# ── T0 tiny model ─────────────────────────────────────────────────────────── +# T0Model loads config.json + model.safetensors through Darts' shared HuggingFaceConnector, +# like the other foundation models. We build a small real model, save it to a temp dir, and +# point tests at it via ``local_dir`` (the gated t0-alpha weights are never downloaded). + + +@functools.lru_cache(maxsize=1) +def tiny_t0_dir() -> str: + """Build a tiny T0 model, save it (config.json + model.safetensors) to a temp dir, and return + the path β€” so tests load it via ``local_dir`` exactly like the other foundation models.""" + from t0 import T0Config, T0Forecaster + + directory = tempfile.mkdtemp(prefix="darts_tiny_t0_") + model = T0Forecaster.from_config( + T0Config( + embed_dim=64, + num_layers=4, + num_heads=8, + mlp_hidden_dim=128, + patch_size=8, + group_every_n=2, + dropout=0.0, + quantile_levels=(0.1, 0.25, 0.5, 0.75, 0.9), + scaler_use_arcsinh=True, + ) + ) + model.save_pretrained(directory) + return directory + + +def tiny_t0(): + """The tiny model the wrapper loads from ``tiny_t0_dir`` (identical weights).""" + from t0 import T0Forecaster + + return T0Forecaster.from_pretrained(tiny_t0_dir()) + + # ── TimesFM 2.5 tiny model ───────────────────────────────────────────────── # # The production ``_TimesFM2p5Module`` uses a hardcoded diff --git a/darts/tests/models/forecasting/test_foundation.py b/darts/tests/models/forecasting/test_foundation.py index ce920b6b40..81a6238631 100644 --- a/darts/tests/models/forecasting/test_foundation.py +++ b/darts/tests/models/forecasting/test_foundation.py @@ -8,7 +8,12 @@ import pytest from darts import TimeSeries, concatenate -from darts.tests.conftest import TIREX_AVAILABLE, TORCH_AVAILABLE, tfm_kwargs +from darts.tests.conftest import ( + T0_AVAILABLE, + TIREX_AVAILABLE, + TORCH_AVAILABLE, + tfm_kwargs, +) from darts.utils.likelihood_models import QuantileRegression from darts.utils.timeseries_generation import linear_timeseries @@ -18,7 +23,13 @@ allow_module_level=True, ) -from darts.models import Chronos2Model, PatchTSTFMModel, TimesFM2p5Model, TiRexModel +from darts.models import ( + Chronos2Model, + PatchTSTFMModel, + T0Model, + TimesFM2p5Model, + TiRexModel, +) from darts.tests.models.forecasting.foundation_test_utils import ( CHRONOS2_TINY_DIR, HF_HUB_DOWNLOAD_PATCH_TARGET, @@ -27,6 +38,7 @@ TiRexStub, mock_hf_hub_download, timesfm2p5_tiny_context, + tiny_t0_dir, ) @@ -427,6 +439,17 @@ def test_finetuning_misconfiguration(self, mock_method): ] if TIREX_AVAILABLE else [] + ) + + ( + [ + ( + T0Model, + "*decoder*", + {"local_dir": tiny_t0_dir()}, + ) + ] + if T0_AVAILABLE + else [] ), ) def test_finetuning_all_models(self, config): diff --git a/darts/tests/models/forecasting/test_t0.py b/darts/tests/models/forecasting/test_t0.py new file mode 100644 index 0000000000..db154ef901 --- /dev/null +++ b/darts/tests/models/forecasting/test_t0.py @@ -0,0 +1,214 @@ +from contextlib import contextmanager +from unittest.mock import patch + +import numpy as np +import pytest + +from darts.tests.conftest import T0_AVAILABLE, TORCH_AVAILABLE, tfm_kwargs + +if not TORCH_AVAILABLE: + pytest.skip( + f"Torch not available. {__name__} tests will be skipped.", + allow_module_level=True, + ) + +if not T0_AVAILABLE: + pytest.skip( + f"tfc-t0 not available. {__name__} tests will be skipped.", + allow_module_level=True, + ) + +import torch + +from darts import TimeSeries, concatenate +from darts.models import T0Model +from darts.tests.models.forecasting.foundation_test_utils import tiny_t0_dir +from darts.utils.likelihood_models import GaussianLikelihood, QuantileRegression +from darts.utils.timeseries_generation import ( + gaussian_timeseries, + linear_timeseries, + sine_timeseries, +) + +# `T0Model` rebuilds T0 from the Hub config and loads the weights through the connector; point it at +# a tiny local checkpoint (no gated download) and swap the rebuilt model for a stub. +_LOCAL = {"local_dir": tiny_t0_dir()} +_PATCH_T0_FROM_CONFIG = "darts.models.forecasting.t0_model.T0Forecaster.from_config" +_PATCH_LOAD_WEIGHTS = "darts.models.components.huggingface_connector.HuggingFaceConnector.load_model_weights" + + +@contextmanager +def _stub_t0(stub: "_StubT0Forecaster | None" = None): + """Build `stub` instead of the real T0 forecaster; weight loading becomes a no-op.""" + stub = _StubT0Forecaster() if stub is None else stub + with patch(_PATCH_T0_FROM_CONFIG, return_value=stub), patch(_PATCH_LOAD_WEIGHTS): + yield stub + + +class _StubForecast: + def __init__(self, quantiles: torch.Tensor): + self.quantiles = quantiles + + +class _StubT0Forecaster(torch.nn.Module): + """Stub emulating the `tfc-t0` ``T0Forecaster`` API used by the wrapper. + + ``predict(context, horizon, quantiles, future_covariates)`` returns a ``Forecast``-like object whose + ``quantiles`` is shaped ``(B, V, horizon, Q)`` β€” matching ``T0Forecaster`` for ndim-3 (multivariate) context. + """ + + def __init__(self): + super().__init__() + # a parameter so `next(self.parameters()).device` works like the real model + self._p = torch.nn.Parameter(torch.zeros(1)) + + def predict(self, context, horizon, quantiles, future_covariates=None): + assert torch.is_tensor(context) and context.ndim == 3 # (B, V, T) + batch, n_variates, _ = context.shape + n_q = len(quantiles) + if future_covariates is not None: + # covariates must span context + horizon + assert future_covariates.shape[0] == batch + assert future_covariates.shape[2] == context.shape[-1] + horizon + base = torch.arange(1, horizon + 1, dtype=torch.float32, device=context.device) + quantile_offsets = torch.tensor( + [float(q) - 0.5 for q in quantiles], device=context.device + ) + # (B, V, horizon, Q) + out = base.view(1, 1, horizon, 1) + quantile_offsets.view(1, 1, 1, n_q) + return _StubForecast(out.expand(batch, n_variates, horizon, n_q).contiguous()) + + +class TestT0Model: + np.random.seed(42) + + series = linear_timeseries(length=200, dtype=np.float32, column_name="A") + series_multi = concatenate( + [ + linear_timeseries(length=200, dtype=np.float32, column_name="A"), + sine_timeseries(length=200, dtype=np.float32, column_name="B"), + gaussian_timeseries(length=200, dtype=np.float32, column_name="C"), + ], + axis=1, + ) + cov = sine_timeseries(length=400, dtype=np.float32, column_name="cov") + + def test_creation(self): + # only QuantileRegression likelihood is supported + with pytest.raises(ValueError, match="Only QuantileRegression likelihood is"): + T0Model( + input_chunk_length=12, + output_chunk_length=6, + likelihood=GaussianLikelihood(), + **tfm_kwargs, + ) + + # fine-tuning is supported + model = T0Model( + input_chunk_length=12, + output_chunk_length=6, + enable_finetuning=True, + **tfm_kwargs, + ) + assert model.enable_finetuning is True + + def test_default(self): + model = T0Model( + input_chunk_length=24, output_chunk_length=12, **_LOCAL, **tfm_kwargs + ) + with _stub_t0(): + model.fit(self.series) + + # deterministic, single component + pred = model.predict(n=10, series=self.series) + assert isinstance(pred, TimeSeries) + assert len(pred) == 10 + assert pred.n_components == 1 + + # autoregressive prediction (n > output_chunk_length) + pred_ar = model.predict(n=20, series=self.series) + assert len(pred_ar) == 20 + + def test_probabilistic(self): + model = T0Model( + input_chunk_length=24, + output_chunk_length=12, + likelihood=QuantileRegression(quantiles=[0.1, 0.5, 0.9]), + **_LOCAL, + **tfm_kwargs, + ) + with _stub_t0(): + model.fit(self.series) + assert model.model_created + assert model.supports_probabilistic_prediction + + pred = model.predict( + n=6, series=self.series, predict_likelihood_parameters=True + ) + assert pred.n_components == 3 # 3 quantiles + + @pytest.mark.parametrize("probabilistic", [True, False]) + def test_multivariate(self, probabilistic: bool): + model = T0Model( + input_chunk_length=24, + output_chunk_length=8, + likelihood=( + QuantileRegression(quantiles=[0.1, 0.5, 0.9]) if probabilistic else None + ), + **_LOCAL, + **tfm_kwargs, + ) + with _stub_t0(): + model.fit(series=self.series_multi) + pred = model.predict(n=7, predict_likelihood_parameters=probabilistic) + assert len(pred) == 7 + if probabilistic: + assert pred.n_components == 9 # 3 variables x 3 quantiles + else: + assert pred.n_components == 3 + + @pytest.mark.parametrize("which", ["future", "past", "both"]) + def test_covariates(self, which: str): + # past covariates are forecast jointly with the target and dropped from the output; + # future covariates are passed to T0's covariate branch ([B, F, context+horizon], asserted by the stub). + model = T0Model( + input_chunk_length=24, output_chunk_length=12, **_LOCAL, **tfm_kwargs + ) + past_cov = self.cov if which in ("past", "both") else None + future_cov = self.cov if which in ("future", "both") else None + + with _stub_t0(): + model.fit( + series=self.series, + past_covariates=past_cov, + future_covariates=future_cov, + ) + pred = model.predict( + n=12, + series=self.series, + past_covariates=past_cov, + future_covariates=future_cov, + ) + assert isinstance(pred, TimeSeries) + assert len(pred) == 12 + # only the single target component is returned, never the past covariate + assert pred.n_components == 1 + + def test_multiple_series(self): + model = T0Model( + input_chunk_length=24, output_chunk_length=8, **_LOCAL, **tfm_kwargs + ) + series_multi_2 = concatenate( + [ + linear_timeseries(length=150, dtype=np.float32, column_name="A"), + sine_timeseries(length=150, dtype=np.float32, column_name="B"), + gaussian_timeseries(length=150, dtype=np.float32, column_name="C"), + ], + axis=1, + ) + with _stub_t0(): + model.fit(series=[self.series_multi, series_multi_2]) + pred = model.predict(n=5, series=[self.series_multi, series_multi_2]) + assert isinstance(pred, list) and len(pred) == 2 + assert all(len(p) == 5 for p in pred) + assert all(p.n_components == 3 for p in pred) diff --git a/docs/source/index.rst b/docs/source/index.rst index 2da67e15c2..4d156dbaf8 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -695,6 +695,12 @@ Our regression models are designed to predict continuous numerical values, makin - βœ… βœ… - βœ… - `PatchTST-FM paper `_, `PatchTST-FM Github `_ + * - `T0Model `_ + - βœ… βœ… + - βœ… βœ… πŸ”΄ + - βœ… βœ… + - βœ… + - `T0 model card `_, `tfc-t0 GitHub `_ * - **Ensemble Models** (`GlobalForecastingModel `_): Model support is dependent on ensembled forecasting models and the ensemble model itself - - diff --git a/examples/27-Torch-and-Foundation-Model-Fine-Tuning-examples.ipynb b/examples/27-Torch-and-Foundation-Model-Fine-Tuning-examples.ipynb index 3ce2e6bf61..9610a004bd 100644 --- a/examples/27-Torch-and-Foundation-Model-Fine-Tuning-examples.ipynb +++ b/examples/27-Torch-and-Foundation-Model-Fine-Tuning-examples.ipynb @@ -60,7 +60,7 @@ "from darts import set_option\n", "from darts.datasets import AirPassengersDataset, AusBeerDataset\n", "from darts.metrics import mae\n", - "from darts.models import Chronos2Model, TiDEModel\n", + "from darts.models import Chronos2Model, T0Model, TiDEModel\n", "from darts.utils.callbacks import TFMProgressBar\n", "\n", "warnings.filterwarnings(\"ignore\")\n", @@ -516,6 +516,111 @@ ");" ] }, + { + "cell_type": "markdown", + "id": "b1f4a0c7", + "metadata": {}, + "source": [ + "#### Another foundation model: T0\n", + "\n", + "The same `enable_finetuning` API works for every foundation model. Let's repeat the exercise with [T0](https://huggingface.co/theforecastingcompany/t0-alpha), The Forecasting Company's open-weights ~100M-parameter patch transformer, on the same beer production data.\n", + "\n", + "
\n", + " Getting T0: the weights are open and published on the Hugging Face Hub, but the repository is gated. Open the model card while signed in and accept the access conditions, then authenticate your environment with huggingface-cli login or by setting HF_TOKEN. Signing in on the website alone does not authenticate your Python environment. T0 also needs its optional package: pip install tfc-t0.\n", + "
\n", + "\n", + "
\n", + "\n", + "T0: Fine-tunable layers\n", + "\n", + "T0 is a patch transformer: a patch encoder, 24 transformer blocks, and a decoder that maps each position to the next patch's quantiles. The output-side decoder is the cheapest useful surface, and is the direct analogue of Chronos-2's `output_patch_embedding`:\n", + "\n", + "| Layer group | Trainable params | Notes |\n", + "|------------------------------------------|------------------|------------------------------------------|\n", + "| `*decoder*` | ~0.43M | Output head, a good default |\n", + "| `*decoder*` + last transformer block(s) | ~21M | More capacity, slower, easier to overfit |\n", + "| Full model (`enable_finetuning=True`) | ~101.6M | Most flexible, most expensive |\n", + "\n", + "Three T0-specific notes:\n", + "\n", + "1. **Patterns must match the full parameter name.** T0's weights live under the `t0.` prefix inside the Darts module, so use `\"*decoder*\"` rather than `\"decoder*\"`. If a pattern matches nothing, *every* parameter stays frozen and training fails with an autograd error. Check with `[n for n, p in model.model.named_parameters() if p.requires_grad]`.\n", + "2. **Fine-tuning runs a single forward pass**, so `output_chunk_length` plus `output_chunk_shift` must not exceed T0's `max_horizon`. Longer horizons still work for zero-shot `predict()`, which continues auto-regressively.\n", + "3. **T0 reads its context in patches of 32 steps.** With `input_chunk_length=24` below it sees less than one full patch, which is well outside the regime it was trained for. Expect a weak zero-shot baseline here and a correspondingly large gain from fine-tuning. A longer input window is the first thing to try on real data.\n", + "\n", + "Dropout is disabled automatically during fine-tuning, so a frozen backbone stays deterministic.\n", + "\n", + "
" + ] + }, + { + "cell_type": "code", + "id": "c2e5b1d8", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "t0_params = dict(\n", + " input_chunk_length=input_chunk_length,\n", + " output_chunk_length=output_chunk_length,\n", + " random_state=42,\n", + ")\n", + "model = T0Model(**t0_params)\n", + "\n", + "# as before, `fit()` only loads the model without fine-tuning\n", + "model.fit(series=train_beer)\n", + "\n", + "# predict\n", + "pred_beer = model.predict(n=output_chunk_length, series=val_beer)\n", + "\n", + "# plot\n", + "series_beer[-3 * output_chunk_length :].plot(label=\"Ground truth\")\n", + "pred_beer.plot(\n", + " label=\"Forecast\",\n", + " title=f\"Pre-trained T0; MAE {mae(series_beer, pred_beer):.2f}\",\n", + ");" + ] + }, + { + "cell_type": "code", + "id": "d3a6c2e9", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# fine-tune only the output decoder (\"*decoder*\" matches the `t0.decoder.*` parameters)\n", + "model_finetune = T0Model(\n", + " enable_finetuning={\"unfreeze\": [\"*decoder*\"]},\n", + " save_checkpoints=True,\n", + " model_name=\"t0\",\n", + " force_reset=True,\n", + " optimizer_kwargs={\"lr\": 1e-3},\n", + " pl_trainer_kwargs=dict(\n", + " check_val_every_n_epoch=5,\n", + " gradient_clip_val=1,\n", + " callbacks=[TFMProgressBar(enable_train_bar_only=True)],\n", + " ),\n", + " **t0_params,\n", + ")\n", + "\n", + "# fine-tune for 15 epochs\n", + "model_finetune.fit(\n", + " series=train_beer,\n", + " val_series=val_beer,\n", + " load_best=True,\n", + " epochs=15,\n", + ")\n", + "\n", + "# predict\n", + "pred_beer = model_finetune.predict(n=output_chunk_length, series=val_beer)\n", + "\n", + "# plot\n", + "series_beer[-3 * output_chunk_length :].plot(label=\"Ground truth\")\n", + "pred_beer.plot(\n", + " label=\"Forecast\",\n", + " title=f\"Fine-tuned T0; MAE {mae(series_beer, pred_beer):.2f}\",\n", + ");" + ] + }, { "cell_type": "markdown", "id": "996456e0", diff --git a/pyproject.toml b/pyproject.toml index 1e1857ce42..5960831d18 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -180,6 +180,7 @@ optional = [ "plotly>=6.5.2", "neuralforecast>=3.0.0", "tirex-ts>=1.4.0", + "tfc-t0>=0.4.0", ] release = [ "bump-my-version==1.3.0", From 7c79c59b192ed5bf56ecbb418de979e64d6508b7 Mon Sep 17 00:00:00 2001 From: Lucas Date: Wed, 16 Sep 2026 00:00:58 +0200 Subject: [PATCH 2/7] fix: remove duplicated t0 tiny model from test fixtures keep only the smaller model --- .../forecasting/foundation_test_utils.py | 37 ------------------- 1 file changed, 37 deletions(-) diff --git a/darts/tests/models/forecasting/foundation_test_utils.py b/darts/tests/models/forecasting/foundation_test_utils.py index 4167842f64..6d4f312121 100644 --- a/darts/tests/models/forecasting/foundation_test_utils.py +++ b/darts/tests/models/forecasting/foundation_test_utils.py @@ -103,43 +103,6 @@ def _forecast_quantiles(self, context, prediction_length: int, **_kwargs): return quantiles, mean -# ── T0 tiny model ─────────────────────────────────────────────────────────── -# T0Model loads config.json + model.safetensors through Darts' shared HuggingFaceConnector, -# like the other foundation models. We build a small real model, save it to a temp dir, and -# point tests at it via ``local_dir`` (the gated t0-alpha weights are never downloaded). - - -@functools.lru_cache(maxsize=1) -def tiny_t0_dir() -> str: - """Build a tiny T0 model, save it (config.json + model.safetensors) to a temp dir, and return - the path β€” so tests load it via ``local_dir`` exactly like the other foundation models.""" - from t0 import T0Config, T0Forecaster - - directory = tempfile.mkdtemp(prefix="darts_tiny_t0_") - model = T0Forecaster.from_config( - T0Config( - embed_dim=64, - num_layers=4, - num_heads=8, - mlp_hidden_dim=128, - patch_size=8, - group_every_n=2, - dropout=0.0, - quantile_levels=(0.1, 0.25, 0.5, 0.75, 0.9), - scaler_use_arcsinh=True, - ) - ) - model.save_pretrained(directory) - return directory - - -def tiny_t0(): - """The tiny model the wrapper loads from ``tiny_t0_dir`` (identical weights).""" - from t0 import T0Forecaster - - return T0Forecaster.from_pretrained(tiny_t0_dir()) - - # ── TimesFM 2.5 tiny model ───────────────────────────────────────────────── # # The production ``_TimesFM2p5Module`` uses a hardcoded From 1962b3cf006e7b3eefa37a388b4bb6e10162c5d8 Mon Sep 17 00:00:00 2001 From: dennisbader Date: Fri, 18 Sep 2026 14:08:52 +0200 Subject: [PATCH 3/7] update dependency --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f5278526b0..8cc98be964 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,7 +79,7 @@ all = ["darts[torch,notorch]"] [tool.uv] # Cooldown period for PyPI releases to mitigate supply chain attacks. exclude-newer = "7 days" -exclude-newer-package = {} +exclude-newer-package = { tfc-t0 = "0 days" } # Must resolve dependencies for major platforms in uv.lock required-environments = [ "sys_platform == 'darwin'", @@ -180,7 +180,7 @@ optional = [ "plotly>=6.5.2", "neuralforecast>=3.0.0", "tirex-ts>=1.4.0", - "tfc-t0>=0.4.0", + "tfc-t0>=0.5.0", ] release = [ "bump-my-version==1.3.0", From 11a49141aa3d560ef492ca26224ca48c5b12408d Mon Sep 17 00:00:00 2001 From: dennisbader Date: Fri, 18 Sep 2026 14:55:11 +0200 Subject: [PATCH 4/7] make model work with new torch module I/O --- INSTALL.md | 4 +- README.md | 2 +- darts/models/forecasting/t0_model.py | 72 ++++++++++++----------- darts/tests/models/forecasting/test_t0.py | 6 +- 4 files changed, 43 insertions(+), 41 deletions(-) diff --git a/INSTALL.md b/INSTALL.md index ad4813bf63..f5638f92a4 100644 --- a/INSTALL.md +++ b/INSTALL.md @@ -22,7 +22,7 @@ Some models have additional dependencies that are not included in the `all` inst |-----------------------|-----------------------| | `NeuralForecastModel` | neuralforecast>=3.0.0 | | `TiRexModel` | tirex-ts>=1.4.0 | -| `T0Model` | tfc-t0>=0.4.0 | +| `T0Model` | tfc-t0>=0.5.0 | Some optional integrations also require additional dependencies: @@ -60,7 +60,7 @@ Some models have dependencies not available on conda-forge. To use them, you nee | Model | Dependencies | |-----------------------|-----------------------| | `TiRexModel` | tirex-ts>=1.4.0 | -| `T0Model` | tfc-t0>=0.4.0 | +| `T0Model` | tfc-t0>=0.5.0 | ## Other Information diff --git a/README.md b/README.md index d471a1d8ce..4856c0182e 100644 --- a/README.md +++ b/README.md @@ -284,7 +284,7 @@ Here's a breakdown of the forecasting models currently implemented in Darts. Our | [TimesFM3Model](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.timesfm3_model.html#darts.models.forecasting.timesfm3_model.TimesFM3Model) | [TimesFM 1.0 paper](https://arxiv.org/abs/2310.10688), [TimesFM 3.0 model card](https://huggingface.co/google/timesfm-3.0-pytorch) | βœ… βœ… | βœ… βœ… πŸ”΄ | βœ… βœ… | βœ… | | [TiRexModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.tirex_model.html#darts.models.forecasting.tirex_model.TiRexModel) | [TiRex paper](https://arxiv.org/abs/2505.23719), [TiRex GitHub](https://github.com/NX-AI/tirex) | βœ… βœ… | πŸ”΄ πŸ”΄ πŸ”΄ | βœ… βœ… | βœ… | | [PatchTSTFMModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.patchtst_fm_model.html#darts.models.forecasting.patchtst_fm_model.PatchTSTFMModel) | [PatchTST-FM paper](https://arxiv.org/abs/2602.06909), [PatchTST-FM GitHub](https://github.com/ibm-granite/granite-tsfm) | βœ… βœ… | πŸ”΄ πŸ”΄ πŸ”΄ | βœ… βœ… | βœ… | -| [T0Model](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.t0_model.html#darts.models.forecasting.t0_model.T0Model) | [T0 model card](https://huggingface.co/theforecastingcompany/t0-alpha), [tfc-t0 GitHub](https://github.com/theforecastingcompany/tfc-t0) | βœ… βœ… | βœ… βœ… πŸ”΄ | βœ… βœ… | βœ… | +| [T0Model](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.t0_model.html#darts.models.forecasting.t0_model.T0Model) | [T0 model card](https://huggingface.co/theforecastingcompany/t0-alpha), [tfc-t0 GitHub](https://github.com/theforecastingcompany/tfc-t0) | βœ… βœ… | βœ… βœ… πŸ”΄ | βœ… βœ… | βœ… | | **Ensemble Models**
([GlobalForecastingModel](https://unit8co.github.io/darts/userguide/covariates.html#global-forecasting-models-gfms)): Model support is dependent on ensembled forecasting models and the ensemble model itself | | | | | | | [NaiveEnsembleModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.naive_ensemble_model.html#darts.models.forecasting.naive_ensemble_model.NaiveEnsembleModel) | | βœ… βœ… | βœ… βœ… βœ… | βœ… βœ… | βœ… | | [RegressionEnsembleModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.regression_ensemble_model.html#darts.models.forecasting.regression_ensemble_model.RegressionEnsembleModel) | | βœ… βœ… | βœ… βœ… βœ… | βœ… βœ… | βœ… | diff --git a/darts/models/forecasting/t0_model.py b/darts/models/forecasting/t0_model.py index fba6dc4ad1..fe49f48a3f 100644 --- a/darts/models/forecasting/t0_model.py +++ b/darts/models/forecasting/t0_model.py @@ -29,6 +29,7 @@ from darts.models.components.huggingface_connector import HuggingFaceConnector from darts.models.forecasting.foundation_model import FoundationModel from darts.models.forecasting.pl_forecasting_module import PLForecastingModule +from darts.utils.data import ModuleStage, PLModuleOutput, TorchTrainingSample from darts.utils.data.torch_datasets.utils import PLModuleInput from darts.utils.likelihood_models.torch import QuantileRegression @@ -123,37 +124,36 @@ def forward(self, x_in: PLModuleInput, *args, **kwargs): # Qp: pre-trained quantiles (returned during fine-tuning) # N: likelihood quantiles (user-specified, 1 if deterministic) - # `x_past`: (B, L, C + P + F) stack of [past_target, past_covariates, historic_future_covariates]; - # `x_future`: (B, T, F) future covariates, or None. - x_past, x_future, _, _ = x_in - batch_size = x_past.shape[0] - - # Past covariates are forecast jointly with the target (T0 is variate-agnostic) and dropped from the - # output. Future covariates are the trailing columns of `x_past` (their historic part) and are passed to - # T0's covariate branch instead. `x_future` width gives the number of future covariates. - n_future_covs = x_future.shape[-1] if x_future is not None else 0 - n_context = x_past.shape[-1] - n_future_covs - + # Past covariates are forecast jointly with the target (T0 is variate-agnostic) and dropped from the output. # context: (B, V, L) - context = x_past[:, :, :n_context].transpose(1, 2) - - # T0 expects covariates over context + horizon. Re-assemble them from the historic part (in `x_past`) - # and the future chunk (`x_future`); the `output_chunk_shift` gap is left NaN (T0 treats NaN as missing). - future_covariates = None - if n_future_covs > 0: - historic = x_past[:, :, n_context:] # (B, L, F) - future = torch.full( - (batch_size, self.future_len, n_future_covs), - torch.nan, - device=x_past.device, - dtype=x_past.dtype, - ) - if x_future is not None: - future[:, -(self.output_chunk_length or 0) :, :] = x_future - # (B, L + H, F) -> (B, F, L + H) - future_covariates = torch.cat([historic, future], dim=1).transpose(1, 2) + context = torch.cat( + [el for el in [x_in.past_target, x_in.past_covariates] if el is not None], + dim=2, + ).transpose(1, 2) + batch_size, n_context, _ = context.shape + + # Future covariates and their historic part are passed to T0's covariate branch instead. + # T0 expects covariates over context (historic) + horizon (future); + # the `output_chunk_shift` gap is left NaN (T0 treats NaN as missing). + historic_future_covariates = x_in.historic_future_covariates + future_covariates = x_in.future_covariates + if historic_future_covariates is not None and future_covariates is not None: + if self.output_chunk_shift == 0: + future_covariates = x_in.concatenate_future_along_time() + else: + n_future_covs = future_covariates.shape[-1] + gap = torch.full( + (batch_size, self.output_chunk_shift, n_future_covs), + torch.nan, + device=context.device, + dtype=context.dtype, + ) + future_covariates = torch.cat( + [historic_future_covariates, gap, future_covariates], dim=1 + ) + future_covariates = future_covariates.transpose(1, 2) - if self.training and self._enable_finetuning: + if x_in.stage is ModuleStage.TRAIN: # scale the input, run one forward pass, and rescale the predictions patch_size = self.t0.patch_size context_length = context.shape[-1] @@ -188,21 +188,23 @@ def forward(self, x_in: PLModuleInput, *args, **kwargs): quantiles = self.t0.predict( context, horizon=self.future_len, - quantiles=user_q, + quantile_levels=user_q, future_covariates=future_covariates, ).quantiles # drop the past-covariate variates, keep targets: (B, V, H, N) -> (B, C, H, N) quantiles = quantiles[:, : self.n_targets] # (B, C, H, N) -> (B, H, C, N) -> slice output shift -> (B, T, C, N) - return quantiles.permute(0, 2, 1, 3)[:, self.output_chunk_shift :, :, :] + quantiles = quantiles.permute(0, 2, 1, 3)[:, self.output_chunk_shift :, :, :] + return PLModuleOutput(prediction=quantiles) def _compute_loss(self, output, target, criterion, sample_weight): - if self.training and self._enable_finetuning: + if self.training: # compute loss on pre-trained quantiles return self._finetuning_likelihood.compute_loss( - output, target, sample_weight + output.prediction, target, sample_weight ) - return super()._compute_loss(output, target, criterion, sample_weight) + else: + return super()._compute_loss(output, target, criterion, sample_weight) class T0Model(FoundationModel): @@ -413,7 +415,7 @@ def supports_past_covariates(self) -> bool: def supports_future_covariates(self) -> bool: return True - def _create_model(self, train_sample) -> PLForecastingModule: + def _create_model(self, train_sample: TorchTrainingSample) -> PLForecastingModule: # enable_finetuning is injected into pl_module_params by the base class; # _T0Module accepts it as an explicit parameter and converts dict form to bool return _T0Module( diff --git a/darts/tests/models/forecasting/test_t0.py b/darts/tests/models/forecasting/test_t0.py index db154ef901..12bd272684 100644 --- a/darts/tests/models/forecasting/test_t0.py +++ b/darts/tests/models/forecasting/test_t0.py @@ -62,17 +62,17 @@ def __init__(self): # a parameter so `next(self.parameters()).device` works like the real model self._p = torch.nn.Parameter(torch.zeros(1)) - def predict(self, context, horizon, quantiles, future_covariates=None): + def predict(self, context, horizon, quantile_levels, future_covariates=None): assert torch.is_tensor(context) and context.ndim == 3 # (B, V, T) batch, n_variates, _ = context.shape - n_q = len(quantiles) + n_q = len(quantile_levels) if future_covariates is not None: # covariates must span context + horizon assert future_covariates.shape[0] == batch assert future_covariates.shape[2] == context.shape[-1] + horizon base = torch.arange(1, horizon + 1, dtype=torch.float32, device=context.device) quantile_offsets = torch.tensor( - [float(q) - 0.5 for q in quantiles], device=context.device + [float(q) - 0.5 for q in quantile_levels], device=context.device ) # (B, V, horizon, Q) out = base.view(1, 1, horizon, 1) + quantile_offsets.view(1, 1, 1, n_q) From ebace3b35ea42bb37f7150127d3d412f2a3bdf7d Mon Sep 17 00:00:00 2001 From: dennisbader Date: Fri, 18 Sep 2026 15:52:48 +0200 Subject: [PATCH 5/7] update changelog --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ff3af4d75..874990fdae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,14 +11,14 @@ but cannot always guarantee backwards compatibility. Changes that may **break co **Improved** -- πŸš€ Added new forecasting model `T0Model` : [The Forecasting Company's open-weights ~100M-parameter foundation model](https://huggingface.co/theforecastingcompany/t0-alpha) for zero-shot forecasting. It supports univariate, multivariate, and multiple time series as well as past and future covariates, without training, and can output deterministic or probabilistic forecasts. It can also be fine-tuned (full or partial) with `enable_finetuning`. [#3142](https://github.com/unit8co/darts/pull/3142) by [Geoffrey NΓ©giar](https://github.com/GeoffNN), [Huikan Xiang](https://github.com/huikan-tfc) and [Lucas Meyer](https://github.com/LTMeyer). - Improvements to `TorchForecastingModel` : [#3204](https://github.com/unit8co/darts/pull/3204) by [Dennis Bader](https://github.com/dennisbader). - πŸš€πŸš€ ONNX export and inference are substantially more capable: train a model in PyTorch, export it once, then run forecasts in a lightweight environment with only ONNX Runtime, NumPy, and Darts β€” no PyTorch required. `run_onnx_prediction()` mirrors `predict()` (including auto-regressive horizons and RNN warm-up); `RNNModel` and probabilistic models are now supported as well. - πŸ”΄ Removed `darts.utils.onnx_utils`; use `darts.utils.onnx.inference` instead. Custom ONNX loops should load graph metadata via `OnnxModelSpec.from_session()`. - Custom PyTorch datasets and Lightning modules are easier to read, extend, and debug: samples use named fields (`past_target`, `future_covariates`, ...) instead of positional tuples, modules receive each feature as a separate tensor rather than one concatenated input, and recurrent state is returned in a structured output. Models saved with previous Darts versions continue to load for inference. - πŸ”΄ Custom `TorchTrainingDataset` / `TorchInferenceDataset` implementations must return `TorchTrainingSample` / `TorchInferenceSample`. - πŸ”΄ Custom module `forward()` methods must accept `PLModuleInput` and return `PLModuleOutput`. -- πŸš€πŸš€ Added new forecasting model `TimesFM3Model` : Google's pre-trained 330M-parameter foundation model for zero-shot forecasting. Unlike previous versions, it natively supports multivariate time series, past covariates, and future covariates, and can output deterministic or probabilistic forecasts without training. The TimesFM 3.0 pre-trained weights are non-commercial: users must accept the license with `accept_license=True` when creating the model. [#3199](https://github.com/unit8co/darts/pull/3199) by [JuanCruzC97](https://github.com/JuanCruzC97). +- πŸš€ Added new forecasting model `TimesFM3Model` : Google's pre-trained 330M-parameter foundation model for zero-shot forecasting. Unlike previous versions, it natively supports multivariate time series, past covariates, and future covariates, and can output deterministic or probabilistic forecasts without training. The TimesFM 3.0 pre-trained weights are non-commercial: users must accept the license with `accept_license=True` when creating the model. [#3199](https://github.com/unit8co/darts/pull/3199) by [JuanCruzC97](https://github.com/JuanCruzC97). +- πŸš€ Added new forecasting model `T0Model` : The Forecasting Company's open-weights ~100M-parameter foundation model (see [here](https://huggingface.co/theforecastingcompany/t0-alpha)) for zero-shot forecasting. It supports univariate, multivariate, and multiple time series as well as past and future covariates, without training, and can output deterministic or probabilistic forecasts. It can also be fine-tuned (full or partial) with `enable_finetuning`. [#3142](https://github.com/unit8co/darts/pull/3142) by [Geoffrey NΓ©giar](https://github.com/GeoffNN), [Huikan Xiang](https://github.com/huikan-tfc) and [Lucas Meyer](https://github.com/LTMeyer). - `FittableAnomalyScorer.fit_from_prediction()` now returns the fitted scorer object similar to `fit()`. [#3202](https://github.com/unit8co/darts/pull/3202) by [Venish Paneliya](https://github.com/VenishPaneliya). - Calling `ForecastingModel.historical_forecasts()` with a `start` value that is later than what is forecastable given the supplied covariates now raises an informative exception. [#3207](https://github.com/unit8co/darts/pull/3207) by [Dennis Bader](https://github.com/dennisbader). - Calling `ForecastingModel.gridsearch()` with a sequence of `TimeSeries` now raises an informative exception. [#3191](https://github.com/unit8co/darts/pull/3191) by [Geovanny Basantes](https://github.com/COMPUMAX-EC). From 06a02bd51647dcec68b8709723a8aa7879d06fcb Mon Sep 17 00:00:00 2001 From: dennisbader Date: Fri, 18 Sep 2026 16:18:57 +0200 Subject: [PATCH 6/7] add missing tests --- darts/tests/models/forecasting/test_t0.py | 93 +++++++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/darts/tests/models/forecasting/test_t0.py b/darts/tests/models/forecasting/test_t0.py index 12bd272684..7010e34bc9 100644 --- a/darts/tests/models/forecasting/test_t0.py +++ b/darts/tests/models/forecasting/test_t0.py @@ -79,6 +79,13 @@ def predict(self, context, horizon, quantile_levels, future_covariates=None): return _StubForecast(out.expand(batch, n_variates, horizon, n_q).contiguous()) +def series_with_nans(series: TimeSeries, start: int, end: int | None) -> TimeSeries: + """Returns a copy of `series` with the values in [start, end) set to NaN.""" + values = series.values(copy=True).astype(np.float32) + values[start:end, :] = np.nan + return TimeSeries.from_times_and_values(series.time_index, values) + + class TestT0Model: np.random.seed(42) @@ -118,6 +125,8 @@ def test_default(self): ) with _stub_t0(): model.fit(self.series) + assert model.model_created + assert not model.supports_probabilistic_prediction # deterministic, single component pred = model.predict(n=10, series=self.series) @@ -127,7 +136,9 @@ def test_default(self): # autoregressive prediction (n > output_chunk_length) pred_ar = model.predict(n=20, series=self.series) + assert isinstance(pred_ar, TimeSeries) assert len(pred_ar) == 20 + assert pred_ar.n_components == 1 def test_probabilistic(self): model = T0Model( @@ -145,8 +156,21 @@ def test_probabilistic(self): pred = model.predict( n=6, series=self.series, predict_likelihood_parameters=True ) + assert isinstance(pred, TimeSeries) + assert len(pred) == 6 assert pred.n_components == 3 # 3 quantiles + # probabilistic model allows autoregressive predictions (8 > 6) + pred_ar = model.predict( + n=14, + series=self.series, + num_samples=10, + ) + assert isinstance(pred_ar, TimeSeries) + assert len(pred_ar) == 14 + assert pred_ar.n_components == 1 # sampling yields single component + assert pred_ar.n_samples == 10 + @pytest.mark.parametrize("probabilistic", [True, False]) def test_multivariate(self, probabilistic: bool): model = T0Model( @@ -194,6 +218,75 @@ def test_covariates(self, which: str): # only the single target component is returned, never the past covariate assert pred.n_components == 1 + def test_missing_values(self): + """NaNs in target and covariates are handled via the masking logic of the + ported `decode()`, instead of the linear interpolation applied by the + upstream `TimesFM3Forecaster.predict_batch()`. Predictions must not + contain NaNs for any of the missing value locations. + """ + + def make_model() -> T0Model: + return T0Model( + input_chunk_length=8, + output_chunk_length=4, + **_LOCAL, + **tfm_kwargs, + ) + + # NaNs inside the target series, incl. autoregressive prediction + series_nan = series_with_nans(self.series, 20, 26) + model = make_model() + model.fit(series=series_nan) + pred = model.predict(n=6, series=series_nan) + assert isinstance(pred, TimeSeries) + assert not np.isnan(pred.all_values(copy=False)).any() + + # NaNs at the end of the target series (trailing missing values) + series_trailing_nan = series_with_nans(self.series, len(self.series) - 3, None) + model = make_model() + model.fit(series=series_trailing_nan) + pred = model.predict(n=4, series=series_trailing_nan) + assert isinstance(pred, TimeSeries) + assert not np.isnan(pred.all_values(copy=False)).any() + + # NaNs in the past covariates + past_cov_nan = series_with_nans(self.cov, 5, 10) + model = make_model() + model.fit(series=self.series, past_covariates=past_cov_nan) + pred = model.predict(n=4, series=self.series, past_covariates=past_cov_nan) + assert isinstance(pred, TimeSeries) + assert not np.isnan(pred.all_values(copy=False)).any() + + # NaNs in the future covariates + future_cov_nan = series_with_nans(self.cov, 204, 208) + model = make_model() + model.fit(series=self.series, future_covariates=future_cov_nan) + pred = model.predict(n=4, series=self.series, future_covariates=future_cov_nan) + assert isinstance(pred, TimeSeries) + assert not np.isnan(pred.all_values(copy=False)).any() + + def test_output_chunk_shift(self): + """With `output_chunk_shift`, predictions start after the shifted gap and + auto-regressive prediction (`n > output_chunk_length`) is not allowed.""" + model = T0Model( + input_chunk_length=8, + output_chunk_length=4, + output_chunk_shift=2, + **_LOCAL, + **tfm_kwargs, + ) + model.fit(series=self.series, future_covariates=self.cov) + pred = model.predict(n=4, series=self.series, future_covariates=self.cov) + assert isinstance(pred, TimeSeries) + assert len(pred) == 4 + # predictions start `output_chunk_shift + 1` steps after the end of the series + assert pred.start_time() == self.series.end_time() + self.series.freq * 3 + assert not np.isnan(pred.all_values(copy=False)).any() + + # auto-regression is not allowed with an output chunk shift + with pytest.raises(ValueError, match="output_chunk_shift > 0"): + model.predict(n=5, series=self.series, future_covariates=self.cov) + def test_multiple_series(self): model = T0Model( input_chunk_length=24, output_chunk_length=8, **_LOCAL, **tfm_kwargs From 6e79b18ba7c42b2cb108696b4b88c840f2b481d6 Mon Sep 17 00:00:00 2001 From: dennisbader Date: Mon, 21 Sep 2026 11:58:15 +0200 Subject: [PATCH 7/7] update docs --- README.md | 2 +- .../components/huggingface_connector.py | 10 +- darts/models/forecasting/t0_model.py | 281 +++++++++++------- docs/source/index.rst | 2 +- ...oundation-Model-Fine-Tuning-examples.ipynb | 107 +------ 5 files changed, 187 insertions(+), 215 deletions(-) diff --git a/README.md b/README.md index 4856c0182e..55b73c5e46 100644 --- a/README.md +++ b/README.md @@ -284,7 +284,7 @@ Here's a breakdown of the forecasting models currently implemented in Darts. Our | [TimesFM3Model](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.timesfm3_model.html#darts.models.forecasting.timesfm3_model.TimesFM3Model) | [TimesFM 1.0 paper](https://arxiv.org/abs/2310.10688), [TimesFM 3.0 model card](https://huggingface.co/google/timesfm-3.0-pytorch) | βœ… βœ… | βœ… βœ… πŸ”΄ | βœ… βœ… | βœ… | | [TiRexModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.tirex_model.html#darts.models.forecasting.tirex_model.TiRexModel) | [TiRex paper](https://arxiv.org/abs/2505.23719), [TiRex GitHub](https://github.com/NX-AI/tirex) | βœ… βœ… | πŸ”΄ πŸ”΄ πŸ”΄ | βœ… βœ… | βœ… | | [PatchTSTFMModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.patchtst_fm_model.html#darts.models.forecasting.patchtst_fm_model.PatchTSTFMModel) | [PatchTST-FM paper](https://arxiv.org/abs/2602.06909), [PatchTST-FM GitHub](https://github.com/ibm-granite/granite-tsfm) | βœ… βœ… | πŸ”΄ πŸ”΄ πŸ”΄ | βœ… βœ… | βœ… | -| [T0Model](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.t0_model.html#darts.models.forecasting.t0_model.T0Model) | [T0 model card](https://huggingface.co/theforecastingcompany/t0-alpha), [tfc-t0 GitHub](https://github.com/theforecastingcompany/tfc-t0) | βœ… βœ… | βœ… βœ… πŸ”΄ | βœ… βœ… | βœ… | +| [T0Model](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.t0_model.html#darts.models.forecasting.t0_model.T0Model) | [T0 report](https://www.theforecastingcompany.com/papers/files/t0-technical-report.pdf), [tfc-t0 GitHub](https://github.com/theforecastingcompany/tfc-t0) | βœ… βœ… | βœ… βœ… πŸ”΄ | βœ… βœ… | βœ… | | **Ensemble Models**
([GlobalForecastingModel](https://unit8co.github.io/darts/userguide/covariates.html#global-forecasting-models-gfms)): Model support is dependent on ensembled forecasting models and the ensemble model itself | | | | | | | [NaiveEnsembleModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.naive_ensemble_model.html#darts.models.forecasting.naive_ensemble_model.NaiveEnsembleModel) | | βœ… βœ… | βœ… βœ… βœ… | βœ… βœ… | βœ… | | [RegressionEnsembleModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.regression_ensemble_model.html#darts.models.forecasting.regression_ensemble_model.RegressionEnsembleModel) | | βœ… βœ… | βœ… βœ… βœ… | βœ… βœ… | βœ… | diff --git a/darts/models/components/huggingface_connector.py b/darts/models/components/huggingface_connector.py index d1b48b0344..dab10b50d7 100644 --- a/darts/models/components/huggingface_connector.py +++ b/darts/models/components/huggingface_connector.py @@ -43,18 +43,18 @@ def __init__( Parameters ---------- model_name - The HuggingFace repository name where the model is stored, e.g., "amazon/chronos-2". + The HuggingFace repository name where the model is stored, e.g., "amazon/chronos-2". model_revision - The revision of the model in the HuggingFace repository. Must be a branch name, tag name, or commit hash. - If not provided, the default branch and the latest commit will be used. + The revision of the model in the HuggingFace repository. Must be a branch name, tag name, or commit hash. + If not provided, the default branch and the latest commit will be used. local_dir Optional local directory to load the pre-downloaded model. If specified and the directory is empty, the model will be downloaded from HuggingFace Hub and saved to this directory. Default is ``None``, which will use a cache directory managed by ``huggingface_hub`` instead. config_file - The name of the configuration file. Default is "config.json". + The name of the configuration file. Default is "config.json". model_file - The name of the model weight file. Default is "model.safetensors". + The name of the model weight file. Default is "model.safetensors". """ if local_dir is not None: local_dir_path = Path(local_dir) diff --git a/darts/models/forecasting/t0_model.py b/darts/models/forecasting/t0_model.py index fe49f48a3f..fa21c1b407 100644 --- a/darts/models/forecasting/t0_model.py +++ b/darts/models/forecasting/t0_model.py @@ -6,8 +6,10 @@ multivariate series, it supports past and future covariates. The weights are published on the Hugging Face Hub at `theforecastingcompany/t0-alpha -`__. The repository is gated: request access on -the model page and authenticate (``huggingface-cli login``, or set ``HF_TOKEN``) before first use. +`__ and `theforecastingcompany/t0-beta +`__ (more checkpoints might be available in the future). +The repository is gated: request access on the model page and authenticate (``hf auth login`` via CLI or set +``HF_TOKEN``) before first use. For detailed examples and tutorials, see: @@ -37,26 +39,37 @@ class _T0Module(PLForecastingModule): - """PyTorch Lightning module wrapping a pre-loaded T0 forecaster. - - Targets and past covariates are forecast jointly, and past-covariate predictions are dropped. - Future covariates are mapped to T0's ``[batch, n_covariates, context + horizon]`` format. - - Fine-tuning runs T0's forward pass over all pre-trained quantiles; prediction returns only the - user-specified ones. - """ - def __init__( self, hub_model_name: str, hub_model_revision: str | None, local_dir: str | os.PathLike | None, - all_quantiles: tuple[float, ...], - enable_finetuning: bool | dict = False, **kwargs, ): + """PyTorch Lightning module wrapping a pre-loaded T0 forecaster. + + Targets and past covariates are forecast jointly, and past-covariate predictions are dropped. + Future covariates are mapped to T0's ``[batch, n_covariates, context + horizon]`` format. + + Fine-tuning runs T0's forward pass over all pre-trained quantiles; prediction returns only the + user-specified ones. + + Parameters + ---------- + hub_model_name + The HuggingFace repository name where the model is stored. + hub_model_revision + The revision of the model in the HuggingFace repository. Must be a branch name, tag name, or commit hash. + If not provided, the default branch and the latest commit will be used. + local_dir + Optional local directory to load the pre-downloaded model. If specified and the directory is empty, the + model will be downloaded from HuggingFace Hub and saved to this directory. Default is ``None``, which will + use a cache directory managed by ``huggingface_hub`` instead. + **kwargs + Additional arguments passed to `PLForecastingModule`. + """ + enable_finetuning = kwargs.pop("enable_finetuning", False) super().__init__(**kwargs) - self._enable_finetuning = bool(enable_finetuning) # rebuild T0 from the Hub config, then load the weights into it connector = HuggingFaceConnector( @@ -72,45 +85,39 @@ def __init__( } if "quantile_levels" in config_kwargs: config_kwargs["quantile_levels"] = tuple(config_kwargs["quantile_levels"]) - if self._enable_finetuning: + if enable_finetuning: # set dropout to 0, which would otherwise inject noise into a frozen backbone config_kwargs["dropout"] = 0.0 - self.t0: T0Forecaster = T0Forecaster.from_config(T0Config(**config_kwargs)) - connector.load_model_weights(self.t0) + config = T0Config(**config_kwargs) - # switch the model to train mode if needed. - self.t0.train(self._enable_finetuning) + self.t0: T0Forecaster = T0Forecaster.from_config(config) self.future_len = (self.output_chunk_length or 0) + self.output_chunk_shift + if enable_finetuning and self.future_len > self.t0.max_horizon: + raise_log( + ValueError( + f"T0 fine-tuning only supports up to {self.t0.max_horizon} steps got {self.future_len}" + ), + ) - if self._enable_finetuning: - if self.future_len > self.t0.max_horizon: - raise_log( - ValueError( - f"T0 fine-tuning only supports up to {self.t0.max_horizon} steps got {self.future_len}" - ), - ) + connector.load_model_weights(self.t0) + + # gather user-specified quantiles used at prediction time + self.user_quantiles: list[float] = ( + self.likelihood.quantiles + if isinstance(self.likelihood, QuantileRegression) + else [0.5] + ) + + if enable_finetuning: # loss is computed over all pre-trained quantiles to preserve the distribution; # user-specified quantiles are selected at prediction time - self._finetuning_likelihood = QuantileRegression(list(all_quantiles)) + self._finetuning_likelihood = QuantileRegression( + list(config.quantile_levels) + ) else: self._finetuning_likelihood = None - def forward(self, x_in: PLModuleInput, *args, **kwargs): - """Forward pass returning quantile predictions shaped ``(batch, time, n_targets, n_quantiles)``. - - During training with fine-tuning enabled, all pre-trained quantiles are returned for the loss. - At prediction time, only user-specified quantiles are returned. - - Parameters - ---------- - x_in - ``(x_past, x_future, x_static, future_target)`` the past, future, and static features, as well as - the future target. - *args - Positional arguments passed to the forward method. - **kwargs - Optional keyword arguments. - """ + def forward(self, x_in: PLModuleInput) -> PLModuleOutput: # Dimension notation in comments below: # B: batch size # L: input chunk length @@ -132,7 +139,7 @@ def forward(self, x_in: PLModuleInput, *args, **kwargs): ).transpose(1, 2) batch_size, n_context, _ = context.shape - # Future covariates and their historic part are passed to T0's covariate branch instead. + # Future covariates and their historic part are passed to T0's covariate branch. # T0 expects covariates over context (historic) + horizon (future); # the `output_chunk_shift` gap is left NaN (T0 treats NaN as missing). historic_future_covariates = x_in.historic_future_covariates @@ -179,16 +186,11 @@ def forward(self, x_in: PLModuleInput, *args, **kwargs): .unflatten(0, (batch_size, n_context)) ) else: - user_q: list[float] = ( - self.likelihood.quantiles - if isinstance(self.likelihood, QuantileRegression) - else [0.5] - ) # quantiles: (B, V, H, N) quantiles = self.t0.predict( context, horizon=self.future_len, - quantile_levels=user_q, + quantile_levels=self.user_quantiles, future_covariates=future_covariates, ).quantiles # drop the past-covariate variates, keep targets: (B, V, H, N) -> (B, C, H, N) @@ -208,16 +210,13 @@ def _compute_loss(self, output, target, criterion, sample_weight): class T0Model(FoundationModel): - # Quantile levels T0 was trained on. Other levels are interpolated, so any quantiles in (0, 1) are accepted. - _PRETRAINED_QUANTILES: tuple[float, ...] = (0.1, 0.25, 0.5, 0.75, 0.9) - def __init__( self, input_chunk_length: int, output_chunk_length: int, output_chunk_shift: int = 0, likelihood: QuantileRegression | None = None, - hub_model_name: str = "theforecastingcompany/t0-alpha", + hub_model_name: str = "theforecastingcompany/t0-beta", hub_model_revision: str | None = None, local_dir: str | os.PathLike | None = None, **kwargs, @@ -225,40 +224,49 @@ def __init__( """ T0 foundation model for zero-shot time series forecasting. - This is a Darts wrapper around The Forecasting Company's open-weights T0 model. The implementation delegates - all forecasting logic to the optional `tfc-t0 `_ package and loads the config - and weights through Darts' :class:`HuggingFaceConnector`, while exposing a standard - :class:`TorchForecastingModel` interface. + This is a Darts wrapper around The Forecasting Company's open-weights T0 model [1]_, [2]_. The implementation + delegates all forecasting logic to the optional `tfc-t0 `_ package while + exposing a standard :class:`TorchForecastingModel` interface. + + T0 is a pre-trained patch-transformer foundation model designed for zero-shot forecasting across both short and + long horizons. Multiple checkpoints are available via ``hub_model_name``: - T0 is a ~100M-parameter pre-trained patch-transformer foundation model designed for zero-shot forecasting - across both short and long horizons. + - `theforecastingcompany/t0-alpha `_: The original + ~100M-parameter T0 model. + - `theforecastingcompany/t0-beta `_ (default): The second + 256M-parameter model iteration This model supports univariate and multivariate time series, as well as past and future covariates. Because T0 is variate-agnostic, past covariates are forecast jointly with the target series (and dropped from the output); future covariates are conditioned on but not forecast. + Using this model will automatically download and cache the pre-trained ``t0-beta`` model from HuggingFace Hub. + Alternatively, you can specify a local directory containing the model config and weights using the ``local_dir`` + parameter. + By default, the model is deterministic (median forecast only). To enable probabilistic forecasts, pass a :class:`~darts.utils.likelihood_models.torch.QuantileRegression` instance to the ``likelihood`` parameter. It is recommended to call :func:`predict()` with ``predict_likelihood_parameters=True`` or ``num_samples >> 1`` to get meaningful results. T0 was trained on quantile levels [0.1, 0.25, 0.5, 0.75, 0.9]; other levels are interpolated, so any quantiles in the open interval (0, 1) may be requested. - For more details on the T0 model, see the `model card `_ - and the `tfc-t0 repository `_. The weights are open but the - Hub repository is gated, so request access on the model card and authenticate before first use. - + .. tip:: + You can perform full or partial fine-tuning of the model by setting the ``enable_finetuning`` parameter. + Read more in the parameter description below and in the `Fine-Tuning Examples + `__. + Fine-tuning runs a single forward pass, so ``output_chunk_length`` plus ``output_chunk_shift`` must not + exceed T0's ``max_horizon`` (1024). Dropout is disabled so a frozen backbone stays deterministic. .. note:: - The model can be fine-tuned (full or partial) via ``enable_finetuning``. Fine-tuning runs a - single forward pass, so ``output_chunk_length`` plus ``output_chunk_shift`` must not exceed - T0's ``max_horizon``. The training loss is computed on all pre-trained quantiles to preserve - the pre-trained distribution, and dropout is disabled so a frozen backbone stays - deterministic. + The T0 weights are open but the Hub repository is gated, so request access on the model card and + authenticate before first use. Parameters ---------- input_chunk_length Number of time steps in the past to take as a model input (per chunk). Applies to the target series, and past and/or future covariates (if the model supports it). + Can be either an ``int`` for a fixed input window, or a ``(min_length, max_length)`` tuple to enable + variable-length inputs for inference and fine-tuning. output_chunk_length Number of time steps predicted at once (per chunk) by the internal model. Also, the number of future values from future covariates to use as a model input (if the model supports future covariates). It is not the same @@ -272,16 +280,20 @@ def __init__( input chunk end). This will create a gap between the input and output. If the model supports `future_covariates`, the future values are extracted from the shifted output chunk. Predictions will start `output_chunk_shift` steps after the end of the target `series`. If `output_chunk_shift` is set, the model - cannot generate autoregressive predictions (`n > output_chunk_length`). + cannot generate autoregressive predictions (`n > output_chunk_length`). The covariate values inside the + gap are unknown to the model and are masked out. likelihood The likelihood model to be used for probabilistic forecasts. Must be ``None`` or an instance of - :class:`~darts.utils.likelihood_models.torch.QuantileRegression`. Any quantiles in the open interval - (0, 1) are supported (T0 interpolates levels it was not trained on). Default: ``None``, which will make - the model deterministic (median quantile only). + :class:`~darts.utils.likelihood_models.torch.QuantileRegression`. For zero-shot predicitons, any quantiles + in the open interval (0, 1) are supported (T0 interpolates levels it was not trained on). + Default: ``None``, which will make the model deterministic (median quantile only). + When fine-tuning is enabled, the training loss is always computed on all pre-trained quantiles to + preserve the full distribution, regardless of the ``likelihood`` setting. The ``likelihood`` parameter + only affects prediction output. hub_model_name - The model ID on HuggingFace Hub. Default: ``"theforecastingcompany/t0-alpha"``. + The model ID on HuggingFace Hub. Default: ``"theforecastingcompany/t0-beta"``. hub_model_revision - The model version to use. This can be a branch name, tag name, or commit hash. Default: ``None``, which + The model version to use. This can be a branch name, tag name, or commit hash. Default is ``None``, which will use the default branch from ``hub_model_name``. local_dir Optional local directory to load the pre-downloaded model. If specified and the directory is empty, the @@ -292,11 +304,27 @@ def __init__( Optional arguments to initialize the pytorch_lightning.Module, pytorch_lightning.Trainer, and Darts' :class:`TorchForecastingModel`. + loss_fn + PyTorch loss function used for fine-tuning a deterministic model. Ignored for probabilistic models when + ``likelihood`` is specified. Default: ``nn.MSELoss()``. torch_metrics A torch metric or a ``MetricCollection`` used for evaluation. A full list of available metrics can be found at https://torchmetrics.readthedocs.io/en/latest/. Default: ``None``. + optimizer_cls + The PyTorch optimizer class to be used. Default: ``torch.optim.Adam``. + optimizer_kwargs + Optionally, some keyword arguments for the PyTorch optimizer (e.g., ``{'lr': 1e-3}`` + for specifying a learning rate). Otherwise, the default values of the selected ``optimizer_cls`` + will be used. Default: ``None``. + lr_scheduler_cls + Optionally, the PyTorch learning rate scheduler class to be used. Specifying ``None`` corresponds + to using a constant learning rate. Default: ``None``. + lr_scheduler_kwargs + Optionally, some keyword arguments for the PyTorch learning rate scheduler. Default: ``None``. batch_size - Number of time series (input and output sequences) used in each prediction pass. Default: ``32``. + Number of time series (input and output sequences) used in each training pass. Default: ``32``. + n_epochs + Number of epochs over which to train the model. Default: ``100``. model_name Name of the model. Used for creating checkpoints and saving tensorboard data. If not specified, defaults to the following string ``"YYYY-mm-dd_HH_MM_SS_torch_model_run_PID"``, where the initial part @@ -307,8 +335,11 @@ def __init__( Path of the working directory, where to save checkpoints and Tensorboard summaries. Default: current working directory. log_tensorboard - If set, use Tensorboard to log the different parameters. The logs will be located in: + If set, use Tensorboard to log the different parameters. The logs will be located in ``"{work_dir}/darts_logs/{model_name}/logs/"``. Default: ``False``. + nr_epochs_val_period + Number of epochs to wait before evaluating the validation loss (if a validation + ``TimeSeries`` is passed to the :func:`fit()` method). Default: ``1``. force_reset If set to ``True``, any previously-existing model with the same name will be reset (all checkpoints will be discarded). Default: ``False``. @@ -343,7 +374,7 @@ def encode_year(idx): } .. random_state - Controls the randomness of reproducible forecasting. + Controls the randomness of the weights initialization and reproducible forecasting. pl_trainer_kwargs By default :class:`TorchForecastingModel` creates a PyTorch Lightning Trainer with several useful presets that performs the training, validation and prediction processes. These presets include automatic @@ -362,14 +393,53 @@ def encode_year(idx): For more info, see here: https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html#trainer-flags , and - https://pytorch-lightning.readthedocs.io/en/stable/accelerators/gpu_basic.html#train-on-multiple-gpus + https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html#accelerators-gpu-basic + + With parameter ``"callbacks"`` you can add custom or PyTorch-Lightning built-in callbacks to Darts' + :class:`TorchForecastingModel`. Below is an example for adding EarlyStopping to the training process. + The model will stop training early if the validation loss `val_loss` does not improve beyond + specifications. + + .. highlight:: python + .. code-block:: python + + from pytorch_lightning.callbacks.early_stopping import EarlyStopping + + # stop training when validation loss does not decrease more than 0.05 (`min_delta`) over + # a period of 5 epochs (`patience`) + my_stopper = EarlyStopping( + monitor="val_loss", + patience=5, + min_delta=0.05, + mode='min', + ) + + pl_trainer_kwargs={"callbacks": [my_stopper]} + .. + + Note that you can also use a custom PyTorch Lightning Trainer for training and prediction with optional + parameter ``trainer`` in :func:`fit()` and :func:`predict()`. show_warnings whether to show warnings raised from PyTorch Lightning. Useful to detect potential issues of your forecasting use case. Default: ``False``. + enable_finetuning + Enables model fine-tuning. Only effective if not ``None``. + If a bool, specifies whether to perform full fine-tuning / training (all parameters are updated) or keep + all parameters frozen. If a dict, specifies which parameters to fine-tune. Must only contain one key-value + record. Can be used to: + + - Unfreeze specific parameters, while keeping everything else frozen: + ``{"unfreeze": ["param.name.patterns.*"]}`` + - Freeze specific parameters, while keeping everything else unfrozen: + ``{"freeze": ["param.name.patterns.*"]}`` + + Default: ``None``. References ---------- - .. [1] The Forecasting Company, "T0", https://huggingface.co/theforecastingcompany/t0-alpha. + .. [1] The Forecasting Company. "T0". https://huggingface.co/theforecastingcompany/t0-alpha. + .. [2] The Forecasting Company. "t0: A Time-Series Foundation Model for Forecasting with Context". + https://www.theforecastingcompany.com/papers/files/t0-technical-report.pdf. Examples -------- @@ -378,20 +448,36 @@ def encode_year(idx): >>> from darts.models import T0Model >>> from darts.datasets import AirPassengersDataset >>> series = AirPassengersDataset().load().astype("float32") - >>> model = T0Model(input_chunk_length=24, output_chunk_length=12) + >>> model = T0Model(input_chunk_length=12, output_chunk_length=6) >>> model.fit(series) - >>> pred = model.predict(n=12) + >>> pred = model.predict(n=6) + #Passengers + Month + 1961-01-01 449.567627 + 1961-02-01 452.335510 + 1961-03-01 452.222931 + 1961-04-01 455.815308 + 1961-05-01 456.107849 + 1961-06-01 454.145569 Probabilistic forecasting: >>> from darts.utils.likelihood_models import QuantileRegression >>> model = T0Model( - ... input_chunk_length=24, - ... output_chunk_length=12, - ... likelihood=QuantileRegression(quantiles=[0.1, 0.5, 0.9]), - ... ) + >>> input_chunk_length=12, + >>> output_chunk_length=6, + >>> likelihood=QuantileRegression(quantiles=[0.1, 0.5, 0.9]), + >>> ) >>> model.fit(series) - >>> pred = model.predict(n=12, predict_likelihood_parameters=True) + >>> pred = model.predict(n=6, predict_likelihood_parameters=True) + #Passengers_q0.100 #Passengers_q0.500 #Passengers_q0.900 + Month + 1961-01-01 374.172150 449.567627 544.139648 + 1961-02-01 342.856903 452.335510 600.989746 + 1961-03-01 324.758545 452.222931 636.924072 + 1961-04-01 307.714386 455.815308 658.392700 + 1961-05-01 295.556030 456.107849 690.342773 + 1961-06-01 286.393066 454.145569 702.022705 """ if likelihood is not None and not isinstance(likelihood, QuantileRegression): raise_log( @@ -403,25 +489,16 @@ def encode_year(idx): super().__init__(**kwargs) - self.hub_model_name = hub_model_name - self.hub_model_revision = hub_model_revision - self.local_dir = local_dir - - @property - def supports_past_covariates(self) -> bool: - return True - - @property - def supports_future_covariates(self) -> bool: - return True + self._hub_model_name = hub_model_name + self._hub_model_revision = hub_model_revision + self._local_dir = local_dir def _create_model(self, train_sample: TorchTrainingSample) -> PLForecastingModule: # enable_finetuning is injected into pl_module_params by the base class; # _T0Module accepts it as an explicit parameter and converts dict form to bool return _T0Module( - hub_model_name=self.hub_model_name, - hub_model_revision=self.hub_model_revision, - local_dir=self.local_dir, - all_quantiles=self._PRETRAINED_QUANTILES, + hub_model_name=self._hub_model_name, + hub_model_revision=self._hub_model_revision, + local_dir=self._local_dir, **(self.pl_module_params or {}), ) diff --git a/docs/source/index.rst b/docs/source/index.rst index 7f42ecb4ee..83818337e2 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -706,7 +706,7 @@ Our regression models are designed to predict continuous numerical values, makin - βœ… βœ… πŸ”΄ - βœ… βœ… - βœ… - - `T0 model card `_, `tfc-t0 GitHub `_ + - `T0 report `_, `tfc-t0 GitHub `_ * - **Ensemble Models** (`GlobalForecastingModel `_): Model support is dependent on ensembled forecasting models and the ensemble model itself - - diff --git a/examples/27-Torch-and-Foundation-Model-Fine-Tuning-examples.ipynb b/examples/27-Torch-and-Foundation-Model-Fine-Tuning-examples.ipynb index 9610a004bd..3ce2e6bf61 100644 --- a/examples/27-Torch-and-Foundation-Model-Fine-Tuning-examples.ipynb +++ b/examples/27-Torch-and-Foundation-Model-Fine-Tuning-examples.ipynb @@ -60,7 +60,7 @@ "from darts import set_option\n", "from darts.datasets import AirPassengersDataset, AusBeerDataset\n", "from darts.metrics import mae\n", - "from darts.models import Chronos2Model, T0Model, TiDEModel\n", + "from darts.models import Chronos2Model, TiDEModel\n", "from darts.utils.callbacks import TFMProgressBar\n", "\n", "warnings.filterwarnings(\"ignore\")\n", @@ -516,111 +516,6 @@ ");" ] }, - { - "cell_type": "markdown", - "id": "b1f4a0c7", - "metadata": {}, - "source": [ - "#### Another foundation model: T0\n", - "\n", - "The same `enable_finetuning` API works for every foundation model. Let's repeat the exercise with [T0](https://huggingface.co/theforecastingcompany/t0-alpha), The Forecasting Company's open-weights ~100M-parameter patch transformer, on the same beer production data.\n", - "\n", - "
\n", - " Getting T0: the weights are open and published on the Hugging Face Hub, but the repository is gated. Open the model card while signed in and accept the access conditions, then authenticate your environment with huggingface-cli login or by setting HF_TOKEN. Signing in on the website alone does not authenticate your Python environment. T0 also needs its optional package: pip install tfc-t0.\n", - "
\n", - "\n", - "
\n", - "\n", - "T0: Fine-tunable layers\n", - "\n", - "T0 is a patch transformer: a patch encoder, 24 transformer blocks, and a decoder that maps each position to the next patch's quantiles. The output-side decoder is the cheapest useful surface, and is the direct analogue of Chronos-2's `output_patch_embedding`:\n", - "\n", - "| Layer group | Trainable params | Notes |\n", - "|------------------------------------------|------------------|------------------------------------------|\n", - "| `*decoder*` | ~0.43M | Output head, a good default |\n", - "| `*decoder*` + last transformer block(s) | ~21M | More capacity, slower, easier to overfit |\n", - "| Full model (`enable_finetuning=True`) | ~101.6M | Most flexible, most expensive |\n", - "\n", - "Three T0-specific notes:\n", - "\n", - "1. **Patterns must match the full parameter name.** T0's weights live under the `t0.` prefix inside the Darts module, so use `\"*decoder*\"` rather than `\"decoder*\"`. If a pattern matches nothing, *every* parameter stays frozen and training fails with an autograd error. Check with `[n for n, p in model.model.named_parameters() if p.requires_grad]`.\n", - "2. **Fine-tuning runs a single forward pass**, so `output_chunk_length` plus `output_chunk_shift` must not exceed T0's `max_horizon`. Longer horizons still work for zero-shot `predict()`, which continues auto-regressively.\n", - "3. **T0 reads its context in patches of 32 steps.** With `input_chunk_length=24` below it sees less than one full patch, which is well outside the regime it was trained for. Expect a weak zero-shot baseline here and a correspondingly large gain from fine-tuning. A longer input window is the first thing to try on real data.\n", - "\n", - "Dropout is disabled automatically during fine-tuning, so a frozen backbone stays deterministic.\n", - "\n", - "
" - ] - }, - { - "cell_type": "code", - "id": "c2e5b1d8", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "t0_params = dict(\n", - " input_chunk_length=input_chunk_length,\n", - " output_chunk_length=output_chunk_length,\n", - " random_state=42,\n", - ")\n", - "model = T0Model(**t0_params)\n", - "\n", - "# as before, `fit()` only loads the model without fine-tuning\n", - "model.fit(series=train_beer)\n", - "\n", - "# predict\n", - "pred_beer = model.predict(n=output_chunk_length, series=val_beer)\n", - "\n", - "# plot\n", - "series_beer[-3 * output_chunk_length :].plot(label=\"Ground truth\")\n", - "pred_beer.plot(\n", - " label=\"Forecast\",\n", - " title=f\"Pre-trained T0; MAE {mae(series_beer, pred_beer):.2f}\",\n", - ");" - ] - }, - { - "cell_type": "code", - "id": "d3a6c2e9", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# fine-tune only the output decoder (\"*decoder*\" matches the `t0.decoder.*` parameters)\n", - "model_finetune = T0Model(\n", - " enable_finetuning={\"unfreeze\": [\"*decoder*\"]},\n", - " save_checkpoints=True,\n", - " model_name=\"t0\",\n", - " force_reset=True,\n", - " optimizer_kwargs={\"lr\": 1e-3},\n", - " pl_trainer_kwargs=dict(\n", - " check_val_every_n_epoch=5,\n", - " gradient_clip_val=1,\n", - " callbacks=[TFMProgressBar(enable_train_bar_only=True)],\n", - " ),\n", - " **t0_params,\n", - ")\n", - "\n", - "# fine-tune for 15 epochs\n", - "model_finetune.fit(\n", - " series=train_beer,\n", - " val_series=val_beer,\n", - " load_best=True,\n", - " epochs=15,\n", - ")\n", - "\n", - "# predict\n", - "pred_beer = model_finetune.predict(n=output_chunk_length, series=val_beer)\n", - "\n", - "# plot\n", - "series_beer[-3 * output_chunk_length :].plot(label=\"Ground truth\")\n", - "pred_beer.plot(\n", - " label=\"Forecast\",\n", - " title=f\"Fine-tuned T0; MAE {mae(series_beer, pred_beer):.2f}\",\n", - ");" - ] - }, { "cell_type": "markdown", "id": "996456e0",