From ffa6f14af8618ca9f74710d17dd96aff88f4ed27 Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Wed, 12 Aug 2026 16:15:47 +0200 Subject: [PATCH 01/18] Add tslearn.foundation module for re-using pre-trained time series models Introduces ZeroShotForecaster, LinearProbeForecaster, LinearProbeClassifier and the underlying TimeSeriesFoundationEmbedder, letting pre-trained models (e.g. Chronos-2) be used behind the usual tslearn estimator API, either zero-shot or with a linear head fitted on frozen representations. --- CHANGELOG.md | 7 + .../plot_foundation_linear_probe.py | 203 +++++ .../plot_foundation_forecasting.py | 238 ++++++ docs/gen_modules/tslearn.foundation.rst | 8 + docs/reference.rst | 1 + docs/requirements_rtd.txt | 119 ++- pyproject.toml | 7 +- tests/test_estimators.py | 8 + tests/test_foundation.py | 808 ++++++++++++++++++ tslearn/foundation/__init__.py | 37 + tslearn/foundation/_classification.py | 304 +++++++ tslearn/foundation/_embedding.py | 608 +++++++++++++ tslearn/foundation/_forecasting.py | 806 +++++++++++++++++ 13 files changed, 3149 insertions(+), 5 deletions(-) create mode 100644 docs/examples/classification/plot_foundation_linear_probe.py create mode 100644 docs/examples/forecasting/plot_foundation_forecasting.py create mode 100644 docs/gen_modules/tslearn.foundation.rst create mode 100644 tests/test_foundation.py create mode 100644 tslearn/foundation/__init__.py create mode 100644 tslearn/foundation/_classification.py create mode 100644 tslearn/foundation/_embedding.py create mode 100644 tslearn/foundation/_forecasting.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3736cca7..ab7a20c9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,13 @@ Changelogs for this project are recorded in this file since v0.2.0. ### Added +* New `tslearn.foundation` module to re-use pre-trained time series models, such as + the ones published on the Hugging Face Hub, behind the usual tslearn API. It provides + `ZeroShotForecaster`, `LinearProbeForecaster`, `LinearProbeClassifier` and the + underlying `TimeSeriesFoundationEmbedder` feature extractor, which allows one to choose + the probed layer, the tokens taken into account and the pooling applied to their + representations. With `pooling=None`, the latter acts as a time series to time series + transform, returning one representation per token. Requires PyTorch. * soft-dtw related tools now support Itakura and Sakoe-Chiba global constraints. ([#189](https://github.com/tslearn-team/tslearn/issues/189)) * Multithreading support added to `cdist_soft_dtw` and `softdtw_barycenter`. ([#310](https://github.com/tslearn-team/tslearn/issues/310)) diff --git a/docs/examples/classification/plot_foundation_linear_probe.py b/docs/examples/classification/plot_foundation_linear_probe.py new file mode 100644 index 00000000..0ac5c290 --- /dev/null +++ b/docs/examples/classification/plot_foundation_linear_probe.py @@ -0,0 +1,203 @@ +""" +Linear probing a pre-trained model for classification +===================================================== + +Time series foundation models are pre-trained for forecasting, yet the +representations they build along the way carry a lot of information about the +shape of a series, which makes them useful for other tasks as well. Linear +probing [1]_ is the standard way to measure that: the pre-trained model is kept +frozen and used as a feature extractor, and a plain linear classifier is fitted +on top of its representations. Because the head is linear and the backbone is +never updated, the accuracy reached tells us how linearly separable the classes +already are in the representation space. + +This example applies :class:`~tslearn.foundation.LinearProbeClassifier` to a +UCR dataset, using Chronos-2 [2]_ as the frozen backbone, and compares it to +tslearn's classical baselines. Running it requires the ``chronos-forecasting`` +package:: + + pip install "chronos-forecasting>=2.0" + +References +---------- +.. [1] G. Alain and Y. Bengio. Understanding intermediate layers using linear + classifier probes. ICLR Workshop, 2017. +.. [2] A. F. Ansari, O. Shchur, J. Küken, et al. Chronos-2: From Univariate to + Universal Forecasting. arXiv:2510.15821, 2025. +""" + +############################################################################## +# Data +# ---- +# +# We use the ``Trace`` dataset from the UCR archive, which gathers four classes +# of transient signals recorded in a nuclear power plant, with only 100 training +# series in total. + +import numpy as np + +from tslearn.datasets import CachedDatasets +from tslearn.preprocessing import TimeSeriesScalerMeanVariance + +X_train, y_train, X_test, y_test = CachedDatasets().load_dataset("Trace") + +scaler = TimeSeriesScalerMeanVariance() +X_train = scaler.fit_transform(X_train) +X_test = scaler.transform(X_test) + +print(f"{X_train.shape=}, {X_test.shape=}, {len(np.unique(y_train))} classes") + +############################################################################## +# Probing the pre-trained model +# ----------------------------- +# +# The classifier only needs the pre-trained model and, since Chronos-2 returns +# forecasts rather than hidden states, an explicit layer to read +# representations from. A forward hook is placed on the selected block, so no +# modification of the model is needed. +# +# Fitting is fast: each series goes through the model exactly once, and the +# only thing actually trained is a logistic regression over a few hundred +# features. + +from chronos import Chronos2Pipeline + +from tslearn.foundation import LinearProbeClassifier + +pipeline = Chronos2Pipeline.from_pretrained("autogluon/chronos-2-small") + +clf = LinearProbeClassifier( + pipeline.model, + layer=-2, + pooling="mean", + # Chronos-2 appends a register token and a forecast token after its context + # tokens; neither represents the series, so they are left out of the average + tokens=(0, -2), +) +clf.fit(X_train, y_train) + +print(f"Embedding size: {clf.embedder_.embedding_size_}") +print(f"Test accuracy: {clf.score(X_test, y_test):.3f}") + +############################################################################## +# Which layer, which pooling? +# --------------------------- +# +# The representations of a pre-trained model change a lot from one layer to the +# next, and the last ones are usually the most specialized towards the +# pre-training objective, here forecasting. Sweeping over layers is therefore +# worthwhile, and cheap: nothing is being trained beyond the linear head. +# +# ``pooling`` matters just as much. Chronos-2 emits one token per patch of the +# context plus a register token, so we compare averaging those tokens against +# picking the register token alone, the latter playing the role a ``[CLS]`` +# token plays in a text encoder. + +n_layers = len(pipeline.model.encoder.block) +poolings = ["mean", "max", "cls"] +# Every other layer is enough to see the trend, and halves the build time +layers = sorted({*range(0, n_layers, 2), n_layers - 1}) + +accuracies = {} +for layer in layers: + for pooling in poolings: + model = LinearProbeClassifier( + pipeline.model, + layer=layer, + pooling=pooling, + # Chronos-2 places its register token after the context tokens + cls_index=-2 if pooling == "cls" else 0, + tokens=(0, -2), + ).fit(X_train, y_train) + accuracies[layer, pooling] = model.score(X_test, y_test) + +header = "layer " + " ".join(f"{pooling:>6}" for pooling in poolings) +print(header) +for layer in layers: + row = " ".join(f"{accuracies[layer, pooling]:>6.3f}" for pooling in poolings) + print(f"{layer:>5} {row}") + +############################################################################## + +import matplotlib.pyplot as plt + +fig, ax = plt.subplots(figsize=(8, 4), layout="constrained") +for pooling in poolings: + ax.plot( + layers, + [accuracies[layer, pooling] for layer in layers], + marker="o", + label=f"pooling={pooling}", + ) +ax.set_xlabel("probed layer") +ax.set_ylabel("test accuracy") +ax.set_title("Linear probe accuracy across the layers of Chronos-2") +ax.legend() +plt.show() + +############################################################################## +# Using the representations elsewhere +# ----------------------------------- +# +# The feature extractor can also be used on its own, through +# :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. Being a regular +# scikit-learn transformer, it composes with the rest of the ecosystem: here we +# project the frozen representations of the test set onto two dimensions to see +# whether the classes separate. + +from sklearn.decomposition import PCA + +from tslearn.foundation import TimeSeriesFoundationEmbedder + +best_layer, best_pooling = max(accuracies, key=accuracies.get) +embedder = TimeSeriesFoundationEmbedder( + pipeline.model, layer=best_layer, tokens=(0, -2) +) +embeddings = embedder.fit_transform(X_test) +projected = PCA(n_components=2).fit_transform(embeddings) + +fig, ax = plt.subplots(figsize=(6, 5), layout="constrained") +for label in np.unique(y_test): + mask = y_test == label + ax.scatter(projected[mask, 0], projected[mask, 1], label=f"class {label}", s=20) +ax.set_xlabel("first principal component") +ax.set_ylabel("second principal component") +ax.set_title("Frozen representations of the Trace test set") +ax.legend() +plt.show() + +############################################################################## +# Representations as time series +# ------------------------------ +# +# Setting ``pooling=None`` skips the aggregation step altogether. Instead of +# one vector per series, the transformer then returns one vector per token, +# that is, a time series dataset of shape ``(n_ts, n_tokens, dim)``. Combined +# with ``tokens``, which drops the tokens that do not represent the input, this +# turns the pre-trained model into a time series to time series transform whose +# output can be fed to any other tslearn estimator. +# +# Chronos-2 emits one token per patch of 16 timesteps, so the representations +# below are a 16-times subsampled, high-dimensional view of the input series. +# Clustering those sequences under DTW recovers the classes reasonably well, +# without the labels ever being used. + +from tslearn.clustering import TimeSeriesKMeans + +from sklearn.pipeline import Pipeline +from sklearn.metrics import adjusted_rand_score + +embedder = TimeSeriesFoundationEmbedder( + pipeline.model, layer=best_layer, pooling=None, tokens=(0, -2) +) +clusterer = TimeSeriesKMeans( + n_clusters=len(np.unique(y_test)), metric="dtw", max_iter=5, random_state=0 +) +pipeline = Pipeline(steps=[ + ("embed", embedder), + ("cluster", clusterer) +]) +labels = pipeline.fit_predict(X_test) + +print(f"Adjusted Rand index against the true classes: " + f"{adjusted_rand_score(y_test, labels):.3f}") diff --git a/docs/examples/forecasting/plot_foundation_forecasting.py b/docs/examples/forecasting/plot_foundation_forecasting.py new file mode 100644 index 00000000..a9c6e887 --- /dev/null +++ b/docs/examples/forecasting/plot_foundation_forecasting.py @@ -0,0 +1,238 @@ +""" +Forecasting with a pre-trained model +==================================== + +This example showcases the :mod:`tslearn.foundation` module, which allows one +to re-use a pre-trained time series model behind the usual tslearn API. Two +adaptation strategies are compared on the same data: + +* zero-shot forecasting, with :class:`~tslearn.foundation.ZeroShotForecaster`, + where the pre-trained model is used exactly as it is; +* linear probing, with :class:`~tslearn.foundation.LinearProbeForecaster`, + where the pre-trained model is kept frozen and only a linear head mapping its + representations to future values is fitted. + +Both are compared to the classical :class:`~tslearn.forecasting.AutoVARIMA` +baseline. + +The pre-trained model used here is Chronos-2 [1]_, an encoder-only model +pre-trained on a large corpus of real and synthetic series. Running this +example therefore requires the ``chronos-forecasting`` package:: + + pip install "chronos-forecasting>=2.0" + +Note that nothing in :mod:`tslearn.foundation` is specific to Chronos: any +PyTorch model exposing similar conventions can be plugged in the same way. + +References +---------- +.. [1] A. F. Ansari, O. Shchur, J. Küken, et al. Chronos-2: From Univariate to + Universal Forecasting. arXiv:2510.15821, 2025. +""" + +############################################################################## +# Data +# ---- +# +# We use a set of sine waves with varying frequencies, phases and noise levels. +# Using synthetic data keeps the example self-contained, and lets us check that +# the models pick up a periodic structure that a short-context model would miss. + +import numpy as np + +from tslearn.preprocessing import TimeSeriesScalerMeanVariance + +rng = np.random.RandomState(0) + +n_ts, sz, horizon, context_length = 32, 512, 48, 128 +t = np.arange(sz + horizon) + +periods = rng.uniform(30, 50, size=n_ts) +phases = rng.uniform(0, 2 * np.pi, size=n_ts) +trends = rng.uniform(-2e-3, 2e-3, size=n_ts) + +full_series = ( + np.sin(2 * np.pi * t[None, :] / periods[:, None] + phases[:, None]) + + trends[:, None] * t[None, :] + + 0.1 * rng.randn(n_ts, sz + horizon) +) +full_series = TimeSeriesScalerMeanVariance().fit_transform(full_series) + +X_train, X_test = full_series[:, :sz], full_series[:, sz:] +print(f"{X_train.shape=}, {X_test.shape=}") + +############################################################################## +# Zero-shot forecasting +# --------------------- +# +# A time series foundation model is pre-trained to forecast unseen series out +# of the box, so no training is needed at all. Wrapping the inference pipeline +# in a :class:`~tslearn.foundation.ZeroShotForecaster` is enough to obtain a +# tslearn estimator, whose ``predict`` method returns an array of shape +# ``(n_ts, n, d)`` like every other forecaster of the library. +# +# Calling ``fit`` is not needed here. Doing so is allowed, so that the +# estimator can be dropped into scikit-learn tooling that expects it, but it +# does not learn anything and emits a warning to that effect. + +from chronos import Chronos2Pipeline + +from tslearn.foundation import ZeroShotForecaster + +pipeline = Chronos2Pipeline.from_pretrained("autogluon/chronos-2-small") + +zero_shot = ZeroShotForecaster(pipeline, horizon_axis=-1) +y_zero_shot = zero_shot.predict(X_train, n=horizon) +print(f"{y_zero_shot.shape=}") + +############################################################################## +# ``horizon_axis`` says which axis of the model's output holds the forecast +# horizon. There is no shared convention across implementations, Chronos-2 +# alone returning ``(n_series, n_quantiles, horizon)`` from ``predict`` and +# ``(n_series, horizon, n_quantiles)`` from ``predict_quantiles``, so stating +# it is the reliable option. Left to its default of ``"auto"``, the estimator +# infers it from the returned shape and warns when the shape is ambiguous. +# +# The remaining axes are understood as holding quantile levels or sample paths +# and are reduced to a point forecast according to the ``quantile`` parameter, +# which defaults to the median. + +############################################################################## +# Linear probing +# -------------- +# +# Zero-shot use ignores the fact that we do have training data at hand. Linear +# probing exploits it, while leaving the pre-trained weights untouched: the +# model is only used to embed the context window, and a linear map from those +# representations to the next ``horizon`` values is fitted. +# +# Training pairs are cut out of the series with a sliding window. A larger +# ``stride`` yields fewer, less redundant windows and a faster fit, which +# matters because every window has to go through the pre-trained model once. +# +# The probing estimator needs the model itself rather than the inference +# pipeline, since it reads hidden states rather than forecasts. Two options +# drive which representations are read: +# +# * ``layer`` selects the block to probe: an integer places a forward hook on +# the corresponding block, so ``layer=-2`` probes the penultimate one. The +# last layers of a pre-trained model tend to specialize towards its +# pre-training objective, so probing slightly earlier is often beneficial. +# ``layer=None`` reads the model's own output instead, which only works for +# models returning their hidden states; Chronos-2 returns quantile forecasts +# only, so an explicit layer is required here. +# * ``pooling`` says how the token representations are aggregated into a single +# vector. Chronos-2 emits one token per patch of the context, plus a register +# token, so both averaging (``pooling="mean"``) and picking that single token +# (``pooling="cls"``) are sensible. +# * ``tokens`` restricts which tokens take part in that aggregation. Chronos-2 +# appends a register token and a forecast token after its context tokens, and +# neither represents the input series, so ``tokens=(0, -2)`` keeps the average +# clean. + +from tslearn.foundation import LinearProbeForecaster + +probe = LinearProbeForecaster( + pipeline.model, + context_length=context_length, + horizon=horizon, + stride=16, + layer=-2, + pooling="mean", + tokens=(0, -2), +) +probe.fit(X_train) +y_probe = probe.predict(X_train) + +print(f"{probe.n_windows_} training windows, " + f"{probe.embedder_.embedding_size_}-dimensional embeddings") + +############################################################################## +# The head fitted on top of the frozen representations defaults to a +# :class:`sklearn.linear_model.RidgeCV`, which picks its regularization +# strength by cross-validation. Any scikit-learn regressor can be passed +# instead through the ``probe`` parameter, for instance to control the +# regularization path explicitly. + +############################################################################## +# Comparison +# ---------- +# +# Both strategies are compared to AutoVARIMA on the mean absolute error over +# the held-out horizon. + +from tslearn.forecasting import AutoVARIMA +from tslearn.metrics.performance import mae + +varima = AutoVARIMA(seasonal_period=40).fit(X_train) +y_varima = varima.predict(n=horizon) + +scores = { + "AutoVARIMA (with seasonal differentiation)": mae(X_test, y_varima), + "Chronos-2 (zero-shot)": mae(X_test, y_zero_shot), + "Chronos-2 (linear probe)": mae(X_test, y_probe), +} +for name, score in scores.items(): + print(f"{name:>42}: MAE = {score:.4f}") + +############################################################################## +# Let us look at a few series to see where the differences come from. + +import matplotlib.pyplot as plt + +shown = [0, 1, 2] +fig, axes = plt.subplots( + len(shown), 1, sharex=True, figsize=(11, 8), layout="constrained" +) +context = 150 +for ax, i in zip(axes, shown): + ax.plot( + np.arange(sz - context, sz), + X_train[i, -context:, 0], + color="0.4", + label="context", + ) + ax.plot( + np.arange(sz, sz + horizon), + X_test[i, :, 0], + color="k", + label="ground truth", + ) + for values, label in [ + (y_varima, "AutoVARIMA"), + (y_zero_shot, "zero-shot"), + (y_probe, "linear probe"), + ]: + ax.plot(np.arange(sz, sz + horizon), values[i, :, 0], label=label, alpha=0.8) + ax.axvline(sz, color="0.8", linestyle="--") + ax.set_ylabel(f"series {i}") +axes[0].legend(loc="upper left", ncol=5, fontsize="small") +fig.suptitle("Forecasting a held-out horizon", fontsize=14) +plt.show() + +############################################################################## +# Choosing what to probe +# ---------------------- +# +# Which layer and which pooling work best is data-dependent, and cheap enough +# to explore: the frozen model is the expensive part, and it is never trained. +# Below we score a few configurations on the held-out horizon. + +results = {} +for layer in [-1, -2, -4]: + for pooling in ["mean", "cls"]: + model = LinearProbeForecaster( + pipeline.model, + context_length=context_length, + horizon=horizon, + stride=48, + layer=layer, + pooling=pooling, + # Chronos-2 appends its register token after the context tokens + cls_index=-2, + tokens=(0, -2), + ).fit(X_train) + results[(layer, pooling)] = mae(X_test, model.predict(X_train)) + +for (layer, pooling), score in sorted(results.items(), key=lambda kv: kv[1]): + print(f"layer={str(layer):>4}, pooling={pooling:>4}: MAE = {score:.4f}") diff --git a/docs/gen_modules/tslearn.foundation.rst b/docs/gen_modules/tslearn.foundation.rst new file mode 100644 index 00000000..f1e25911 --- /dev/null +++ b/docs/gen_modules/tslearn.foundation.rst @@ -0,0 +1,8 @@ +.. _mod-tslearn.foundation: + +tslearn.foundation +================== + +.. automodule:: tslearn.foundation + + \ No newline at end of file diff --git a/docs/reference.rst b/docs/reference.rst index 3dededac..09e8b81f 100644 --- a/docs/reference.rst +++ b/docs/reference.rst @@ -16,6 +16,7 @@ The complete ``tslearn`` project is automatically documented for every module. datasets early_classification forecasting + foundation generators matrix_profile metrics diff --git a/docs/requirements_rtd.txt b/docs/requirements_rtd.txt index e872cc43..55b73d6c 100644 --- a/docs/requirements_rtd.txt +++ b/docs/requirements_rtd.txt @@ -4,10 +4,16 @@ # # pip-compile --extra=docs --output-file=docs/requirements_rtd.txt # +accelerate==1.14.0 + # via chronos-forecasting accessible-pygments==0.0.5 # via pydata-sphinx-theme alabaster==1.0.0 # via sphinx +annotated-doc==0.0.5 + # via typer +anyio==4.14.2 + # via httpx babel==2.18.0 # via # pydata-sphinx-theme @@ -15,9 +21,16 @@ babel==2.18.0 beautifulsoup4==4.14.3 # via pydata-sphinx-theme certifi==2026.2.25 - # via requests + # via + # httpcore + # httpx + # requests charset-normalizer==3.4.7 # via requests +chronos-forecasting==2.3.1 + # via tslearn (pyproject.toml) +click==8.4.2 + # via huggingface-hub contourpy==1.3.3 # via matplotlib cycler==0.12.1 @@ -26,14 +39,42 @@ docutils==0.22.4 # via # pydata-sphinx-theme # sphinx +einops==0.8.2 + # via chronos-forecasting +filelock==3.32.2 + # via + # huggingface-hub + # torch fonttools==4.62.1 # via matplotlib +fsspec==2026.7.0 + # via + # huggingface-hub + # torch +h11==0.16.0 + # via httpcore +hf-xet==1.6.0 + # via huggingface-hub +httpcore==1.0.9 + # via httpx +httpx==0.28.1 + # via huggingface-hub +huggingface-hub==1.27.0 + # via + # accelerate + # tokenizers + # transformers idna==3.15 - # via requests + # via + # anyio + # httpx + # requests imagesize==2.0.0 # via sphinx jinja2==3.1.6 - # via sphinx + # via + # sphinx + # torch joblib==1.5.3 # via # scikit-learn @@ -42,53 +83,98 @@ kiwisolver==1.5.0 # via matplotlib llvmlite==0.47.0 # via numba +markdown-it-py==4.2.0 + # via rich markupsafe==3.0.3 # via jinja2 matplotlib==3.10.8 # via tslearn (pyproject.toml) +mdurl==0.1.2 + # via markdown-it-py +mpmath==1.3.0 + # via sympy +networkx==3.6.1 + # via torch numba==0.65.0 # via tslearn (pyproject.toml) numpy==2.4.4 # via + # accelerate + # chronos-forecasting # contourpy # matplotlib # numba + # pandas + # patsy # scikit-learn # scipy + # statsmodels + # transformers # tslearn (pyproject.toml) numpydoc==1.10.0 # via tslearn (pyproject.toml) packaging==26.1 # via + # accelerate + # huggingface-hub # matplotlib # sphinx + # statsmodels + # transformers +pandas==3.0.5 + # via + # chronos-forecasting + # statsmodels +patsy==1.0.2 + # via statsmodels pillow==12.3.0 # via # matplotlib # sphinx-gallery +psutil==7.2.2 + # via accelerate pydata-sphinx-theme==0.17.0 # via tslearn (pyproject.toml) pygments==2.20.0 # via # accessible-pygments # pydata-sphinx-theme + # rich # sphinx pypandoc==1.17 # via tslearn (pyproject.toml) pyparsing==3.3.2 # via matplotlib python-dateutil==2.9.0.post0 - # via matplotlib + # via + # matplotlib + # pandas +pyyaml==6.0.3 + # via + # accelerate + # huggingface-hub + # transformers +regex==2026.7.19 + # via transformers requests==2.33.1 # via sphinx +rich==15.0.0 + # via typer roman-numerals==4.1.0 # via sphinx +safetensors==0.8.0 + # via + # accelerate + # transformers scikit-learn==1.8.0 # via tslearn (pyproject.toml) scipy==1.17.1 # via # scikit-learn + # statsmodels # tslearn (pyproject.toml) +shellingham==1.5.4 + # via typer six==1.17.0 # via python-dateutil snowballstemmer==3.0.1 @@ -118,11 +204,36 @@ sphinxcontrib-qthelp==2.0.0 # via sphinx sphinxcontrib-serializinghtml==2.0.0 # via sphinx +statsmodels==0.14.6 + # via tslearn (pyproject.toml) +sympy==1.14.0 + # via torch threadpoolctl==3.6.0 # via scikit-learn +tokenizers==0.22.2 + # via transformers +torch==2.13.0 + # via + # accelerate + # chronos-forecasting + # tslearn (pyproject.toml) +tqdm==4.70.0 + # via + # huggingface-hub + # transformers +transformers==5.15.0 + # via chronos-forecasting +typer==0.27.1 + # via transformers typing-extensions==4.15.0 # via + # anyio # beautifulsoup4 + # huggingface-hub # pydata-sphinx-theme + # torch urllib3==2.7.0 # via requests + +# The following packages are considered to be unsafe in a requirements file: +# setuptools diff --git a/pyproject.toml b/pyproject.toml index c1fe364a..70642669 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,8 +42,10 @@ changelog = "https://github.com/tslearn-team/tslearn/CHANGELOG.md" [project.optional-dependencies] pytorch = ['torch'] +# Re-use of pre-trained models through the tslearn.foundation module +foundation = ['torch'] tests = [ - "pytest", + "pytest", "h5py" ] docs = [ @@ -54,6 +56,9 @@ docs = [ "numpydoc", "matplotlib", "pypandoc", + # Used by the tslearn.foundation examples of the gallery + "torch", + "chronos-forecasting>=2.0", ] all_features = [ "torch", diff --git a/tests/test_estimators.py b/tests/test_estimators.py index 5d78dc99..0cd0f285 100644 --- a/tests/test_estimators.py +++ b/tests/test_estimators.py @@ -91,6 +91,14 @@ def get_estimators(type_filter='all'): # only keep those that are from tslearn all_classes = filter(lambda c: not is_sklearn(c), all_classes) + # tslearn.foundation estimators wrap an externally provided pre-trained + # model, so they cannot be instantiated without arguments, which is what + # the common checks below rely on. They are covered by tests/test_foundation.py. + all_classes = filter( + lambda c: not inspect.getmodule(c).__name__.startswith('tslearn.foundation'), + all_classes + ) + # Now filter out the estimators that are not of the specified type filters = { 'all': [ClassifierMixin, RegressorMixin, diff --git a/tests/test_foundation.py b/tests/test_foundation.py new file mode 100644 index 00000000..5696e339 --- /dev/null +++ b/tests/test_foundation.py @@ -0,0 +1,808 @@ +"""Tests for the :mod:`tslearn.foundation` module. + +These tests run entirely offline, against small hand-written models that +reproduce the calling conventions of the pre-trained models published on the +Hugging Face Hub, so that no download is ever needed. +""" + +import numpy as np + +import pytest + +from sklearn.base import clone +from sklearn.linear_model import Ridge +from sklearn.pipeline import make_pipeline +from sklearn.svm import LinearSVC + +from tslearn.generators import random_walks + +torch = pytest.importorskip("torch") + +from tslearn.foundation import ( # noqa: E402 + LinearProbeClassifier, + LinearProbeForecaster, + TimeSeriesFoundationEmbedder, + ZeroShotForecaster, +) + + +PATCH_SIZE = 4 +D_MODEL = 8 + + +class _Block(torch.nn.Module): + """A single, deliberately simple, encoder block.""" + + def __init__(self, d_model=D_MODEL): + super().__init__() + self.linear = torch.nn.Linear(d_model, d_model) + + def forward(self, hidden_states): + return hidden_states + torch.tanh(self.linear(hidden_states)) + + +class _Encoder(torch.nn.Module): + def __init__(self, n_layers=3, d_model=D_MODEL): + super().__init__() + self.block = torch.nn.ModuleList([_Block(d_model) for _ in range(n_layers)]) + self.final_layer_norm = torch.nn.LayerNorm(d_model) + + def forward(self, hidden_states): + for block in self.block: + hidden_states = block(hidden_states) + return self.final_layer_norm(hidden_states) + + +class _Output: + """Stands in for a ``transformers`` ``ModelOutput``.""" + + def __init__(self, last_hidden_state): + self.last_hidden_state = last_hidden_state + + +class _DummyBackbone(torch.nn.Module): + """A patch-based encoder taking univariate series, like Chronos-2. + + A register token is appended after the context tokens, so that the + ``pooling="cls"`` code path can be exercised on a non-zero ``cls_index``. + """ + + def __init__(self, n_layers=3, d_model=D_MODEL, patch_size=PATCH_SIZE, seed=0): + super().__init__() + # Seeded so that results do not depend on which test ran before + torch.manual_seed(seed) + self.patch_size = patch_size + self.embedding = torch.nn.Linear(patch_size, d_model) + self.register_token = torch.nn.Parameter(torch.randn(1, 1, d_model)) + self.encoder = _Encoder(n_layers, d_model) + + def forward(self, context): + batch_size, sz = context.shape + n_patches = sz // self.patch_size + patches = context[:, : n_patches * self.patch_size] + patches = patches.reshape(batch_size, n_patches, self.patch_size) + hidden_states = self.embedding(patches) + register = self.register_token.expand(batch_size, -1, -1) + hidden_states = torch.cat([hidden_states, register], dim=1) + return _Output(self.encoder(hidden_states)) + + +class _MultivariateBackbone(torch.nn.Module): + """An encoder taking (batch, sz, d) arrays, like ``transformers`` models.""" + + def __init__(self, n_channels, d_model=D_MODEL, seed=0): + super().__init__() + torch.manual_seed(seed) + self.embedding = torch.nn.Linear(n_channels, d_model) + self.encoder = _Encoder(2, d_model) + + def forward(self, past_values): + return _Output(self.encoder(self.embedding(past_values))) + + +class _DummyPipeline: + """A zero-shot forecaster mimicking ``Chronos2Pipeline``. + + It returns one tensor per series, of shape + ``(n_variates, prediction_length, n_quantiles)``. + """ + + n_quantiles = 9 + + def predict(self, inputs, prediction_length=1, batch_size=256): + inputs = np.asarray(inputs, dtype=np.float64) + # Naive forecast: repeat the last observed value + last = inputs[:, -1] + quantile_offsets = np.linspace(-1.0, 1.0, self.n_quantiles) + forecast = ( + last[:, None, None] + + np.zeros((1, prediction_length, 1)) + + quantile_offsets[None, None, :] + ) + return [torch.as_tensor(row[None]) for row in forecast] + + +def _dataset(n_ts=6, sz=32, d=1, seed=0): + return random_walks(n_ts=n_ts, sz=sz, d=d, random_state=seed) + + +# --------------------------------------------------------------------------- +# TimeSeriesFoundationEmbedder +# --------------------------------------------------------------------------- + + +def test_embedder_shapes_and_layouts(): + X = _dataset(n_ts=5, sz=32, d=1) + embedder = TimeSeriesFoundationEmbedder(_DummyBackbone()) + embeddings = embedder.fit_transform(X) + assert embeddings.shape == (5, D_MODEL) + assert embedder.embedding_size_ == D_MODEL + + # Multivariate series are embedded channel per channel and concatenated + X_multi = _dataset(n_ts=5, sz=32, d=3) + embeddings = TimeSeriesFoundationEmbedder(_DummyBackbone()).fit_transform(X_multi) + assert embeddings.shape == (5, 3 * D_MODEL) + + # Natively multivariate models get the whole (n_ts, sz, d) array + embedder = TimeSeriesFoundationEmbedder( + _MultivariateBackbone(n_channels=3), input_layout="channels_last" + ) + assert embedder.fit_transform(X_multi).shape == (5, D_MODEL) + + +def test_embedder_channel_stacking_is_order_preserving(): + # The embedding of a multivariate series must be the concatenation of the + # embeddings of its channels, in channel order. + X = _dataset(n_ts=4, sz=32, d=2) + backbone = _DummyBackbone() + joint = TimeSeriesFoundationEmbedder(backbone).fit_transform(X) + per_channel = [ + TimeSeriesFoundationEmbedder(backbone).fit_transform(X[:, :, k : k + 1]) + for k in range(2) + ] + np.testing.assert_allclose( + joint, np.concatenate(per_channel, axis=1), rtol=1e-5, atol=1e-6 + ) + + +@pytest.mark.parametrize("pooling", ["mean", "max", "cls", "last", "flatten"]) +def test_embedder_poolings(pooling): + X = _dataset(n_ts=4, sz=32, d=1) + embedder = TimeSeriesFoundationEmbedder(_DummyBackbone(), pooling=pooling) + embeddings = embedder.fit_transform(X) + n_tokens = 32 // PATCH_SIZE + 1 # context patches + register token + expected = n_tokens * D_MODEL if pooling == "flatten" else D_MODEL + assert embeddings.shape == (4, expected) + assert embedder.n_tokens_ is None + + +@pytest.mark.parametrize("pooling", [None, "none"]) +def test_embedder_without_pooling_is_a_series_to_series_transform(pooling): + X = _dataset(n_ts=4, sz=32, d=1) + n_tokens = 32 // PATCH_SIZE + 1 # context patches + register token + + embedder = TimeSeriesFoundationEmbedder(_DummyBackbone(), pooling=pooling) + embeddings = embedder.fit_transform(X) + assert embeddings.shape == (4, n_tokens, D_MODEL) + assert embedder.n_tokens_ == n_tokens + assert embedder.embedding_size_ == D_MODEL + + # The output is a valid tslearn dataset, usable by other estimators + from tslearn.clustering import TimeSeriesKMeans + + labels = TimeSeriesKMeans(n_clusters=2, max_iter=2, random_state=0).fit_predict( + embeddings + ) + assert labels.shape == (4,) + + +def test_embedder_token_selection(): + X = _dataset(n_ts=4, sz=32, d=1) + n_context = 32 // PATCH_SIZE # tokens that represent the series + backbone = _DummyBackbone() + + # Dropping the trailing register token leaves the context tokens + embeddings = TimeSeriesFoundationEmbedder( + backbone, pooling=None, tokens=(0, -1) + ).fit_transform(X) + assert embeddings.shape == (4, n_context, D_MODEL) + + # Equivalent spellings + np.testing.assert_allclose( + embeddings, + TimeSeriesFoundationEmbedder( + backbone, pooling=None, tokens=slice(0, -1) + ).fit_transform(X), + ) + np.testing.assert_allclose( + embeddings, + TimeSeriesFoundationEmbedder( + backbone, pooling=None, tokens=(None, n_context) + ).fit_transform(X), + ) + + # Token selection also applies before pooling, and changes the result: + # averaging over the context tokens is not averaging over all of them + without_register = TimeSeriesFoundationEmbedder( + backbone, pooling="mean", tokens=(0, -1) + ).fit_transform(X) + over_everything = TimeSeriesFoundationEmbedder( + backbone, pooling="mean" + ).fit_transform(X) + np.testing.assert_allclose( + without_register, embeddings.mean(axis=1), rtol=1e-5, atol=1e-6 + ) + assert not np.allclose(without_register, over_everything) + + # ... but not to "cls", which is meant to reach an excluded token + np.testing.assert_allclose( + TimeSeriesFoundationEmbedder( + backbone, pooling="cls", cls_index=-1, tokens=(0, -1) + ).fit_transform(X), + TimeSeriesFoundationEmbedder( + backbone, pooling="cls", cls_index=-1 + ).fit_transform(X), + ) + + +def test_embedder_token_selection_multivariate(): + X = _dataset(n_ts=4, sz=32, d=3) + n_context = 32 // PATCH_SIZE + embeddings = TimeSeriesFoundationEmbedder( + _DummyBackbone(), pooling=None, tokens=(0, -1) + ).fit_transform(X) + # Channels are concatenated along the feature axis, one row per token + assert embeddings.shape == (4, n_context, 3 * D_MODEL) + + # Channel k of the output must be the univariate embedding of channel k + backbone = _DummyBackbone() + joint = TimeSeriesFoundationEmbedder( + backbone, pooling=None, tokens=(0, -1) + ).fit_transform(X) + for k in range(3): + alone = TimeSeriesFoundationEmbedder( + backbone, pooling=None, tokens=(0, -1) + ).fit_transform(X[:, :, k : k + 1]) + np.testing.assert_allclose( + joint[:, :, k * D_MODEL : (k + 1) * D_MODEL], + alone, + rtol=1e-5, + atol=1e-6, + ) + + +def test_embedder_token_selection_errors(): + X = _dataset(n_ts=4, sz=32, d=1) + with pytest.raises(ValueError, match="`tokens` must be"): + TimeSeriesFoundationEmbedder(_DummyBackbone(), tokens=3).fit(X) + with pytest.raises(ValueError, match="`tokens` must be"): + TimeSeriesFoundationEmbedder(_DummyBackbone(), tokens=(1, 2, 3)).fit(X) + with pytest.raises(ValueError, match="selects no token"): + TimeSeriesFoundationEmbedder(_DummyBackbone(), tokens=(0, 0)).fit(X) + + +def test_probes_reject_unpooled_representations(): + X, y = _classification_dataset() + for pooling in (None, "none"): + with pytest.raises(ValueError, match="pooling=None"): + LinearProbeClassifier(_DummyBackbone(), pooling=pooling).fit(X, y) + with pytest.raises(ValueError, match="pooling=None"): + LinearProbeForecaster( + _DummyBackbone(), pooling=pooling, context_length=16, horizon=2 + ).fit(X) + + +def test_probes_accept_token_selection(): + X, y = _classification_dataset() + model = LinearProbeClassifier(_DummyBackbone(), tokens=(0, -1)).fit(X, y) + assert model.predict(X).shape == (len(y),) + + X_fc = _dataset(n_ts=6, sz=48, d=1) + model = LinearProbeForecaster( + _DummyBackbone(), context_length=16, horizon=2, stride=8, tokens=(0, -1) + ).fit(X_fc) + assert model.predict(X_fc).shape == (6, 2, 1) + + +def test_embedder_cls_index_selects_the_register_token(): + X = _dataset(n_ts=3, sz=32, d=1) + # The register token is appended last and does not depend on the input, so + # selecting it must yield identical embeddings for all series. + embeddings = TimeSeriesFoundationEmbedder( + _DummyBackbone(), pooling="cls", cls_index=-1, layer=0 + ).fit_transform(X) + np.testing.assert_allclose( + embeddings, np.repeat(embeddings[:1], 3, axis=0), rtol=1e-5, atol=1e-6 + ) + + # Whereas a context token does depend on the input + embeddings = TimeSeriesFoundationEmbedder( + _DummyBackbone(), pooling="cls", cls_index=0, layer=0 + ).fit_transform(X) + assert not np.allclose(embeddings, embeddings[:1]) + + +def test_embedder_layer_selection(): + X = _dataset(n_ts=4, sz=32, d=1) + backbone = _DummyBackbone(n_layers=3) + embeddings = [ + TimeSeriesFoundationEmbedder(backbone, layer=layer).fit_transform(X) + for layer in (0, 1, 2, None) + ] + # Every layer yields a different representation of the same data + for i in range(len(embeddings)): + for j in range(i + 1, len(embeddings)): + assert not np.allclose(embeddings[i], embeddings[j]) + + # Negative indices address the stack from the end + np.testing.assert_allclose( + TimeSeriesFoundationEmbedder(backbone, layer=-1).fit_transform(X), + embeddings[2], + rtol=1e-5, + atol=1e-6, + ) + + # An explicit path to the layer stack gives the same result as autodetection + np.testing.assert_allclose( + TimeSeriesFoundationEmbedder( + backbone, layer=1, layers_path="encoder.block" + ).fit_transform(X), + embeddings[1], + rtol=1e-5, + atol=1e-6, + ) + + +def test_embedder_leaves_the_model_frozen(): + X = _dataset(n_ts=4, sz=32, d=1) + backbone = _DummyBackbone() + before = [p.detach().clone() for p in backbone.parameters()] + TimeSeriesFoundationEmbedder(backbone).fit_transform(X) + for parameter, reference in zip(backbone.parameters(), before): + assert parameter.grad is None + torch.testing.assert_close(parameter.detach(), reference) + + +def test_embedder_context_length_truncation(): + X = _dataset(n_ts=4, sz=64, d=1) + embedder = TimeSeriesFoundationEmbedder(_DummyBackbone(), context_length=32) + np.testing.assert_allclose( + embedder.fit_transform(X), + TimeSeriesFoundationEmbedder(embedder.model).fit_transform(X[:, -32:]), + rtol=1e-5, + atol=1e-6, + ) + + with pytest.raises(ValueError, match="context_length"): + TimeSeriesFoundationEmbedder( + _DummyBackbone(), context_length=128 + ).fit_transform(X) + + +def test_embedder_errors(): + X = _dataset(n_ts=4, sz=32, d=1) + with pytest.raises(ValueError, match="pooling"): + TimeSeriesFoundationEmbedder(_DummyBackbone(), pooling="median").fit(X) + with pytest.raises(ValueError, match="input_layout"): + TimeSeriesFoundationEmbedder(_DummyBackbone(), input_layout="nchw").fit(X) + with pytest.raises(TypeError, match="torch.nn.Module"): + TimeSeriesFoundationEmbedder(_DummyPipeline()).fit(X) + with pytest.raises(ValueError, match="out of range"): + TimeSeriesFoundationEmbedder(_DummyBackbone(n_layers=3), layer=7).fit(X) + with pytest.raises(ValueError, match="cls_index"): + TimeSeriesFoundationEmbedder( + _DummyBackbone(), pooling="cls", cls_index=999 + ).fit(X) + with pytest.raises(ValueError, match="features"): + TimeSeriesFoundationEmbedder(_DummyBackbone()).fit(X).transform( + _dataset(n_ts=4, sz=32, d=2) + ) + + +def test_embedder_in_a_sklearn_pipeline(): + X = _dataset(n_ts=20, sz=32, d=1) + y = np.arange(20) % 2 + pipeline = make_pipeline( + TimeSeriesFoundationEmbedder(_DummyBackbone()), LinearSVC() + ) + assert pipeline.fit(X, y).predict(X).shape == (20,) + assert clone(pipeline) is not pipeline + + +# --------------------------------------------------------------------------- +# ZeroShotForecaster +# --------------------------------------------------------------------------- + + +def test_zero_shot_forecaster(): + X = _dataset(n_ts=5, sz=32, d=1) + model = ZeroShotForecaster(_DummyPipeline()) + + with pytest.warns(UserWarning, match="does not train anything"): + model.fit(X) + + predicted = model.predict(X, n=7) + assert predicted.shape == (5, 7, 1) + # The dummy pipeline repeats the last value, and the median quantile of its + # symmetric output is that value exactly + np.testing.assert_allclose( + predicted, np.repeat(X[:, -1:], 7, axis=1), rtol=1e-5, atol=1e-6 + ) + + +def test_zero_shot_forecaster_no_warning_when_silenced(): + X = _dataset(n_ts=3, sz=32, d=1) + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("error") + ZeroShotForecaster(_DummyPipeline(), warn_on_fit=False).fit(X) + + +def test_zero_shot_forecaster_multivariate(): + X = _dataset(n_ts=5, sz=32, d=3) + predicted = ZeroShotForecaster(_DummyPipeline()).predict(X, n=4) + assert predicted.shape == (5, 4, 3) + # Channels must not get mixed up while being folded in and out of the batch + np.testing.assert_allclose( + predicted, np.repeat(X[:, -1:], 4, axis=1), rtol=1e-5, atol=1e-6 + ) + + +def test_zero_shot_forecaster_quantile_selection(): + X = _dataset(n_ts=3, sz=32, d=1) + low = ZeroShotForecaster(_DummyPipeline(), quantile=0.1).predict(X, n=3) + high = ZeroShotForecaster(_DummyPipeline(), quantile=0.9).predict(X, n=3) + assert np.all(low < high) + + +def test_zero_shot_forecaster_custom_predict_fn(): + X = _dataset(n_ts=4, sz=32, d=1) + + def predict_fn(model, context, horizon): + assert context.shape == (4, 32) + return np.zeros((context.shape[0], horizon)) + + predicted = ZeroShotForecaster(object(), predict_fn=predict_fn).predict(X, n=5) + np.testing.assert_array_equal(predicted, np.zeros((4, 5, 1))) + + +def test_zero_shot_forecaster_context_length(): + X = _dataset(n_ts=4, sz=64, d=1) + seen = {} + + def predict_fn(model, context, horizon): + seen["shape"] = context.shape + return np.zeros((context.shape[0], horizon)) + + ZeroShotForecaster( + object(), predict_fn=predict_fn, context_length=16 + ).predict(X, n=2) + assert seen["shape"] == (4, 16) + + +def test_zero_shot_forecaster_negotiates_the_context_format(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _NeedsVariateAxis: + """Rejects 2d contexts, like ``Chronos2Pipeline`` does.""" + + def predict(self, inputs, prediction_length=1): + inputs = np.asarray(inputs) + if inputs.ndim != 3: + raise ValueError( + "Expected 3-d tensor with shape " + "(n_series, n_variates, history_length)." + ) + return np.zeros((inputs.shape[0], prediction_length)) + + model = ZeroShotForecaster(_NeedsVariateAxis()) + assert model.predict(X, n=3).shape == (4, 3, 1) + assert model.context_format_ == "3d" + # The negotiated format is reused rather than re-discovered + assert model.predict(X, n=3).shape == (4, 3, 1) + + class _NeedsList: + def predict(self, inputs, prediction_length=1): + if not isinstance(inputs, list): + raise TypeError("A list of series is expected.") + return np.zeros((len(inputs), prediction_length)) + + model = ZeroShotForecaster(_NeedsList()) + assert model.predict(X, n=3).shape == (4, 3, 1) + assert model.context_format_ == "list" + + class _RejectsEverything: + def predict(self, inputs, prediction_length=1): + raise ValueError("nope") + + with pytest.raises(ValueError, match="rejected every supported way"): + ZeroShotForecaster(_RejectsEverything()).predict(X, n=3) + + +def test_zero_shot_forecaster_explicit_horizon_axis(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _HorizonFirst: + """Returns (n_series, horizon, n_quantiles), like predict_quantiles.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + levels = np.arange(5)[None, None, :] + return np.zeros((n_rows, prediction_length, 5)) + levels + + class _HorizonLast: + """Returns (n_series, n_quantiles, horizon), like Chronos-2.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + levels = np.arange(5)[None, :, None] + return np.zeros((n_rows, 5, prediction_length)) + levels + + # Both conventions are handled once the horizon axis is stated + for model, axis in [(_HorizonFirst(), 1), (_HorizonLast(), 2)]: + predicted = ZeroShotForecaster(model, horizon_axis=axis).predict(X, n=5) + np.testing.assert_allclose(predicted, np.full((4, 5, 1), 2.0)) + + # Negative indices are accepted + np.testing.assert_allclose( + ZeroShotForecaster(_HorizonLast(), horizon_axis=-1).predict(X, n=5), + np.full((4, 5, 1), 2.0), + ) + + # A wrong axis is reported rather than silently mis-slicing + with pytest.raises(ValueError, match="cannot hold a horizon"): + ZeroShotForecaster(_HorizonLast(), horizon_axis=1).predict(X, n=7) + with pytest.raises(ValueError, match="out of range"): + ZeroShotForecaster(_HorizonLast(), horizon_axis=5).predict(X, n=3) + with pytest.raises(ValueError, match="out of range"): + ZeroShotForecaster(_HorizonLast(), horizon_axis=0).predict(X, n=3) + with pytest.raises(ValueError, match="must be an integer or 'auto'"): + ZeroShotForecaster(_HorizonLast(), horizon_axis="last").predict(X, n=3) + + +def test_zero_shot_forecaster_explicit_horizon_axis_lifts_ambiguity(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _Ambiguous: + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + # (n_series, horizon, horizon), the second axis being the horizon + steps = np.arange(prediction_length)[None, :, None] + return np.zeros((n_rows, prediction_length, prediction_length)) + steps + + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("error") + predicted = ZeroShotForecaster(_Ambiguous(), horizon_axis=1).predict(X, n=4) + np.testing.assert_allclose( + predicted, np.tile(np.arange(4.0)[None, :, None], (4, 1, 1)) + ) + + +def test_zero_shot_forecaster_warns_on_ambiguous_output_shape(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _Ambiguous: + """Returns as many quantiles as there are forecast steps.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + return np.zeros((n_rows, prediction_length, prediction_length)) + + with pytest.warns(UserWarning, match="Set `horizon_axis` explicitly"): + ZeroShotForecaster(_Ambiguous()).predict(X, n=5) + + # A horizon of one is not ambiguous, singleton axes are simply reduced + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert ZeroShotForecaster(_Ambiguous()).predict(X, n=1).shape == (4, 1, 1) + + +def test_zero_shot_forecaster_reduces_quantiles_to_a_point_forecast(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _Quantiles: + """Returns (n_series, n_quantiles, horizon), like Chronos-2.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + levels = np.arange(9)[None, :, None] + return np.zeros((n_rows, 9, prediction_length)) + levels + + predicted = ZeroShotForecaster(_Quantiles()).predict(X, n=4) + # The median of 0..8 is 4, whatever the horizon + np.testing.assert_allclose(predicted, np.full((4, 4, 1), 4.0)) + np.testing.assert_allclose( + ZeroShotForecaster(_Quantiles(), quantile=None).predict(X, n=4), + np.full((4, 4, 1), 4.0), + ) + np.testing.assert_allclose( + ZeroShotForecaster(_Quantiles(), quantile=0.0).predict(X, n=4), + np.zeros((4, 4, 1)), + ) + + +def test_zero_shot_forecaster_errors(): + X = _dataset(n_ts=4, sz=32, d=1) + with pytest.raises(ValueError, match="`X` is required"): + ZeroShotForecaster(_DummyPipeline()).predict(None) + with pytest.raises(ValueError, match="positive integer"): + ZeroShotForecaster(_DummyPipeline()).predict(X, n=0) + with pytest.raises(ValueError, match="forecasting method"): + ZeroShotForecaster(object()).predict(X, n=1) + with pytest.raises(ValueError, match="input_layout"): + ZeroShotForecaster(_DummyPipeline(), input_layout="bogus").predict(X) + + +# --------------------------------------------------------------------------- +# LinearProbeForecaster +# --------------------------------------------------------------------------- + + +def test_linear_probe_forecaster(): + X = _dataset(n_ts=8, sz=64, d=1) + model = LinearProbeForecaster( + _DummyBackbone(), context_length=16, horizon=4, stride=4 + ) + model.fit(X) + # 8 series x windows of span 20 taken every 4 timestamps + assert model.n_windows_ == 8 * len(range(0, 64 - 20 + 1, 4)) + predicted = model.predict(X) + assert predicted.shape == (8, 4, 1) + # A shorter horizon truncates the forecast rather than refitting + np.testing.assert_allclose(model.predict(X, n=2), predicted[:, :2]) + + +def test_linear_probe_forecaster_multivariate(): + X = _dataset(n_ts=6, sz=48, d=2) + model = LinearProbeForecaster( + _DummyBackbone(), context_length=16, horizon=3, stride=8 + ).fit(X) + assert model.predict(X).shape == (6, 3, 2) + assert model.embedder_.embedding_size_ == 2 * D_MODEL + + +def test_linear_probe_forecaster_learns_something(): + # On a dataset of pure sine waves with varying phases, a probe on top of a + # frozen encoder should do clearly better than forecasting the last value. + rng = np.random.RandomState(0) + t = np.linspace(0, 8 * np.pi, 96) + phases = rng.uniform(0, 2 * np.pi, size=40) + X = np.sin(t[None, :] + phases[:, None])[:, :, None] + + model = LinearProbeForecaster( + _DummyBackbone(), context_length=16, horizon=4, stride=2 + ).fit(X[:30]) + assert model.score(X[30:]) > 0.5 + + +def test_linear_probe_forecaster_accepts_any_probe(): + X = _dataset(n_ts=6, sz=48, d=1) + model = LinearProbeForecaster( + _DummyBackbone(), + probe=Ridge(alpha=1.0), + context_length=16, + horizon=2, + stride=8, + ).fit(X) + assert isinstance(model.probe_, Ridge) + # The passed estimator is cloned, not fitted in place + assert model.probe_ is not model.probe + assert not hasattr(model.probe, "coef_") + + +def test_linear_probe_forecaster_errors(): + X = _dataset(n_ts=4, sz=48, d=1) + with pytest.raises(ValueError, match="context_length \\+ horizon"): + LinearProbeForecaster( + _DummyBackbone(), context_length=40, horizon=16 + ).fit(X) + with pytest.raises(ValueError, match="positive integer"): + LinearProbeForecaster(_DummyBackbone(), horizon=0).fit(X) + + model = LinearProbeForecaster( + _DummyBackbone(), context_length=16, horizon=4, stride=8 + ).fit(X) + with pytest.raises(ValueError, match="horizon of 4"): + model.predict(X, n=8) + with pytest.raises(ValueError, match="context_length"): + model.predict(_dataset(n_ts=4, sz=8, d=1)) + with pytest.raises(ValueError, match="features"): + model.predict(_dataset(n_ts=4, sz=48, d=3)) + + +def test_linear_probe_forecaster_fit_predict(): + X = _dataset(n_ts=6, sz=48, d=1) + model = LinearProbeForecaster( + _DummyBackbone(), context_length=16, horizon=2, stride=8 + ) + np.testing.assert_allclose(model.fit_predict(X, n=2), model.predict(X, n=2)) + + +# --------------------------------------------------------------------------- +# LinearProbeClassifier +# --------------------------------------------------------------------------- + + +def _classification_dataset(n_per_class=15, sz=32, seed=0): + rng = np.random.RandomState(seed) + t = np.linspace(0, 4 * np.pi, sz) + sines = np.sin(t)[None, :] + 0.1 * rng.randn(n_per_class, sz) + lines = np.linspace(-1, 1, sz)[None, :] + 0.1 * rng.randn(n_per_class, sz) + X = np.concatenate([sines, lines])[:, :, None] + y = np.array(["sine"] * n_per_class + ["line"] * n_per_class) + return X, y + + +def test_linear_probe_classifier(): + X, y = _classification_dataset() + model = LinearProbeClassifier(_DummyBackbone()).fit(X, y) + + assert set(model.classes_) == {"sine", "line"} + assert model.predict(X).shape == (len(y),) + assert model.predict_proba(X).shape == (len(y), 2) + assert model.decision_function(X).shape == (len(y),) + np.testing.assert_allclose(model.predict_proba(X).sum(axis=1), 1.0) + # These two classes are easily separable + assert model.score(X, y) > 0.9 + + +def test_linear_probe_classifier_multivariate(): + X, y = _classification_dataset() + X = np.concatenate([X, X[::-1]], axis=2) + model = LinearProbeClassifier(_DummyBackbone()).fit(X, y) + assert model.embedder_.embedding_size_ == 2 * D_MODEL + assert model.predict(X).shape == (len(y),) + + +def test_linear_probe_classifier_accepts_any_probe(): + X, y = _classification_dataset() + model = LinearProbeClassifier(_DummyBackbone(), probe=LinearSVC()).fit(X, y) + assert isinstance(model.probe_, LinearSVC) + with pytest.raises(AttributeError, match="predict_proba"): + model.predict_proba(X) + assert model.decision_function(X).shape == (len(y),) + + +def test_linear_probe_classifier_layer_and_pooling(): + X, y = _classification_dataset() + for layer in (0, 1, None): + for pooling in ("mean", "max", "cls"): + model = LinearProbeClassifier( + _DummyBackbone(), layer=layer, pooling=pooling + ).fit(X, y) + assert model.predict(X).shape == (len(y),) + + +def test_linear_probe_classifier_transform(): + X, y = _classification_dataset() + model = LinearProbeClassifier(_DummyBackbone()).fit(X, y) + assert model.transform(X).shape == (len(y), D_MODEL) + + +def test_linear_probe_classifier_errors(): + X, y = _classification_dataset() + with pytest.raises(ValueError, match="inconsistent numbers of samples"): + LinearProbeClassifier(_DummyBackbone()).fit(X, y[:-1]) + with pytest.raises(ValueError, match="input_layout"): + LinearProbeClassifier(_DummyBackbone(), input_layout="bogus").fit(X, y) + + model = LinearProbeClassifier(_DummyBackbone()).fit(X, y) + with pytest.raises(ValueError, match="features"): + model.predict(np.concatenate([X, X], axis=2)) + + +def test_estimators_are_clonable(): + for estimator in ( + TimeSeriesFoundationEmbedder(_DummyBackbone()), + ZeroShotForecaster(_DummyPipeline()), + LinearProbeForecaster(_DummyBackbone()), + LinearProbeClassifier(_DummyBackbone()), + ): + cloned = clone(estimator) + assert cloned is not estimator + assert cloned.get_params().keys() == estimator.get_params().keys() diff --git a/tslearn/foundation/__init__.py b/tslearn/foundation/__init__.py new file mode 100644 index 00000000..ad336514 --- /dev/null +++ b/tslearn/foundation/__init__.py @@ -0,0 +1,37 @@ +""" +The :mod:`tslearn.foundation` module gathers estimators that re-use pre-trained +time series models, such as the ones published on the Hugging Face Hub, behind +the usual tslearn API. + +Three adaptation strategies are covered: + +* zero-shot forecasting, with :class:`ZeroShotForecaster`, which uses a + pre-trained forecaster as-is; +* linear probing for forecasting, with :class:`LinearProbeForecaster`; +* linear probing for classification, with :class:`LinearProbeClassifier`. + +The last two rely on :class:`TimeSeriesFoundationEmbedder`, which turns any +frozen PyTorch model into a feature extractor and lets one choose which layer +to read representations from and how to pool them. + +These estimators are deliberately agnostic to any specific model +implementation: they duck-type the wrapped object and offer escape hatches +(``predict_fn``, ``input_name``, ``layers_path``, ``input_layout``) for models +that depart from the most widespread conventions. + +Notes +----- + This module requires PyTorch, which is an optional dependency of tslearn. + +""" + +from ._classification import LinearProbeClassifier +from ._embedding import TimeSeriesFoundationEmbedder +from ._forecasting import LinearProbeForecaster, ZeroShotForecaster + +__all__ = [ + "LinearProbeClassifier", + "LinearProbeForecaster", + "TimeSeriesFoundationEmbedder", + "ZeroShotForecaster", +] diff --git a/tslearn/foundation/_classification.py b/tslearn/foundation/_classification.py new file mode 100644 index 00000000..abd9e441 --- /dev/null +++ b/tslearn/foundation/_classification.py @@ -0,0 +1,304 @@ +"""Re-use of pre-trained time series models for classification.""" + +import numpy as np + +from sklearn.base import BaseEstimator, ClassifierMixin, clone +from sklearn.linear_model import LogisticRegression +from sklearn.utils.multiclass import check_classification_targets +from sklearn.utils.validation import check_is_fitted + +from tslearn.bases import TimeSeriesMixin +from tslearn.utils import check_array, to_time_series_dataset + +from ._embedding import ( + LAYOUTS, + TimeSeriesFoundationEmbedder, + _check_probe_pooling, + _require_torch, +) + + +class LinearProbeClassifier(TimeSeriesMixin, ClassifierMixin, BaseEstimator): + """Classify time series with a linear head on a frozen pre-trained model. + + The pre-trained model is used as a frozen feature extractor and a linear + classifier is fitted on the resulting representations. This is the standard + protocol used to assess how much class information a pre-trained model has + learnt [1]_: because the head is linear and the backbone is frozen, the + accuracy it reaches measures how linearly separable the classes already are + in the representation space. + + It is also a practical classifier in its own right, as it needs no + backpropagation through the pre-trained model and therefore trains in + seconds even on models counting hundreds of millions of parameters. + + Parameters + ---------- + model : torch.nn.Module + A pre-trained model used as a frozen feature extractor. + probe : sklearn estimator or None (default: None) + The head fitted on top of the frozen representations. When None, a + :class:`sklearn.linear_model.LogisticRegression` is used. Any + scikit-learn classifier can be passed instead; note that using a + non-linear one makes the resulting accuracy an estimate of predictive + performance rather than of linear separability. Frozen + representations sometimes have a very small variance, in which case + the default regularization is too strong; passing + ``make_pipeline(StandardScaler(), LogisticRegression())`` is then + usually enough to fix it. + context_length : int or None (default: None) + When set, only the last ``context_length`` timestamps of each series + are fed to the model. + layer : int or None (default: None) + Layer to probe. When None, the model's output hidden state is used. + When an integer, a forward hook is placed on the corresponding block + of the model's layer stack. Intermediate layers often carry more + class information than the last ones, which tend to specialize + towards the pre-training objective, so this is worth tuning. + layers_path : str or None (default: None) + Dotted path to the stack of layers, e.g. ``"encoder.block"``. When + None, the stack is auto-detected. + pooling : {"mean", "max", "cls", "last", "flatten"} (default: "mean") + How token representations are aggregated into one vector per series. + A probe needs a flat feature matrix, so ``pooling=None`` is not + accepted here. + tokens : slice, (start, stop) pair or None (default: None) + Which tokens to keep before pooling, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. Restricting + the average to the tokens that actually represent the series, e.g. + ``tokens=(0, -2)`` for Chronos-2, avoids diluting it with class or + forecast tokens. + cls_index : int (default: 0) + Index of the token selected when ``pooling="cls"``. + input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") + How multivariate series are handed over to the model. + input_name : str or None (default: None) + Name of the ``forward`` argument receiving the context values. + model_kwargs : dict or None (default: None) + Extra keyword arguments passed to every ``forward`` call. + batch_size : int (default: 32) + Number of series embedded at once. + device : str or None (default: None) + Device on which inference is run. + verbose : int (default: 0) + When positive, prints progress information. + + Attributes + ---------- + embedder_ : TimeSeriesFoundationEmbedder + The frozen feature extractor. + probe_ : sklearn estimator + The fitted linear head. + classes_ : array of shape=(n_classes,) + Class labels known to the classifier. + n_features_in_ : int + Number of features (channels) of the series seen during fit. + + See Also + -------- + LinearProbeForecaster: Linear probing for time series forecasting. + TimeSeriesFoundationEmbedder: The underlying feature extractor. + + Examples + -------- + >>> model = LinearProbeClassifier(backbone, layer=-2) # doctest: +SKIP + >>> model.fit(X_train, y_train).score(X_test, y_test) # doctest: +SKIP + 0.93 + + References + ---------- + .. [1] G. Alain and Y. Bengio. Understanding intermediate layers using + linear classifier probes. ICLR Workshop, 2017. + + """ + + def __init__( + self, + model, + probe=None, + context_length=None, + layer=None, + layers_path=None, + pooling="mean", + tokens=None, + cls_index=0, + input_layout="univariate", + input_name=None, + model_kwargs=None, + batch_size=32, + device=None, + verbose=0, + ): + self.model = model + self.probe = probe + self.context_length = context_length + self.layer = layer + self.layers_path = layers_path + self.pooling = pooling + self.tokens = tokens + self.cls_index = cls_index + self.input_layout = input_layout + self.input_name = input_name + self.model_kwargs = model_kwargs + self.batch_size = batch_size + self.device = device + self.verbose = verbose + + def _make_embedder(self): + return TimeSeriesFoundationEmbedder( + model=self.model, + layer=self.layer, + layers_path=self.layers_path, + pooling=self.pooling, + tokens=self.tokens, + cls_index=self.cls_index, + input_layout=self.input_layout, + context_length=self.context_length, + input_name=self.input_name, + model_kwargs=self.model_kwargs, + batch_size=self.batch_size, + device=self.device, + verbose=self.verbose, + ) + + def _make_probe(self): + if self.probe is None: + return LogisticRegression(max_iter=1000) + return clone(self.probe) + + def _check_input(self, X): + if self.input_layout not in LAYOUTS: + raise ValueError( + f"`input_layout` must be one of {LAYOUTS}, got " + f"'{self.input_layout}'." + ) + _check_probe_pooling(self.pooling) + X = check_array(X, allow_nd=True, force_all_finite=True) + return to_time_series_dataset(X) + + def fit(self, X, y): + """Fit a linear classifier on top of the frozen pre-trained model. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + y : array-like of shape=(n_ts,) + Class labels. + + Returns + ------- + self + The fitted estimator + + """ + _require_torch() + X = self._check_input(X) + y = np.asarray(y) + check_classification_targets(y) + if len(y) != X.shape[0]: + raise ValueError( + f"X and y have inconsistent numbers of samples: {X.shape[0]} " + f"and {len(y)}." + ) + + self.embedder_ = self._make_embedder() + embeddings = self.embedder_.fit_transform(X) + + self.probe_ = self._make_probe() + self.probe_.fit(embeddings, y) + + self.classes_ = np.asarray(self.probe_.classes_) + self.n_features_in_ = X.shape[2] + return self + + def _transform(self, X): + check_is_fitted(self, "probe_") + X = self._check_input(X) + if X.shape[2] != self.n_features_in_: + raise ValueError( + f"Series with {self.n_features_in_} features were expected, got " + f"{X.shape[2]}." + ) + return self.embedder_.transform(X) + + def transform(self, X): + """Return the frozen representations the classifier operates on. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + + Returns + ------- + array of shape=(n_ts, embedding_size) + Frozen representations of the input series. + + """ + return self._transform(X) + + def predict(self, X): + """Predict the class of each time series. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + + Returns + ------- + array of shape=(n_ts,) + Predicted class labels. + + """ + return self.probe_.predict(self._transform(X)) + + def predict_proba(self, X): + """Predict class probabilities for each time series. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + + Returns + ------- + array of shape=(n_ts, n_classes) + Predicted class probabilities, ordered as ``classes_``. + + """ + check_is_fitted(self, "probe_") + if not hasattr(self.probe_, "predict_proba"): + raise AttributeError( + f"The probe {type(self.probe_).__name__} does not expose " + "`predict_proba`." + ) + return self.probe_.predict_proba(self._transform(X)) + + def decision_function(self, X): + """Return the decision function of the linear head. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + + Returns + ------- + array of shape=(n_ts,) or (n_ts, n_classes) + Confidence scores. + + """ + check_is_fitted(self, "probe_") + if not hasattr(self.probe_, "decision_function"): + raise AttributeError( + f"The probe {type(self.probe_).__name__} does not expose " + "`decision_function`." + ) + return self.probe_.decision_function(self._transform(X)) + + def __sklearn_tags__(self): + tags = super().__sklearn_tags__() + tags.input_tags.allow_nan = False + return tags diff --git a/tslearn/foundation/_embedding.py b/tslearn/foundation/_embedding.py new file mode 100644 index 00000000..d321a6a9 --- /dev/null +++ b/tslearn/foundation/_embedding.py @@ -0,0 +1,608 @@ +"""Model-agnostic extraction of representations from pre-trained time series models.""" + +import inspect +import warnings + +import numpy as np + +from sklearn.base import BaseEstimator, TransformerMixin +from sklearn.utils.validation import check_is_fitted + +from tslearn.bases import TimeSeriesMixin +from tslearn.utils import check_array, to_time_series_dataset + +try: + import torch +except ImportError: # pragma: no cover + torch = None + + +#: Names of ``forward`` arguments that are commonly used by pre-trained time +#: series models to receive the raw context values. +CANDIDATE_INPUT_NAMES = ( + "context", + "past_values", + "past_target", + "input_values", + "inputs", + "x_enc", + "series", + "x", +) + +#: Attribute names under which model outputs commonly expose hidden states. +CANDIDATE_HIDDEN_STATE_NAMES = ( + "last_hidden_state", + "hidden_states", + "encoder_last_hidden_state", +) + +#: Supported ways of aggregating token representations. ``None``, for which +#: ``"none"`` is accepted as an alias, keeps them all and turns the estimator +#: into a time series to time series transform. +POOLINGS = ("mean", "max", "cls", "last", "flatten", "none", None) + +#: Supported layouts for the array handed over to the wrapped model. +LAYOUTS = ("univariate", "channels_last", "channels_first") + + +def _normalize_pooling(pooling): + return None if pooling == "none" else pooling + + +def _check_probe_pooling(pooling): + """Reject the poolings that do not yield a flat feature matrix.""" + if _normalize_pooling(pooling) is None: + raise ValueError( + "A linear probe needs one feature vector per series, so " + "`pooling=None` is not supported here. Use another pooling, or " + "TimeSeriesFoundationEmbedder directly if you want to keep one " + "representation per token." + ) + + +def _token_slice(tokens): + """Turn the ``tokens`` parameter into a slice over the token axis.""" + if tokens is None: + return slice(None) + if isinstance(tokens, slice): + return tokens + if isinstance(tokens, (tuple, list)) and len(tokens) == 2: + start, stop = tokens + if all(bound is None or isinstance(bound, (int, np.integer)) + for bound in (start, stop)): + return slice(start, stop) + raise ValueError( + "`tokens` must be None, a slice, or a (start, stop) pair of integers, " + f"got {tokens!r}." + ) + + +def _layout_to_model_input(X, layout): + """Lay a ``(n_ts, sz, d)`` dataset out as expected by the wrapped model. + + With the ``"univariate"`` layout, channels are unfolded along the batch + axis so that the returned array has shape ``(n_ts * d, sz)``, series + ``i`` and channel ``k`` sitting at row ``i * d + k``. + """ + n_ts, sz, d = X.shape + if layout == "univariate": + return np.ascontiguousarray(np.swapaxes(X, 1, 2)).reshape(n_ts * d, sz) + if layout == "channels_first": + return np.ascontiguousarray(np.swapaxes(X, 1, 2)) + if layout == "channels_last": + return np.ascontiguousarray(X) + raise ValueError(f"`input_layout` must be one of {LAYOUTS}, got '{layout}'.") + + +def _require_torch(): + if torch is None: # pragma: no cover + raise ImportError( + "PyTorch is required by the tslearn.foundation module. " + "Install it with `pip install torch`." + ) + + +def _resolve_attribute_path(model, path): + """Return the sub-module of ``model`` designated by a dotted ``path``.""" + module = model + for attribute in path.split("."): + if attribute.isdigit(): + module = module[int(attribute)] + else: + if not hasattr(module, attribute): + raise AttributeError( + f"Could not resolve `layers_path='{path}'`: the model has no " + f"attribute '{attribute}'." + ) + module = getattr(module, attribute) + return module + + +def _autodetect_layers(model): + """Find the container of repeated blocks of a transformer-like model. + + The heuristic looks for the longest :class:`torch.nn.ModuleList` whose + children all share the same type, which is how stacks of identical + encoder/decoder blocks are declared in essentially every implementation + published on the Hugging Face Hub. + """ + _require_torch() + best_path, best_modules = None, None + for path, module in model.named_modules(): + if not isinstance(module, torch.nn.ModuleList) or len(module) < 2: + continue + types = {type(block) for block in module} + if len(types) > 1: + continue + if best_modules is None or len(module) > len(best_modules): + best_path, best_modules = path, module + if best_modules is None: + raise ValueError( + "Could not automatically locate the stack of layers of the provided " + "model. Pass an explicit `layers_path` (e.g. 'encoder.block') to " + "select the layers to probe." + ) + return best_path, best_modules + + +def _as_tensor(output): + """Extract the hidden state tensor out of an arbitrary module output.""" + _require_torch() + if isinstance(output, torch.Tensor): + return output + for name in CANDIDATE_HIDDEN_STATE_NAMES: + value = getattr(output, name, None) + if isinstance(value, torch.Tensor): + return value + if isinstance(value, (tuple, list)) and len(value) > 0: + return value[-1] + if isinstance(output, dict): + for name in CANDIDATE_HIDDEN_STATE_NAMES: + value = output.get(name) + if isinstance(value, torch.Tensor): + return value + if isinstance(output, (tuple, list)): + for item in output: + if isinstance(item, torch.Tensor) and item.ndim == 3: + return item + for item in output: + if isinstance(item, torch.Tensor): + return item + raise TypeError( + f"Could not extract a hidden state tensor out of an output of type " + f"{type(output).__name__}. Models that only return their task-specific " + "outputs, such as forecasts, expose no hidden state to pool: select an " + "intermediate layer with the `layer` parameter, so that " + "representations are read through a forward hook instead." + ) + + +class TimeSeriesFoundationEmbedder(TimeSeriesMixin, TransformerMixin, BaseEstimator): + """Turn a frozen pre-trained time series model into a feature extractor. + + This transformer runs a (typically pre-trained) PyTorch model in inference + mode and pools its hidden states into a single vector per time series. It + is the building block used by + :class:`~tslearn.foundation.LinearProbeForecaster` and + :class:`~tslearn.foundation.LinearProbeClassifier`, and can also be used on + its own, for instance inside a :class:`sklearn.pipeline.Pipeline`. + + The transformer is deliberately agnostic to any particular model + implementation. It only assumes that: + + * ``model`` is a :class:`torch.nn.Module`; + * its ``forward`` method accepts the raw context values through one + argument (auto-detected among common names, see ``input_name``); + * it returns, or internally computes, hidden states of shape + ``(batch, n_tokens, dim)``. + + Parameters + ---------- + model : torch.nn.Module + A pre-trained model used as a frozen feature extractor. Its + parameters are never updated by this class. + layer : int or None (default: None) + Which layer to read representations from. When None, the output of + the model itself is used, which is usually the last hidden state + after the final normalization layer. When an integer is given, a + forward hook is placed on the corresponding block of the model's + layer stack, so that ``layer=0`` probes the first block and + ``layer=-1`` the last one. Probing intermediate layers is often + preferable, as the last layers of a pre-trained model tend to + specialize towards its pre-training objective [1]_. + layers_path : str or None (default: None) + Dotted path to the module holding the stack of layers, e.g. + ``"encoder.block"``. Only used when ``layer`` is an integer. When + None, the stack is auto-detected as the longest list of identical + sub-modules. + pooling : {"mean", "max", "cls", "last", "flatten", None} (default: "mean") + How to aggregate the ``n_tokens`` representations of a series into a + single vector. ``"cls"`` selects the single token at index + ``cls_index``, which is how models exposing a class (or register) + token are usually probed. ``"last"`` selects the last token and + ``"flatten"`` concatenates them all, which yields a much larger, and + context-length dependent, feature vector. + + ``None``, for which the string ``"none"`` is accepted as an alias, + applies no pooling at all: :meth:`transform` then returns a time + series dataset of shape ``(n_ts, n_tokens, d * dim)`` rather than a + flat feature matrix, which turns this estimator into a time series to + time series transform. Combined with ``tokens``, this is a way to + obtain one representation per timestep, or per patch of timesteps, + that can be fed to any other tslearn estimator. Note that the + resulting series are not accepted by the linear probes of this + module, which need a flat feature matrix. + tokens : slice, (start, stop) pair or None (default: None) + Which tokens to keep before pooling. Models often emit tokens that do + not represent the input series, such as class, register or forecast + tokens; excluding them is usually what one wants, both to build clean + per-timestep representations and to avoid polluting an average. + Chronos-2, for instance, emits its context tokens first, so + ``tokens=(0, -2)`` drops its trailing register and forecast tokens. + When None, all tokens are kept. This parameter does not affect + ``pooling="cls"``, whose whole purpose is to reach a token that + ``tokens`` would typically exclude. + cls_index : int (default: 0) + Index of the token to select when ``pooling="cls"``, applied to the + model's full token sequence, before any ``tokens`` selection. Note + that not all models place their class token first: Chronos-2, for + instance, inserts a register token between the context and forecast + tokens. + input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") + How multivariate series are handed over to the model. + ``"univariate"`` embeds every channel independently, as a + ``(n_ts * d, sz)`` array, and concatenates the resulting per-channel + embeddings, which is what most time series foundation models expect. + The other two layouts feed a ``(n_ts, sz, d)`` or ``(n_ts, d, sz)`` + array respectively, for models that are natively multivariate. + context_length : int or None (default: None) + When set, only the last ``context_length`` timestamps of each series + are fed to the model. + input_name : str or None (default: None) + Name of the ``forward`` argument receiving the context values. When + None, it is auto-detected among common names. + model_kwargs : dict or None (default: None) + Extra keyword arguments passed to every ``forward`` call. + batch_size : int (default: 32) + Number of series embedded at once. + device : str or None (default: None) + Device on which inference is run, e.g. ``"cuda"``. When None, the + device the model already lives on is used. + verbose : int (default: 0) + When positive, prints progress information. + + Attributes + ---------- + n_features_in_ : int + Number of features (channels) of the series seen during fit. + embedding_size_ : int + Dimension of the produced embeddings. When ``pooling`` is None, this + is the number of features of each timestep of the produced series. + n_tokens_ : int or None + Number of tokens kept per series when ``pooling`` is None, that is, + the length of the produced series. None otherwise. + + Notes + ----- + This estimator does not support variable length time series, as the + NaN padding used by tslearn to represent them would propagate through + most models. + + Examples + -------- + >>> import torch # doctest: +SKIP + >>> embedder = TimeSeriesFoundationEmbedder(model, layer=-2) # doctest: +SKIP + >>> embedder.fit_transform(X).shape # doctest: +SKIP + (10, 512) + + Keeping one representation per token, as a time series dataset: + + >>> embedder = TimeSeriesFoundationEmbedder( # doctest: +SKIP + ... model, layer=-2, pooling=None, tokens=(0, -2)) + >>> embedder.fit_transform(X).shape # doctest: +SKIP + (10, 16, 512) + + References + ---------- + .. [1] A. Kolesnikov, L. Beyer, X. Zhai, et al. Big Transfer (BiT): General + Visual Representation Learning. ECCV, 2020. + + """ + + def __init__( + self, + model, + layer=None, + layers_path=None, + pooling="mean", + tokens=None, + cls_index=0, + input_layout="univariate", + context_length=None, + input_name=None, + model_kwargs=None, + batch_size=32, + device=None, + verbose=0, + ): + self.model = model + self.layer = layer + self.layers_path = layers_path + self.pooling = pooling + self.tokens = tokens + self.cls_index = cls_index + self.input_layout = input_layout + self.context_length = context_length + self.input_name = input_name + self.model_kwargs = model_kwargs + self.batch_size = batch_size + self.device = device + self.verbose = verbose + + def _validate_params_(self): + if self.pooling not in POOLINGS: + raise ValueError( + f"`pooling` must be one of {POOLINGS}, got '{self.pooling}'." + ) + _token_slice(self.tokens) + if self.input_layout not in LAYOUTS: + raise ValueError( + f"`input_layout` must be one of {LAYOUTS}, got " + f"'{self.input_layout}'." + ) + _require_torch() + if not isinstance(self.model, torch.nn.Module): + raise TypeError( + "`model` must be a torch.nn.Module, got " + f"{type(self.model)}. Pre-trained models are typically obtained " + "through a `from_pretrained` call." + ) + + @property + def _device(self): + if self.device is not None: + return torch.device(self.device) + try: + return next(self.model.parameters()).device + except StopIteration: # pragma: no cover + return torch.device("cpu") + + def _resolve_input_name(self): + if self.input_name is not None: + return self.input_name + try: + signature = inspect.signature(self.model.forward) + except (TypeError, ValueError): # pragma: no cover + return CANDIDATE_INPUT_NAMES[0] + parameters = signature.parameters + for name in CANDIDATE_INPUT_NAMES: + if name in parameters: + return name + for name, parameter in parameters.items(): + if name == "self": + continue + if parameter.kind in ( + parameter.POSITIONAL_ONLY, + parameter.POSITIONAL_OR_KEYWORD, + ): + return name + raise ValueError( # pragma: no cover + "Could not determine which argument of the model's `forward` method " + "should receive the time series. Pass an explicit `input_name`." + ) + + def _resolve_layers(self): + if self.layers_path is not None: + return self.layers_path, _resolve_attribute_path( + self.model, self.layers_path + ) + return _autodetect_layers(self.model) + + def _forward_hidden_states(self, batch): + """Run the model on a batch and return hidden states of shape + (batch, n_tokens, dim).""" + kwargs = dict(self.model_kwargs or {}) + kwargs[self._input_name_] = batch + + if self.layer is None: + with torch.no_grad(): + output = self.model(**kwargs) + return _as_tensor(output) + + _, layers = self._resolve_layers() + try: + layer_module = layers[self.layer] + except IndexError: + raise ValueError( + f"`layer={self.layer}` is out of range: the model's layer stack " + f"holds {len(layers)} layers." + ) + + captured = {} + + def hook(_module, _inputs, output): + captured["hidden_states"] = _as_tensor(output) + + handle = layer_module.register_forward_hook(hook) + try: + with torch.no_grad(): + self.model(**kwargs) + finally: + handle.remove() + + if "hidden_states" not in captured: # pragma: no cover + raise RuntimeError( + "The forward hook was never triggered; the selected layer does " + "not seem to take part in the model's forward pass." + ) + return captured["hidden_states"] + + def _pool(self, hidden_states): + if hidden_states.ndim == 2: + # Already pooled by the model itself + return hidden_states + if hidden_states.ndim != 3: + raise ValueError( + "Expected hidden states of shape (batch, n_tokens, dim), got " + f"shape {tuple(hidden_states.shape)}." + ) + pooling = _normalize_pooling(self.pooling) + + # The class token is deliberately read before any token selection, as + # it is usually one of the tokens `tokens` is meant to filter out. + if pooling == "cls": + n_tokens = hidden_states.shape[1] + if not -n_tokens <= self.cls_index < n_tokens: + raise ValueError( + f"`cls_index={self.cls_index}` is out of range: the model " + f"produced {n_tokens} tokens." + ) + return hidden_states[:, self.cls_index] + + hidden_states = hidden_states[:, _token_slice(self.tokens)] + if hidden_states.shape[1] == 0: + raise ValueError( + f"`tokens={self.tokens!r}` selects no token at all out of the " + "sequence produced by the model." + ) + if pooling is None: + return hidden_states + if pooling == "mean": + return hidden_states.mean(dim=1) + if pooling == "max": + return hidden_states.max(dim=1).values + if pooling == "last": + return hidden_states[:, -1] + return hidden_states.reshape(hidden_states.shape[0], -1) + + def _check_input(self, X): + X = check_array(X, allow_nd=True, force_all_finite=True) + X = to_time_series_dataset(X) + if self.context_length is not None: + if X.shape[1] < self.context_length: + raise ValueError( + f"Series of length at least context_length=" + f"{self.context_length} are required, got {X.shape[1]}." + ) + X = X[:, -self.context_length :] + return X + + def _embed(self, X): + """Embed a checked (n_ts, sz, d) dataset. + + Returns a ``(n_ts, D)`` feature matrix, or a ``(n_ts, n_tokens, D)`` + time series dataset when ``pooling`` is None. + """ + n_ts, sz, d = X.shape + flat = _layout_to_model_input(X, self.input_layout) + + embeddings = [] + device = self._device + for start in range(0, flat.shape[0], self.batch_size): + chunk = flat[start : start + self.batch_size] + batch = torch.as_tensor(np.asarray(chunk, dtype=np.float32), device=device) + hidden_states = self._forward_hidden_states(batch) + embeddings.append(self._pool(hidden_states).to(torch.float32).cpu().numpy()) + if self.verbose: + print( + f"Embedded {min(start + self.batch_size, flat.shape[0])}" + f"/{flat.shape[0]} series" + ) + embeddings = np.concatenate(embeddings, axis=0) + + if self.input_layout == "univariate": + if embeddings.ndim == 3: + # (n_ts * d, n_tokens, dim) -> (n_ts, n_tokens, d * dim), so + # that the per-channel representations of a same timestep sit + # side by side, as in any other tslearn time series dataset. + n_tokens, dim = embeddings.shape[1:] + embeddings = embeddings.reshape(n_ts, d, n_tokens, dim) + embeddings = np.swapaxes(embeddings, 1, 2) + embeddings = embeddings.reshape(n_ts, n_tokens, d * dim) + else: + # Concatenate the per-channel embeddings of a same series + embeddings = embeddings.reshape(n_ts, -1) + return embeddings + + def fit(self, X, y=None): + """Check the model and the input data. No parameter is learnt. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + y : Ignored + + Returns + ------- + self + The fitted estimator + + """ + self._validate_params_() + X = self._check_input(X) + self.model.eval() + self._input_name_ = self._resolve_input_name() + self.n_features_in_ = X.shape[2] + self._sz_fit_ = X.shape[1] + # A single series is enough to know the shape of the representations + embedded = self._embed(X[:1]) + self.embedding_size_ = embedded.shape[-1] + self.n_tokens_ = embedded.shape[1] if embedded.ndim == 3 else None + return self + + def transform(self, X): + """Embed a time series dataset. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + + Returns + ------- + array of shape=(n_ts, embedding_size_), or + (n_ts, n_tokens_, embedding_size_) when ``pooling`` is None + Frozen representations of the input series. + + """ + check_is_fitted(self, "embedding_size_") + X = self._check_input(X) + if X.shape[2] != self.n_features_in_: + raise ValueError( + f"Series with {self.n_features_in_} features were expected, got " + f"{X.shape[2]}." + ) + if self.pooling == "flatten" and X.shape[1] != self._sz_fit_: + warnings.warn( + "pooling='flatten' produces context-length dependent features; " + f"series of length {self._sz_fit_} were seen at fit time and " + f"series of length {X.shape[1]} are being transformed." + ) + return self._embed(X) + + def fit_transform(self, X, y=None, **fit_params): + """Fit the estimator and embed a time series dataset. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + y : Ignored + + Returns + ------- + array of shape=(n_ts, embedding_size_), or + (n_ts, n_tokens_, embedding_size_) when ``pooling`` is None + Frozen representations of the input series. + + """ + return self.fit(X, y).transform(X) + + def __sklearn_tags__(self): + tags = super().__sklearn_tags__() + tags.target_tags.required = False + tags.non_deterministic = False + tags.input_tags.allow_nan = False + return tags diff --git a/tslearn/foundation/_forecasting.py b/tslearn/foundation/_forecasting.py new file mode 100644 index 00000000..2911976b --- /dev/null +++ b/tslearn/foundation/_forecasting.py @@ -0,0 +1,806 @@ +"""Re-use of pre-trained time series models for forecasting.""" + +import inspect +import warnings + +import numpy as np + +from sklearn.base import BaseEstimator, clone +from sklearn.linear_model import RidgeCV +from sklearn.utils.validation import check_is_fitted + +from tslearn.bases import TimeSeriesMixin +from tslearn.utils import check_array, to_time_series_dataset + +from ._embedding import ( + LAYOUTS, + TimeSeriesFoundationEmbedder, + _check_probe_pooling, + _layout_to_model_input, + _require_torch, +) + +try: + import torch +except ImportError: # pragma: no cover + torch = None + + +#: Method names commonly exposed by pre-trained forecasters, by decreasing +#: order of preference. +CANDIDATE_PREDICT_METHODS = ( + "predict", + "forecast", + "predict_quantiles", + "generate", +) + +#: Names of the argument through which a forecast horizon is commonly passed. +CANDIDATE_HORIZON_NAMES = ( + "prediction_length", + "horizon", + "forecast_horizon", + "prediction_horizon", + "num_steps", + "n_steps", + "steps", + "h", + "n", +) + +#: Attribute names under which forecasts are commonly returned. +CANDIDATE_FORECAST_NAMES = ( + "quantile_preds", + "prediction_outputs", + "predictions", + "sequences", + "mean", +) + + +#: Ways of presenting a batch of univariate contexts to a model, by decreasing +#: order of preference. Implementations disagree on this: Chronos-Bolt expects a +#: ``(batch, length)`` array while Chronos-2 requires an explicit variate axis, +#: and GluonTS-derived pipelines take a list of series. The first representation +#: the model accepts is used, and remembered in the ``context_format_`` +#: attribute. +CONTEXT_FORMATS = ("2d", "3d", "list") + + +def _format_context(context, context_format): + """Present a ``(n_rows, sz)`` batch of univariate contexts to a model.""" + if context_format == "2d": + return context + if context_format == "3d": + # (n_rows, n_variates=1, sz) + return context[:, None, :] + if context_format == "list": + return [row for row in context] + raise ValueError( # pragma: no cover + f"`context_format` must be one of {CONTEXT_FORMATS}, got " + f"'{context_format}'." + ) + + +def _to_numpy(value): + if torch is not None and isinstance(value, torch.Tensor): + return value.detach().to(torch.float32).cpu().numpy() + return np.asarray(value, dtype=np.float64) + + +def _unwrap_forecast(output): + """Extract a numeric array out of an arbitrary forecasting output.""" + for name in CANDIDATE_FORECAST_NAMES: + value = getattr(output, name, None) + if value is not None and not callable(value): + return _to_numpy(value) + if isinstance(output, dict): + for name in CANDIDATE_FORECAST_NAMES: + if name in output: + return _to_numpy(output[name]) + if isinstance(output, (list, tuple)): + if len(output) == 0: + raise ValueError("The model returned an empty forecast.") + items = [_to_numpy(item) for item in output] + shapes = {item.shape for item in items} + if len(shapes) == 1: + # A list with one entry per series, as returned e.g. by pipelines + return np.stack(items, axis=0) + # A tuple of heterogeneous outputs: keep the first one + return items[0] + return _to_numpy(output) + + +def _resolve_horizon_axis(shape, horizon, horizon_axis): + """Locate the axis of a forecast array that holds the forecast horizon.""" + ndim = len(shape) + if horizon_axis != "auto": + axis = horizon_axis + ndim if horizon_axis < 0 else horizon_axis + if not 1 <= axis < ndim: + raise ValueError( + f"`horizon_axis={horizon_axis}` is out of range for a forecast " + f"of shape {shape}; it must designate one of the axes 1 to " + f"{ndim - 1}, the axis 0 being the series." + ) + if shape[axis] < horizon: + raise ValueError( + f"`horizon_axis={horizon_axis}` designates an axis of size " + f"{shape[axis]} in a forecast of shape {shape}, which cannot " + f"hold a horizon of {horizon}." + ) + return axis + + # An axis whose size is exactly the horizon is a much safer guess than one + # that is merely large enough; among those, the last one is retained, as + # forecasts are most often returned with time as their trailing axis. + exact_axes = [axis for axis in range(1, ndim) if shape[axis] == horizon] + if exact_axes: + ambiguous = [axis for axis in exact_axes if shape[axis] > 1] + if len(ambiguous) > 1: + warnings.warn( + f"The forecast returned by the model has shape {shape}, in " + f"which several axes could be the horizon axis of length " + f"{horizon}; axis {exact_axes[-1]} was assumed. Set " + "`horizon_axis` explicitly to remove the ambiguity.", + UserWarning, + stacklevel=3, + ) + return exact_axes[-1] + + longer_axes = [axis for axis in range(1, ndim) if shape[axis] > horizon] + if not longer_axes: + raise ValueError( + f"Could not identify the forecast horizon axis in an output of " + f"shape {shape} for a horizon of {horizon}. Set `horizon_axis` " + "explicitly, or pass a `predict_fn` to control how the model is " + "called." + ) + return min(longer_axes, key=lambda axis: shape[axis]) + + +def _reduce_to_point_forecast(forecast, n_rows, horizon, quantile, horizon_axis): + """Reduce an arbitrarily shaped univariate forecast to (n_rows, horizon). + + Any axis that is neither the batch axis nor the horizon axis is understood + as holding quantile levels or sample paths, and is therefore reduced. + """ + forecast = np.asarray(forecast, dtype=np.float64) + if forecast.shape[0] != n_rows: + raise ValueError( + f"The model returned forecasts for {forecast.shape[0]} series while " + f"{n_rows} were provided." + ) + if forecast.ndim == 1: # pragma: no cover + forecast = forecast[:, None] + if forecast.ndim == 2: + if forecast.shape[1] < horizon: + raise ValueError( + f"The model returned forecasts of length {forecast.shape[1]}, " + f"which is shorter than the requested horizon {horizon}." + ) + return forecast[:, :horizon] + + axis = _resolve_horizon_axis(forecast.shape, horizon, horizon_axis) + forecast = np.moveaxis(forecast, axis, -1) + forecast = forecast.reshape(n_rows, -1, forecast.shape[-1]) + if forecast.shape[1] == 1: + forecast = forecast[:, 0] + elif quantile is None: + forecast = forecast.mean(axis=1) + else: + forecast = np.quantile(forecast, quantile, axis=1) + return forecast[:, :horizon] + + +class _BaseFoundationForecaster(TimeSeriesMixin, BaseEstimator): + """Shared input handling for pre-trained forecasters.""" + + def _check_input(self, X): + X = check_array(X, allow_nd=True, force_all_finite=True) + return to_time_series_dataset(X) + + def _check_layout(self): + if self.input_layout not in LAYOUTS: + raise ValueError( + f"`input_layout` must be one of {LAYOUTS}, got " + f"'{self.input_layout}'." + ) + horizon_axis = getattr(self, "horizon_axis", "auto") + if horizon_axis != "auto" and not isinstance( + horizon_axis, (int, np.integer) + ): + raise ValueError( + f"`horizon_axis` must be an integer or 'auto', got " + f"{horizon_axis!r}." + ) + + def fit_predict(self, X, y=None, n=1): + """Fit the estimator and forecast ``n`` timestamps for the given data. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + y : Ignored + n : int (default: 1) + The number of timestamps to forecast, a.k.a. the horizon. + + Returns + ------- + array of shape=(n_ts, n, d) + Array of forecasted timestamps + + """ + return self.fit(X, y).predict(X, n=n) + + def __sklearn_tags__(self): + tags = super().__sklearn_tags__() + tags.target_tags.required = False + tags.input_tags.allow_nan = False + return tags + + +class ZeroShotForecaster(_BaseFoundationForecaster): + """Forecast with a pre-trained model, without any training. + + Time series foundation models are pre-trained on large corpora of series + and are meant to forecast unseen series out of the box. This estimator + wraps such a model behind the usual tslearn forecasting API, so that it can + be compared to, or combined with, the other estimators of the library. + Calling :meth:`fit` is therefore not required and does not learn anything. + + The estimator makes no assumption about a particular implementation. It + only assumes that ``model`` exposes a forecasting method (auto-detected + among ``predict_quantiles``, ``predict``, ``forecast`` and ``generate``) + accepting the context values as its first argument and a forecast horizon + as a keyword argument. When this is not the case, or when finer control is + needed, pass an explicit ``predict_fn``. + + Parameters + ---------- + model : object + A pre-trained forecasting model or inference pipeline. + predict_fn : callable or None (default: None) + Called as ``predict_fn(model, context, horizon)``, where ``context`` + is a NumPy array laid out as prescribed by ``input_layout``, and + expected to return forecasts that can be broadcast to + ``(n_rows, horizon)``. When None, the model's own forecasting method + is auto-detected. + input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") + How the context is laid out when handed over to the model. + ``"univariate"`` forecasts every channel of every series + independently and feeds a ``(n_ts * d, sz)`` array, which is what + most time series foundation models expect. The other two layouts feed + a ``(n_ts, sz, d)`` or ``(n_ts, d, sz)`` array respectively, for + models that are natively multivariate. + context_length : int or None (default: None) + When set, only the last ``context_length`` timestamps of each series + are used as context. + horizon_axis : int or "auto" (default: "auto") + Which axis of the array returned by the model holds the forecast + horizon, axis 0 being the series. There is no shared convention here: + Chronos-2 returns ``(n_series, n_quantiles, horizon)`` while its own + ``predict_quantiles`` returns ``(n_series, horizon, n_quantiles)``, + so setting this explicitly, e.g. ``horizon_axis=-1``, is the reliable + option. When ``"auto"``, the axis is inferred from the returned + shape, preferring an axis whose size is exactly the horizon and, among + those, the last one. A warning is raised if the shape leaves the + choice ambiguous, which happens when the model returns as many + quantiles as there are forecast steps. All the axes that are neither + the series axis nor the horizon axis are taken to hold quantile levels + or sample paths, and are reduced according to ``quantile``. + quantile : float or None (default: 0.5) + Probabilistic models return several quantiles or sample paths; this + is the quantile used to derive a point forecast from them. When None, + the mean is used instead. Note that the quantile is taken over the + values the model returned, without assuming which level each of them + corresponds to, since that information is not exposed in any + standard way. For a model returning evenly spread quantile levels, + the default therefore recovers the median forecast, up to any + quantile crossing. + model_kwargs : dict or None (default: None) + Extra keyword arguments passed to the model's forecasting method. + warn_on_fit : bool (default: True) + Whether calling :meth:`fit` should emit a warning reminding that + nothing is being learnt. + + Attributes + ---------- + n_features_in_ : int + Number of features (channels) of the series seen during fit. + context_format_ : str + How the batch of contexts ended up being presented to the model, + among ``"2d"``, ``"3d"`` and ``"list"``. Set after the first call to + :meth:`predict` with the ``"univariate"`` layout, as implementations + disagree on this and the accepted one is discovered by trial. + + Notes + ----- + Unlike :class:`~tslearn.forecasting.VARIMA`, this estimator holds no + state about the fitted data, so :meth:`predict` requires ``X``. + + See Also + -------- + LinearProbeForecaster: Fit a linear head on top of a frozen pre-trained model. + + Examples + -------- + >>> from chronos import Chronos2Pipeline # doctest: +SKIP + >>> pipeline = Chronos2Pipeline.from_pretrained("amazon/chronos-2") # doctest: +SKIP + >>> model = ZeroShotForecaster(pipeline) # doctest: +SKIP + >>> model.predict(X, n=24).shape # doctest: +SKIP + (10, 24, 1) + + References + ---------- + .. [1] A. F. Ansari, L. Stella, C. Turkmen, et al. Chronos: Learning the + Language of Time Series. Transactions on Machine Learning Research, 2024. + + """ + + def __init__( + self, + model, + predict_fn=None, + input_layout="univariate", + context_length=None, + horizon_axis="auto", + quantile=0.5, + model_kwargs=None, + warn_on_fit=True, + ): + self.model = model + self.predict_fn = predict_fn + self.input_layout = input_layout + self.context_length = context_length + self.horizon_axis = horizon_axis + self.quantile = quantile + self.model_kwargs = model_kwargs + self.warn_on_fit = warn_on_fit + + def _resolve_predict_fn(self): + """Build a ``(context, horizon) -> forecast`` callable from the model.""" + if self.predict_fn is not None: + return lambda context, horizon: self.predict_fn( + self.model, context, horizon + ) + + for method_name in CANDIDATE_PREDICT_METHODS: + method = getattr(self.model, method_name, None) + if not callable(method): + continue + try: + parameters = inspect.signature(method).parameters + except (TypeError, ValueError): # pragma: no cover + continue + horizon_name = next( + (name for name in CANDIDATE_HORIZON_NAMES if name in parameters), + None, + ) + if horizon_name is None: + continue + + def call(context, horizon, method=method, horizon_name=horizon_name): + kwargs = dict(self.model_kwargs or {}) + kwargs[horizon_name] = horizon + return method(context, **kwargs) + + return call + + raise ValueError( + "Could not find a forecasting method on the provided model. The " + f"model is expected to expose one of {CANDIDATE_PREDICT_METHODS} " + "accepting a forecast horizon. Pass an explicit `predict_fn` " + "otherwise." + ) + + def _call_univariate(self, call, context, horizon): + """Call the model, negotiating how the context batch is presented. + + Implementations disagree on whether a batch of univariate series should + be handed over as a 2d array, as a 3d array with an explicit variate + axis, or as a list of series. Rather than requiring the user to know, + the candidate representations are tried in turn and the first one the + model accepts is remembered for subsequent calls. + """ + if self.predict_fn is not None: + return call(context, horizon) + + known_format = getattr(self, "context_format_", None) + formats = ( + (known_format,) if known_format is not None else CONTEXT_FORMATS + ) + errors = {} + for context_format in formats: + try: + forecast = call(_format_context(context, context_format), horizon) + except (ValueError, TypeError, IndexError, RuntimeError) as error: + errors[context_format] = f"{type(error).__name__}: {error}" + continue + self.context_format_ = context_format + return forecast + + details = "\n".join(f" - as a {key} input: {value}" for key, value in errors.items()) + raise ValueError( + "The model rejected every supported way of passing a batch of " + f"univariate contexts:\n{details}\nPass an explicit `predict_fn` to " + "control how the model is called." + ) + + def fit(self, X, y=None): + """Do nothing, as a zero-shot model needs no training. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + y : Ignored + + Returns + ------- + self + The estimator, ready to be used for prediction + + """ + self._check_layout() + X = self._check_input(X) + if self.warn_on_fit: + warnings.warn( + "ZeroShotForecaster.fit does not train anything: the wrapped " + "model is used as-is for zero-shot forecasting. Use " + "LinearProbeForecaster to fit a head on top of it, or pass " + "warn_on_fit=False to silence this warning.", + UserWarning, + stacklevel=2, + ) + self.n_features_in_ = X.shape[2] + self._is_fitted_ = True + return self + + def predict(self, X, n=1): + """Forecast ``n`` timestamps ahead of the given series. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset to forecast. + n : int (default: 1) + The number of timestamps to forecast, a.k.a. the horizon. + + Returns + ------- + array of shape=(n_ts, n, d) + Array of forecasted timestamps + + """ + self._check_layout() + if X is None: + raise ValueError( + "A zero-shot model keeps no state about the data seen at fit " + "time, so `X` is required at predict time." + ) + if n < 1: + raise ValueError(f"`n` must be a positive integer, got {n}.") + X = self._check_input(X) + if getattr(self, "n_features_in_", X.shape[2]) != X.shape[2]: + raise ValueError( + f"Series with {self.n_features_in_} features were expected, got " + f"{X.shape[2]}." + ) + if self.context_length is not None: + X = X[:, -self.context_length :] + + n_ts, sz, d = X.shape + context = _layout_to_model_input(X, self.input_layout) + call = self._resolve_predict_fn() + + if self.input_layout == "univariate": + forecast = _unwrap_forecast(self._call_univariate(call, context, n)) + forecast = _reduce_to_point_forecast( + forecast, n_ts * d, n, self.quantile, self.horizon_axis + ) + # (n_ts * d, n) -> (n_ts, n, d) + return np.swapaxes(forecast.reshape(n_ts, d, n), 1, 2) + + forecast = np.asarray( + _unwrap_forecast(call(context, n)), dtype=np.float64 + ) + if forecast.shape[:1] != (n_ts,): + raise ValueError( + f"The model returned forecasts for {forecast.shape[0]} series " + f"while {n_ts} were provided." + ) + if self.input_layout == "channels_first": + forecast = np.swapaxes(forecast, 1, 2) + if forecast.ndim != 3 or forecast.shape[2] != d: + raise ValueError( + f"Expected forecasts of shape (n_ts, n, d) = ({n_ts}, {n}, {d}) " + f"with input_layout='{self.input_layout}', got shape " + f"{forecast.shape}. Pass an explicit `predict_fn` to control " + "how the model is called." + ) + return forecast[:, :n] + + +class LinearProbeForecaster(_BaseFoundationForecaster): + """Forecast by fitting a linear head on a frozen pre-trained model. + + Linear probing keeps the pre-trained model entirely frozen and only fits a + linear map from its representations to the quantity of interest. Compared + to zero-shot use, it adapts the model to the data at hand at a very small + computational cost, and compared to fine-tuning, it leaves the pre-trained + weights untouched, which makes it both cheap and hard to overfit. + + Training pairs are cut out of the provided series with a sliding window: + the ``context_length`` timestamps preceding a cut point are embedded, and + the ``horizon`` timestamps following it are used as the target. + + Parameters + ---------- + model : torch.nn.Module + A pre-trained model used as a frozen feature extractor. + probe : sklearn estimator or None (default: None) + The head fitted on top of the frozen representations. When None, a + :class:`sklearn.linear_model.RidgeCV` is used, which selects its + regularization strength by cross-validation and admits a closed-form + solution. Any multi-output capable regressor can be passed instead; + note that using a non-linear one makes the name "probe" a misnomer, + and the resulting scores no longer measure linear separability of the + representations. Frozen representations sometimes have a very small + variance, in which case prepending a + :class:`sklearn.preprocessing.StandardScaler` to the head helps. + context_length : int (default: 128) + Number of timestamps fed to the pre-trained model. + horizon : int (default: 1) + Number of timestamps forecasted at once. This is fixed at fit time, + as one linear head is fitted for the whole horizon. + stride : int (default: 1) + Step between two consecutive training windows. Larger values yield + fewer, less redundant training pairs and a faster fit. + layer : int or None (default: None) + Layer to probe, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. + layers_path : str or None (default: None) + Dotted path to the stack of layers, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. + pooling : {"mean", "max", "cls", "last", "flatten"} (default: "mean") + How token representations are aggregated, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. Unlike + that class, a probe needs a flat feature matrix, so ``pooling=None`` + is not accepted here. + tokens : slice, (start, stop) pair or None (default: None) + Which tokens to keep before pooling, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. Restricting + the average to the tokens that actually represent the context, e.g. + ``tokens=(0, -2)`` for Chronos-2, avoids diluting it with class or + forecast tokens. + cls_index : int (default: 0) + Index of the token selected when ``pooling="cls"``. + input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") + How the context is laid out when handed over to the model. + input_name : str or None (default: None) + Name of the ``forward`` argument receiving the context values. + model_kwargs : dict or None (default: None) + Extra keyword arguments passed to every ``forward`` call. + batch_size : int (default: 32) + Number of series embedded at once. + device : str or None (default: None) + Device on which inference is run. + verbose : int (default: 0) + When positive, prints progress information. + + Attributes + ---------- + embedder_ : TimeSeriesFoundationEmbedder + The frozen feature extractor. + probe_ : sklearn estimator + The fitted linear head. + n_features_in_ : int + Number of features (channels) of the series seen during fit. + n_windows_ : int + Number of training windows cut out of the fit data. + + See Also + -------- + ZeroShotForecaster: Use a pre-trained model without any training. + LinearProbeClassifier: Linear probing for time series classification. + + Examples + -------- + >>> model = LinearProbeForecaster(backbone, horizon=24) # doctest: +SKIP + >>> model.fit(X_train).predict(X_test, n=24).shape # doctest: +SKIP + (10, 24, 1) + + References + ---------- + .. [1] G. Alain and Y. Bengio. Understanding intermediate layers using + linear classifier probes. ICLR Workshop, 2017. + + """ + + def __init__( + self, + model, + probe=None, + context_length=128, + horizon=1, + stride=1, + layer=None, + layers_path=None, + pooling="mean", + tokens=None, + cls_index=0, + input_layout="univariate", + input_name=None, + model_kwargs=None, + batch_size=32, + device=None, + verbose=0, + ): + self.model = model + self.probe = probe + self.context_length = context_length + self.horizon = horizon + self.stride = stride + self.layer = layer + self.layers_path = layers_path + self.pooling = pooling + self.tokens = tokens + self.cls_index = cls_index + self.input_layout = input_layout + self.input_name = input_name + self.model_kwargs = model_kwargs + self.batch_size = batch_size + self.device = device + self.verbose = verbose + + def _make_embedder(self): + return TimeSeriesFoundationEmbedder( + model=self.model, + layer=self.layer, + layers_path=self.layers_path, + pooling=self.pooling, + tokens=self.tokens, + cls_index=self.cls_index, + input_layout=self.input_layout, + context_length=None, + input_name=self.input_name, + model_kwargs=self.model_kwargs, + batch_size=self.batch_size, + device=self.device, + verbose=self.verbose, + ) + + def _make_probe(self): + if self.probe is None: + return RidgeCV(alphas=np.logspace(-3, 3, 13)) + return clone(self.probe) + + def _make_windows(self, X): + """Cut (context, target) pairs out of a (n_ts, sz, d) dataset.""" + n_ts, sz, d = X.shape + span = self.context_length + self.horizon + if sz < span: + raise ValueError( + f"Series of length at least context_length + horizon = {span} " + f"are required to build training windows, got {sz}." + ) + starts = np.arange(0, sz - span + 1, self.stride) + contexts = np.stack( + [X[:, start : start + self.context_length] for start in starts], axis=1 + ) + targets = np.stack( + [ + X[:, start + self.context_length : start + span] + for start in starts + ], + axis=1, + ) + # (n_ts, n_windows, ...) -> (n_ts * n_windows, ...) + contexts = contexts.reshape(n_ts * len(starts), self.context_length, d) + targets = targets.reshape(n_ts * len(starts), self.horizon * d) + return contexts, targets + + def _validate_params_(self): + self._check_layout() + _check_probe_pooling(self.pooling) + for name in ("context_length", "horizon", "stride"): + value = getattr(self, name) + if not isinstance(value, (int, np.integer)) or value < 1: + raise ValueError( + f"`{name}` must be a positive integer, got {value}." + ) + + def fit(self, X, y=None): + """Fit a linear head on top of the frozen pre-trained model. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset, with ``sz >= context_length + horizon``. + y : Ignored + + Returns + ------- + self + The fitted estimator + + """ + self._validate_params_() + _require_torch() + X = self._check_input(X) + contexts, targets = self._make_windows(X) + + self.embedder_ = self._make_embedder() + embeddings = self.embedder_.fit_transform(contexts) + + self.probe_ = self._make_probe() + self.probe_.fit(embeddings, targets) + + self.n_features_in_ = X.shape[2] + self.n_windows_ = contexts.shape[0] + return self + + def predict(self, X, n=None): + """Forecast ``horizon`` timestamps ahead of the given series. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset to forecast, with + ``sz >= context_length``. + n : int or None (default: None) + The number of timestamps to forecast. A single head is fitted for + the whole horizon, so ``n`` may not exceed the ``horizon`` value + used at fit time. When None, ``horizon`` is used. + + Returns + ------- + array of shape=(n_ts, n, d) + Array of forecasted timestamps + + """ + check_is_fitted(self, "probe_") + n = self.horizon if n is None else n + if not 1 <= n <= self.horizon: + raise ValueError( + f"This estimator was fitted for a horizon of {self.horizon}, so " + f"`n` must lie in [1, {self.horizon}], got {n}. Refit with a " + "larger `horizon` to forecast further ahead." + ) + X = self._check_input(X) + if X.shape[2] != self.n_features_in_: + raise ValueError( + f"Series with {self.n_features_in_} features were expected, got " + f"{X.shape[2]}." + ) + if X.shape[1] < self.context_length: + raise ValueError( + f"Series of length at least context_length=" + f"{self.context_length} are required, got {X.shape[1]}." + ) + contexts = X[:, -self.context_length :] + predictions = self.probe_.predict(self.embedder_.transform(contexts)) + predictions = np.asarray(predictions, dtype=np.float64) + predictions = predictions.reshape(X.shape[0], self.horizon, self.n_features_in_) + return predictions[:, :n] + + def score(self, X, y=None): + """Return the coefficient of determination of the forecast. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset, with ``sz >= context_length + horizon``. + y : Ignored + + Returns + ------- + float + :math:`R^2` of the forecast over all windows cut out of ``X``. + + """ + check_is_fitted(self, "probe_") + X = self._check_input(X) + contexts, targets = self._make_windows(X) + return self.probe_.score(self.embedder_.transform(contexts), targets) From c06668cbd001afa4d36c4f8a41ee6ee6856c05da Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Wed, 12 Aug 2026 17:38:26 +0200 Subject: [PATCH 02/18] better docs --- docs/gen_modules/tslearn.foundation.rst | 11 +++++++++++ tslearn/foundation/_classification.py | 12 ++++++------ tslearn/foundation/_embedding.py | 5 +++-- tslearn/foundation/_forecasting.py | 4 +++- 4 files changed, 23 insertions(+), 9 deletions(-) diff --git a/docs/gen_modules/tslearn.foundation.rst b/docs/gen_modules/tslearn.foundation.rst index f1e25911..a3787ed9 100644 --- a/docs/gen_modules/tslearn.foundation.rst +++ b/docs/gen_modules/tslearn.foundation.rst @@ -5,4 +5,15 @@ tslearn.foundation .. automodule:: tslearn.foundation + .. rubric:: Classes + + .. autosummary:: + :toctree: foundation + :template: class.rst + + ZeroShotForecaster + LinearProbeForecaster + LinearProbeClassifier + TimeSeriesFoundationEmbedder + \ No newline at end of file diff --git a/tslearn/foundation/_classification.py b/tslearn/foundation/_classification.py index abd9e441..518bb996 100644 --- a/tslearn/foundation/_classification.py +++ b/tslearn/foundation/_classification.py @@ -41,11 +41,9 @@ class LinearProbeClassifier(TimeSeriesMixin, ClassifierMixin, BaseEstimator): :class:`sklearn.linear_model.LogisticRegression` is used. Any scikit-learn classifier can be passed instead; note that using a non-linear one makes the resulting accuracy an estimate of predictive - performance rather than of linear separability. Frozen - representations sometimes have a very small variance, in which case - the default regularization is too strong; passing - ``make_pipeline(StandardScaler(), LogisticRegression())`` is then - usually enough to fix it. + performance rather than of linear separability. Frozen representations + sometimes have a very small variance, in which case prepending a + :class:`sklearn.preprocessing.StandardScaler` to the head helps. context_length : int or None (default: None) When set, only the last ``context_length`` timestamps of each series are fed to the model. @@ -101,7 +99,9 @@ class information than the last ones, which tend to specialize Examples -------- - >>> model = LinearProbeClassifier(backbone, layer=-2) # doctest: +SKIP + >>> from chronos import Chronos2Pipeline # doctest: +SKIP + >>> pipeline = Chronos2Pipeline.from_pretrained("amazon/chronos-2") # doctest: +SKIP + >>> model = LinearProbeClassifier(pipeline.model, layer=-2) # doctest: +SKIP >>> model.fit(X_train, y_train).score(X_test, y_test) # doctest: +SKIP 0.93 diff --git a/tslearn/foundation/_embedding.py b/tslearn/foundation/_embedding.py index d321a6a9..f16d7833 100644 --- a/tslearn/foundation/_embedding.py +++ b/tslearn/foundation/_embedding.py @@ -291,8 +291,9 @@ class TimeSeriesFoundationEmbedder(TimeSeriesMixin, TransformerMixin, BaseEstima Examples -------- - >>> import torch # doctest: +SKIP - >>> embedder = TimeSeriesFoundationEmbedder(model, layer=-2) # doctest: +SKIP + >>> from chronos import Chronos2Pipeline # doctest: +SKIP + >>> pipeline = Chronos2Pipeline.from_pretrained("amazon/chronos-2") # doctest: +SKIP + >>> embedder = TimeSeriesFoundationEmbedder(pipeline.model, layer=-2) # doctest: +SKIP >>> embedder.fit_transform(X).shape # doctest: +SKIP (10, 512) diff --git a/tslearn/foundation/_forecasting.py b/tslearn/foundation/_forecasting.py index 2911976b..a4ed1ac8 100644 --- a/tslearn/foundation/_forecasting.py +++ b/tslearn/foundation/_forecasting.py @@ -607,7 +607,9 @@ class LinearProbeForecaster(_BaseFoundationForecaster): Examples -------- - >>> model = LinearProbeForecaster(backbone, horizon=24) # doctest: +SKIP + >>> from chronos import Chronos2Pipeline # doctest: +SKIP + >>> pipeline = Chronos2Pipeline.from_pretrained("amazon/chronos-2") # doctest: +SKIP + >>> model = LinearProbeForecaster(pipeline.model, horizon=24) # doctest: +SKIP >>> model.fit(X_train).predict(X_test, n=24).shape # doctest: +SKIP (10, 24, 1) From 3135cdaee4c62c59ea293a36e947dc21ad6a50ed Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Wed, 12 Aug 2026 17:44:53 +0200 Subject: [PATCH 03/18] pick a better thumbnail for the gallery --- docs/examples/classification/plot_foundation_linear_probe.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/examples/classification/plot_foundation_linear_probe.py b/docs/examples/classification/plot_foundation_linear_probe.py index 0ac5c290..3c732832 100644 --- a/docs/examples/classification/plot_foundation_linear_probe.py +++ b/docs/examples/classification/plot_foundation_linear_probe.py @@ -26,6 +26,10 @@ Universal Forecasting. arXiv:2510.15821, 2025. """ +# Author: Romain Tavenard +# License: BSD 3 clause +# sphinx_gallery_thumbnail_number = 2 + ############################################################################## # Data # ---- From 108c6b2b4edf3ccd594ba1ad400c9a3683d5a491 Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Thu, 13 Aug 2026 09:40:39 +0200 Subject: [PATCH 04/18] better (doc)tests --- tests/test_foundation.py | 458 +++++++++++++++++++++++++++++ tslearn/foundation/_forecasting.py | 9 +- 2 files changed, 466 insertions(+), 1 deletion(-) diff --git a/tests/test_foundation.py b/tests/test_foundation.py index 5696e339..6f638c66 100644 --- a/tests/test_foundation.py +++ b/tests/test_foundation.py @@ -11,6 +11,7 @@ from sklearn.base import clone from sklearn.linear_model import Ridge +from sklearn.naive_bayes import GaussianNB from sklearn.pipeline import make_pipeline from sklearn.svm import LinearSVC @@ -409,6 +410,249 @@ def test_embedder_in_a_sklearn_pipeline(): assert clone(pipeline) is not pipeline +def test_embedder_channels_first_layout(): + X = _dataset(n_ts=4, sz=32, d=3) + seen = {} + + class _ChannelsFirstBackbone(torch.nn.Module): + """Expects a (batch, d, sz) input, like some ``transformers`` models.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, series): + seen["shape"] = tuple(series.shape) + return _Output(self.linear(series)) + + embedder = TimeSeriesFoundationEmbedder( + _ChannelsFirstBackbone(), input_layout="channels_first", pooling="mean" + ) + embeddings = embedder.fit_transform(X) + assert seen["shape"] == (4, 3, 32) + assert embeddings.shape == (4, D_MODEL) + + +def test_embedder_layers_path_with_numeric_segment(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _WrappedBackbone(torch.nn.Module): + """Nests the encoder in a top-level ``ModuleList``, as some models do, + requiring a numeric index to reach the stack of layers.""" + + def __init__(self, n_layers=3, d_model=D_MODEL, patch_size=PATCH_SIZE, seed=0): + super().__init__() + torch.manual_seed(seed) + self.patch_size = patch_size + self.embedding = torch.nn.Linear(patch_size, d_model) + self.register_token = torch.nn.Parameter(torch.randn(1, 1, d_model)) + self.stages = torch.nn.ModuleList([_Encoder(n_layers, d_model)]) + + def forward(self, context): + batch_size, sz = context.shape + n_patches = sz // self.patch_size + patches = context[:, : n_patches * self.patch_size] + patches = patches.reshape(batch_size, n_patches, self.patch_size) + hidden_states = self.embedding(patches) + register = self.register_token.expand(batch_size, -1, -1) + hidden_states = torch.cat([hidden_states, register], dim=1) + return _Output(self.stages[0](hidden_states)) + + backbone = _WrappedBackbone() + embeddings = TimeSeriesFoundationEmbedder( + backbone, layer=1, layers_path="stages.0.block" + ).fit_transform(X) + assert embeddings.shape == (3, D_MODEL) + + with pytest.raises(AttributeError, match="Could not resolve"): + TimeSeriesFoundationEmbedder( + backbone, layer=0, layers_path="stages.0.bogus" + ).fit(X) + + +def test_embedder_autodetect_layers_failure(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _NoLayerStack(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, x): + return _Output(self.linear(x).unsqueeze(1)) + + with pytest.raises(ValueError, match="Could not automatically locate"): + TimeSeriesFoundationEmbedder(_NoLayerStack(), layer=0).fit(X) + + +def test_embedder_accepts_dict_model_outputs(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _DictOutputBackbone(torch.nn.Module): + """Returns a plain dict, as some pipelines do instead of a ModelOutput.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, x): + return {"hidden_states": self.linear(x).unsqueeze(1)} + + embeddings = TimeSeriesFoundationEmbedder(_DictOutputBackbone()).fit_transform(X) + assert embeddings.shape == (3, D_MODEL) + + +def test_embedder_accepts_hidden_states_tuple_attribute(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _LayerStackOutput: + """Stands in for a ``transformers`` output exposing all-layer states.""" + + def __init__(self, layers): + self.hidden_states = layers + + class _AllLayersBackbone(torch.nn.Module): + def __init__(self): + super().__init__() + torch.manual_seed(0) + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, x): + hidden = self.linear(x).unsqueeze(1) + return _LayerStackOutput((hidden, 2 * hidden)) + + backbone = _AllLayersBackbone() + embeddings = TimeSeriesFoundationEmbedder(backbone).fit_transform(X) + with torch.no_grad(): + raw = backbone.linear(torch.as_tensor(X[:, :, 0], dtype=torch.float32)) + # Only the last element of the `hidden_states` tuple is read, matching + # transformers' convention that it holds one tensor per layer + np.testing.assert_allclose(embeddings, 2 * raw.numpy(), rtol=1e-5, atol=1e-6) + + +def test_embedder_accepts_plain_tuple_outputs(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _PlainTupleBackbone(torch.nn.Module): + """Returns a raw tuple, with no ``ModelOutput`` wrapper at all.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, x): + hidden = self.linear(x) + # The 3d sequence of hidden states must be preferred over the + # pooled 2d tensor that precedes it in the tuple + return (hidden, hidden.unsqueeze(1)) + + embeddings = TimeSeriesFoundationEmbedder(_PlainTupleBackbone()).fit_transform(X) + assert embeddings.shape == (3, D_MODEL) + + +def test_embedder_plain_tuple_falls_back_to_any_tensor(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _NoSequenceBackbone(torch.nn.Module): + """Returns an already-pooled 2d tensor wrapped in a 1-tuple.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, x): + return (self.linear(x),) + + embeddings = TimeSeriesFoundationEmbedder(_NoSequenceBackbone()).fit_transform(X) + assert embeddings.shape == (3, D_MODEL) + + +def test_embedder_rejects_unrecognized_model_output(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _BogusOutputBackbone(torch.nn.Module): + def forward(self, x): + return "not a tensor" + + with pytest.raises(TypeError, match="Could not extract a hidden state"): + TimeSeriesFoundationEmbedder(_BogusOutputBackbone()).fit(X) + + +def test_embedder_rejects_invalid_hidden_state_rank(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _WeirdRankBackbone(torch.nn.Module): + def forward(self, x): + return x[:, None, None, :] # 4d, not a valid hidden-state rank + + with pytest.raises(ValueError, match="Expected hidden states of shape"): + TimeSeriesFoundationEmbedder(_WeirdRankBackbone()).fit(X) + + +def test_embedder_explicit_device(): + X = _dataset(n_ts=3, sz=32, d=1) + embeddings = TimeSeriesFoundationEmbedder( + _DummyBackbone(), device="cpu" + ).fit_transform(X) + assert embeddings.shape == (3, D_MODEL) + + +def test_embedder_explicit_input_name(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _NamedArgBackbone(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, my_custom_arg): + return _Output(self.linear(my_custom_arg).unsqueeze(1)) + + embedder = TimeSeriesFoundationEmbedder( + _NamedArgBackbone(), input_name="my_custom_arg" + ) + embeddings = embedder.fit_transform(X) + assert embeddings.shape == (3, D_MODEL) + + +def test_embedder_resolves_input_name_by_position_when_unrecognized(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _PositionalArgBackbone(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + # `data` is not one of CANDIDATE_INPUT_NAMES, forcing the fallback to + # the first positional parameter of `forward` + def forward(self, data): + return _Output(self.linear(data).unsqueeze(1)) + + embedder = TimeSeriesFoundationEmbedder(_PositionalArgBackbone()) + embeddings = embedder.fit_transform(X) + assert embedder._input_name_ == "data" + assert embeddings.shape == (3, D_MODEL) + + +def test_embedder_verbose_prints_progress(capsys): + X = _dataset(n_ts=5, sz=32, d=1) + TimeSeriesFoundationEmbedder( + _DummyBackbone(), batch_size=2, verbose=1 + ).fit_transform(X) + captured = capsys.readouterr() + assert "Embedded 5/5 series" in captured.out + + +def test_embedder_flatten_pooling_warns_on_length_mismatch(): + X_fit = _dataset(n_ts=3, sz=32, d=1) + X_other = _dataset(n_ts=3, sz=48, d=1, seed=1) + embedder = TimeSeriesFoundationEmbedder( + _DummyBackbone(), pooling="flatten" + ).fit(X_fit) + with pytest.warns(UserWarning, match="pooling='flatten' produces context-length"): + embedder.transform(X_other) + + # --------------------------------------------------------------------------- # ZeroShotForecaster # --------------------------------------------------------------------------- @@ -638,6 +882,211 @@ def test_zero_shot_forecaster_errors(): ZeroShotForecaster(_DummyPipeline(), input_layout="bogus").predict(X) +def test_zero_shot_forecaster_feature_mismatch_after_fit(): + X = _dataset(n_ts=4, sz=32, d=1) + model = ZeroShotForecaster(_DummyPipeline(), warn_on_fit=False).fit(X) + with pytest.raises(ValueError, match="features"): + model.predict(_dataset(n_ts=4, sz=32, d=2), n=2) + + +def test_zero_shot_forecaster_skips_methods_without_a_horizon_argument(): + X = _dataset(n_ts=4, sz=32, d=1) + calls = [] + + class _PartialModel: + """Exposes a `predict` with no forecast-horizon argument, which must + be skipped in favor of `forecast`.""" + + def predict(self, inputs): + calls.append("predict") + raise AssertionError("predict should never be called") + + def forecast(self, inputs, prediction_length=1): + calls.append("forecast") + n_rows = len(inputs) + return np.zeros((n_rows, prediction_length)) + + predicted = ZeroShotForecaster(_PartialModel()).predict(X, n=3) + assert predicted.shape == (4, 3, 1) + assert calls == ["forecast"] + + +def test_zero_shot_forecaster_accepts_object_attribute_output(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _AttributeOutput: + """Stands in for an output object exposing the forecast as an attribute.""" + + def __init__(self, mean): + self.mean = mean + + class _ObjectOutputModel: + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + return _AttributeOutput(np.zeros((n_rows, prediction_length))) + + predicted = ZeroShotForecaster(_ObjectOutputModel()).predict(X, n=3) + assert predicted.shape == (4, 3, 1) + + +def test_zero_shot_forecaster_univariate_row_mismatch_error(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _WrongRowCountModel: + def predict(self, inputs, prediction_length=1): + return np.zeros((len(inputs) - 1, prediction_length)) + + with pytest.raises( + ValueError, match="forecasts for 3 series while 4 were provided" + ): + ZeroShotForecaster(_WrongRowCountModel()).predict(X, n=2) + + +def test_zero_shot_forecaster_univariate_forecast_too_short_error(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _TooShortModel: + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + return np.zeros((n_rows, prediction_length - 1)) # one step short + + with pytest.raises(ValueError, match="shorter than the requested horizon"): + ZeroShotForecaster(_TooShortModel()).predict(X, n=3) + + +def test_zero_shot_forecaster_accepts_dict_output(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _DictOutputModel: + """Returns a plain dict, as some pipelines do instead of a ModelOutput.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + return {"predictions": np.zeros((n_rows, prediction_length))} + + predicted = ZeroShotForecaster(_DictOutputModel()).predict(X, n=3) + assert predicted.shape == (4, 3, 1) + + +def test_zero_shot_forecaster_empty_output_error(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _EmptyOutputModel: + def predict(self, inputs, prediction_length=1): + return [] + + with pytest.raises(ValueError, match="empty forecast"): + ZeroShotForecaster(_EmptyOutputModel()).predict(X, n=3) + + +def test_zero_shot_forecaster_heterogeneous_output_keeps_first_entry(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _HeterogeneousOutputModel: + """Returns (forecast, extra_diagnostics), as some pipelines do.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + forecast = np.zeros((n_rows, prediction_length)) + diagnostics = np.zeros((n_rows, 2)) # a different shape + return (forecast, diagnostics) + + predicted = ZeroShotForecaster(_HeterogeneousOutputModel()).predict(X, n=3) + assert predicted.shape == (4, 3, 1) + + +def test_zero_shot_forecaster_horizon_axis_prefers_smallest_longer_axis(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _PaddedForecastModel: + """Pads the forecast length to a fixed block size larger than asked, + with no axis exactly matching the horizon.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + n_samples = 3 # sample paths, always shorter than the horizon here + padded_length = prediction_length + 5 + return np.zeros((n_rows, n_samples, padded_length)) + + predicted = ZeroShotForecaster(_PaddedForecastModel()).predict(X, n=4) + np.testing.assert_allclose(predicted, np.zeros((4, 4, 1))) + + +def test_zero_shot_forecaster_horizon_axis_cannot_be_identified(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _TooShortEverywhereModel: + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + return np.zeros((n_rows, 2, 2)) + + with pytest.raises( + ValueError, match="Could not identify the forecast horizon axis" + ): + ZeroShotForecaster(_TooShortEverywhereModel()).predict(X, n=4) + + +def test_zero_shot_forecaster_natively_multivariate_channels_last(): + X = _dataset(n_ts=4, sz=32, d=3) + + class _NativelyMultivariateModel: + """Returns per-channel forecasts directly, given a (n_ts, sz, d) input.""" + + def predict(self, inputs, prediction_length=1): + inputs = np.asarray(inputs) + last = inputs[:, -1, :] # (n_ts, d) + return np.repeat(last[:, None, :], prediction_length, axis=1) + + predicted = ZeroShotForecaster( + _NativelyMultivariateModel(), input_layout="channels_last" + ).predict(X, n=4) + assert predicted.shape == (4, 4, 3) + np.testing.assert_allclose(predicted, np.repeat(X[:, -1:], 4, axis=1)) + + +def test_zero_shot_forecaster_natively_multivariate_channels_first(): + X = _dataset(n_ts=4, sz=32, d=3) + + class _ChannelsFirstModel: + """Expects a (n_ts, d, sz) input and returns (n_ts, d, horizon).""" + + def predict(self, inputs, prediction_length=1): + inputs = np.asarray(inputs) + last = inputs[:, :, -1] # (n_ts, d) + return np.repeat(last[:, :, None], prediction_length, axis=2) + + predicted = ZeroShotForecaster( + _ChannelsFirstModel(), input_layout="channels_first" + ).predict(X, n=4) + assert predicted.shape == (4, 4, 3) + np.testing.assert_allclose(predicted, np.repeat(X[:, -1:], 4, axis=1)) + + +def test_zero_shot_forecaster_natively_multivariate_errors(): + X = _dataset(n_ts=4, sz=32, d=3) + + class _WrongBatchModel: + def predict(self, inputs, prediction_length=1): + return np.zeros((len(inputs) - 1, prediction_length, 3)) + + with pytest.raises( + ValueError, match="forecasts for 3 series while 4 were provided" + ): + ZeroShotForecaster( + _WrongBatchModel(), input_layout="channels_last" + ).predict(X, n=2) + + class _WrongChannelsModel: + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + return np.zeros((n_rows, prediction_length, 5)) # wrong d + + with pytest.raises(ValueError, match="Expected forecasts of shape"): + ZeroShotForecaster( + _WrongChannelsModel(), input_layout="channels_last" + ).predict(X, n=2) + + # --------------------------------------------------------------------------- # LinearProbeForecaster # --------------------------------------------------------------------------- @@ -768,6 +1217,15 @@ def test_linear_probe_classifier_accepts_any_probe(): assert model.decision_function(X).shape == (len(y),) +def test_linear_probe_classifier_no_decision_function(): + X, y = _classification_dataset() + model = LinearProbeClassifier(_DummyBackbone(), probe=GaussianNB()).fit(X, y) + assert isinstance(model.probe_, GaussianNB) + with pytest.raises(AttributeError, match="decision_function"): + model.decision_function(X) + assert model.predict_proba(X).shape == (len(y), 2) + + def test_linear_probe_classifier_layer_and_pooling(): X, y = _classification_dataset() for layer in (0, 1, None): diff --git a/tslearn/foundation/_forecasting.py b/tslearn/foundation/_forecasting.py index a4ed1ac8..5342cfa9 100644 --- a/tslearn/foundation/_forecasting.py +++ b/tslearn/foundation/_forecasting.py @@ -610,8 +610,15 @@ class LinearProbeForecaster(_BaseFoundationForecaster): >>> from chronos import Chronos2Pipeline # doctest: +SKIP >>> pipeline = Chronos2Pipeline.from_pretrained("amazon/chronos-2") # doctest: +SKIP >>> model = LinearProbeForecaster(pipeline.model, horizon=24) # doctest: +SKIP - >>> model.fit(X_train).predict(X_test, n=24).shape # doctest: +SKIP + >>> model.fit(X_train) # doctest: +SKIP + >>> model.predict(X_test).shape # doctest: +SKIP (10, 24, 1) + >>> model.predict(X_test, n=12).shape # doctest: +SKIP + (10, 12, 1) + >>> model.predict(X_test, n=36) # doctest: +SKIP + Traceback (most recent call last): + ... + ValueError: This estimator was fitted for a horizon of 24, so `n` must lie in [1, 24], got 36. Refit with a larger `horizon` to forecast further ahead. References ---------- From b80b7708d026cd7b8b5ed2f35b974c894842d432 Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Thu, 13 Aug 2026 10:14:07 +0200 Subject: [PATCH 05/18] fix reference --- tslearn/foundation/_embedding.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tslearn/foundation/_embedding.py b/tslearn/foundation/_embedding.py index f16d7833..5c37e442 100644 --- a/tslearn/foundation/_embedding.py +++ b/tslearn/foundation/_embedding.py @@ -306,8 +306,8 @@ class TimeSeriesFoundationEmbedder(TimeSeriesMixin, TransformerMixin, BaseEstima References ---------- - .. [1] A. Kolesnikov, L. Beyer, X. Zhai, et al. Big Transfer (BiT): General - Visual Representation Learning. ECCV, 2020. + .. [1] G. Alain and Y. Bengio. Understanding intermediate layers using + linear classifier probes. ICLR Workshop, 2017. """ From abf1a03b9d8928c1b835e9928bd16452c6a29247 Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Thu, 13 Aug 2026 12:26:33 +0200 Subject: [PATCH 06/18] better docs --- docs/_static/img/foundation_tokens.svg | 54 +++++++++++ .../plot_foundation_forecasting.py | 29 +++--- tests/test_foundation.py | 16 ---- tslearn/foundation/_classification.py | 71 +++++++------- tslearn/foundation/_embedding.py | 39 ++++---- tslearn/foundation/_forecasting.py | 95 +++++++------------ 6 files changed, 154 insertions(+), 150 deletions(-) create mode 100644 docs/_static/img/foundation_tokens.svg diff --git a/docs/_static/img/foundation_tokens.svg b/docs/_static/img/foundation_tokens.svg new file mode 100644 index 00000000..cac9a3fd --- /dev/null +++ b/docs/_static/img/foundation_tokens.svg @@ -0,0 +1,54 @@ + + + Token layout of a patch-based foundation model, such as Chronos-2 + + + + 0 + 1 + N-1 + -2 + -1 + + + + + + + + + … + + + Token 1 + Token 2 + Token N + + + + + + + Register + token + Forecast + token + + + + + + tokens=(0, -2) + kept, and averaged over + + + + + + excluded + do not represent the series + + diff --git a/docs/examples/forecasting/plot_foundation_forecasting.py b/docs/examples/forecasting/plot_foundation_forecasting.py index a9c6e887..de82ffb0 100644 --- a/docs/examples/forecasting/plot_foundation_forecasting.py +++ b/docs/examples/forecasting/plot_foundation_forecasting.py @@ -35,8 +35,6 @@ # ---- # # We use a set of sine waves with varying frequencies, phases and noise levels. -# Using synthetic data keeps the example self-contained, and lets us check that -# the models pick up a periodic structure that a short-context model would miss. import numpy as np @@ -71,9 +69,9 @@ # tslearn estimator, whose ``predict`` method returns an array of shape # ``(n_ts, n, d)`` like every other forecaster of the library. # -# Calling ``fit`` is not needed here. Doing so is allowed, so that the -# estimator can be dropped into scikit-learn tooling that expects it, but it -# does not learn anything and emits a warning to that effect. +# Calling ``fit`` is not needed here since the purpose of a +# :class:`~tslearn.foundation.ZeroShotForecaster` is to reuse a pre-trained +# model without any training. from chronos import Chronos2Pipeline @@ -88,10 +86,9 @@ ############################################################################## # ``horizon_axis`` says which axis of the model's output holds the forecast # horizon. There is no shared convention across implementations, Chronos-2 -# alone returning ``(n_series, n_quantiles, horizon)`` from ``predict`` and -# ``(n_series, horizon, n_quantiles)`` from ``predict_quantiles``, so stating -# it is the reliable option. Left to its default of ``"auto"``, the estimator -# infers it from the returned shape and warns when the shape is ambiguous. +# returns ``(n_series, n_quantiles, horizon)`` from ``predict`` so ``-1`` is +# used. Left to its default of ``"auto"``, the estimator would +# infer it from the returned shape. # # The remaining axes are understood as holding quantile levels or sample paths # and are reduced to a point forecast according to the ``quantile`` parameter, @@ -111,8 +108,8 @@ # matters because every window has to go through the pre-trained model once. # # The probing estimator needs the model itself rather than the inference -# pipeline, since it reads hidden states rather than forecasts. Two options -# drive which representations are read: +# pipeline, since it reads hidden states rather than forecasts. Three options +# drive which representations are used: # # * ``layer`` selects the block to probe: an integer places a forward hook on # the corresponding block, so ``layer=-2`` probes the penultimate one. The @@ -123,12 +120,18 @@ # only, so an explicit layer is required here. # * ``pooling`` says how the token representations are aggregated into a single # vector. Chronos-2 emits one token per patch of the context, plus a register -# token, so both averaging (``pooling="mean"``) and picking that single token -# (``pooling="cls"``) are sensible. +# token and a forecast token, so both averaging (``pooling="mean"``) and +# picking that forecast token (``pooling="cls"``) are sensible. # * ``tokens`` restricts which tokens take part in that aggregation. Chronos-2 # appends a register token and a forecast token after its context tokens, and # neither represents the input series, so ``tokens=(0, -2)`` keeps the average # clean. +# +# .. image:: /_static/img/foundation_tokens.svg +# :width: 700 +# :align: center +# :alt: Chronos-2 appends a register token and a forecast token after its +# context tokens; tokens=(0, -2) keeps only the context tokens. from tslearn.foundation import LinearProbeForecaster diff --git a/tests/test_foundation.py b/tests/test_foundation.py index 6f638c66..ec0fd8a6 100644 --- a/tests/test_foundation.py +++ b/tests/test_foundation.py @@ -674,15 +674,6 @@ def test_zero_shot_forecaster(): ) -def test_zero_shot_forecaster_no_warning_when_silenced(): - X = _dataset(n_ts=3, sz=32, d=1) - import warnings - - with warnings.catch_warnings(): - warnings.simplefilter("error") - ZeroShotForecaster(_DummyPipeline(), warn_on_fit=False).fit(X) - - def test_zero_shot_forecaster_multivariate(): X = _dataset(n_ts=5, sz=32, d=3) predicted = ZeroShotForecaster(_DummyPipeline()).predict(X, n=4) @@ -882,13 +873,6 @@ def test_zero_shot_forecaster_errors(): ZeroShotForecaster(_DummyPipeline(), input_layout="bogus").predict(X) -def test_zero_shot_forecaster_feature_mismatch_after_fit(): - X = _dataset(n_ts=4, sz=32, d=1) - model = ZeroShotForecaster(_DummyPipeline(), warn_on_fit=False).fit(X) - with pytest.raises(ValueError, match="features"): - model.predict(_dataset(n_ts=4, sz=32, d=2), n=2) - - def test_zero_shot_forecaster_skips_methods_without_a_horizon_argument(): X = _dataset(n_ts=4, sz=32, d=1) calls = [] diff --git a/tslearn/foundation/_classification.py b/tslearn/foundation/_classification.py index 518bb996..c4c75bfb 100644 --- a/tslearn/foundation/_classification.py +++ b/tslearn/foundation/_classification.py @@ -19,18 +19,14 @@ class LinearProbeClassifier(TimeSeriesMixin, ClassifierMixin, BaseEstimator): - """Classify time series with a linear head on a frozen pre-trained model. + """Classify time series with a head on a frozen pre-trained model. - The pre-trained model is used as a frozen feature extractor and a linear - classifier is fitted on the resulting representations. This is the standard - protocol used to assess how much class information a pre-trained model has - learnt [1]_: because the head is linear and the backbone is frozen, the - accuracy it reaches measures how linearly separable the classes already are - in the representation space. - - It is also a practical classifier in its own right, as it needs no - backpropagation through the pre-trained model and therefore trains in - seconds even on models counting hundreds of millions of parameters. + The pre-trained model is used as a frozen feature extractor and a + classifier is fitted on the resulting representations. + + Linear probing keeps the pre-trained model entirely frozen and only fits a + map from its representations to the class predictions. Compared to fine-tuning, + it leaves the pre-trained weights untouched, which makes it much cheaper. Parameters ---------- @@ -39,45 +35,50 @@ class LinearProbeClassifier(TimeSeriesMixin, ClassifierMixin, BaseEstimator): probe : sklearn estimator or None (default: None) The head fitted on top of the frozen representations. When None, a :class:`sklearn.linear_model.LogisticRegression` is used. Any - scikit-learn classifier can be passed instead; note that using a - non-linear one makes the resulting accuracy an estimate of predictive - performance rather than of linear separability. Frozen representations - sometimes have a very small variance, in which case prepending a - :class:`sklearn.preprocessing.StandardScaler` to the head helps. + scikit-learn classifier can be passed instead. + Frozen representations sometimes have a very small + variance, in which case prepending a + :class:`sklearn.preprocessing.StandardScaler` to the head might help. context_length : int or None (default: None) When set, only the last ``context_length`` timestamps of each series are fed to the model. layer : int or None (default: None) - Layer to probe. When None, the model's output hidden state is used. - When an integer, a forward hook is placed on the corresponding block - of the model's layer stack. Intermediate layers often carry more - class information than the last ones, which tend to specialize - towards the pre-training objective, so this is worth tuning. + Layer to probe, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. layers_path : str or None (default: None) - Dotted path to the stack of layers, e.g. ``"encoder.block"``. When - None, the stack is auto-detected. + Dotted path to the stack of layers, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. pooling : {"mean", "max", "cls", "last", "flatten"} (default: "mean") - How token representations are aggregated into one vector per series. - A probe needs a flat feature matrix, so ``pooling=None`` is not - accepted here. + How token representations are aggregated, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` and + note that ``pooling=None`` is not accepted here. + Default is mean pooling, but if the pre-trained model outputs a CLS + token, it should be better to set ``cls_index`` and use ``"cls"`` + as a pooling strategy. tokens : slice, (start, stop) pair or None (default: None) Which tokens to keep before pooling, see :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. Restricting - the average to the tokens that actually represent the series, e.g. + the average to the tokens that actually represent the context, e.g. ``tokens=(0, -2)`` for Chronos-2, avoids diluting it with class or forecast tokens. cls_index : int (default: 0) Index of the token selected when ``pooling="cls"``. input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") - How multivariate series are handed over to the model. + How the context is laid out when handed over to the model. + ``"univariate"`` forecasts every channel of every series + independently. The other two layouts feed + a ``(n_ts, sz, d)`` or ``(n_ts, d, sz)`` array respectively, for + models that are natively multivariate. input_name : str or None (default: None) Name of the ``forward`` argument receiving the context values. model_kwargs : dict or None (default: None) Extra keyword arguments passed to every ``forward`` call. batch_size : int (default: 32) - Number of series embedded at once. + Number of series embedded at once, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. device : str or None (default: None) - Device on which inference is run. + Device on which inference is run, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. verbose : int (default: 0) When positive, prints progress information. @@ -86,7 +87,7 @@ class information than the last ones, which tend to specialize embedder_ : TimeSeriesFoundationEmbedder The frozen feature extractor. probe_ : sklearn estimator - The fitted linear head. + The fitted head. classes_ : array of shape=(n_classes,) Class labels known to the classifier. n_features_in_ : int @@ -104,12 +105,6 @@ class information than the last ones, which tend to specialize >>> model = LinearProbeClassifier(pipeline.model, layer=-2) # doctest: +SKIP >>> model.fit(X_train, y_train).score(X_test, y_test) # doctest: +SKIP 0.93 - - References - ---------- - .. [1] G. Alain and Y. Bengio. Understanding intermediate layers using - linear classifier probes. ICLR Workshop, 2017. - """ def __init__( @@ -177,7 +172,7 @@ def _check_input(self, X): return to_time_series_dataset(X) def fit(self, X, y): - """Fit a linear classifier on top of the frozen pre-trained model. + """Fit a classifier on top of the frozen pre-trained model. Parameters ---------- diff --git a/tslearn/foundation/_embedding.py b/tslearn/foundation/_embedding.py index 5c37e442..a4454506 100644 --- a/tslearn/foundation/_embedding.py +++ b/tslearn/foundation/_embedding.py @@ -194,8 +194,9 @@ class TimeSeriesFoundationEmbedder(TimeSeriesMixin, TransformerMixin, BaseEstima * ``model`` is a :class:`torch.nn.Module`; * its ``forward`` method accepts the raw context values through one argument (auto-detected among common names, see ``input_name``); - * it returns, or internally computes, hidden states of shape - ``(batch, n_tokens, dim)``. + * it returns hidden states of shape + ``(batch, n_tokens, dim)`` if ``pooling=None`` and ``(batch, dim)`` + otherwise. Parameters ---------- @@ -226,7 +227,7 @@ class TimeSeriesFoundationEmbedder(TimeSeriesMixin, TransformerMixin, BaseEstima ``None``, for which the string ``"none"`` is accepted as an alias, applies no pooling at all: :meth:`transform` then returns a time - series dataset of shape ``(n_ts, n_tokens, d * dim)`` rather than a + series dataset of shape ``(n_ts, n_tokens, dim)`` rather than a flat feature matrix, which turns this estimator into a time series to time series transform. Combined with ``tokens``, this is a way to obtain one representation per timestep, or per patch of timesteps, @@ -238,37 +239,29 @@ class TimeSeriesFoundationEmbedder(TimeSeriesMixin, TransformerMixin, BaseEstima not represent the input series, such as class, register or forecast tokens; excluding them is usually what one wants, both to build clean per-timestep representations and to avoid polluting an average. - Chronos-2, for instance, emits its context tokens first, so - ``tokens=(0, -2)`` drops its trailing register and forecast tokens. When None, all tokens are kept. This parameter does not affect - ``pooling="cls"``, whose whole purpose is to reach a token that - ``tokens`` would typically exclude. + ``pooling="cls"``. cls_index : int (default: 0) - Index of the token to select when ``pooling="cls"``, applied to the - model's full token sequence, before any ``tokens`` selection. Note - that not all models place their class token first: Chronos-2, for - instance, inserts a register token between the context and forecast - tokens. + Index of the token to select when ``pooling="cls"``. input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") - How multivariate series are handed over to the model. - ``"univariate"`` embeds every channel independently, as a - ``(n_ts * d, sz)`` array, and concatenates the resulting per-channel - embeddings, which is what most time series foundation models expect. - The other two layouts feed a ``(n_ts, sz, d)`` or ``(n_ts, d, sz)`` - array respectively, for models that are natively multivariate. + How the context is laid out when handed over to the model. + ``"univariate"`` forecasts every channel of every series + independently. The other two layouts feed + a ``(n_ts, sz, d)`` or ``(n_ts, d, sz)`` array respectively, for + models that are natively multivariate. context_length : int or None (default: None) When set, only the last ``context_length`` timestamps of each series are fed to the model. input_name : str or None (default: None) - Name of the ``forward`` argument receiving the context values. When - None, it is auto-detected among common names. + Name of the ``forward`` argument receiving the context values. model_kwargs : dict or None (default: None) Extra keyword arguments passed to every ``forward`` call. batch_size : int (default: 32) - Number of series embedded at once. + Number of series embedded at once, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. device : str or None (default: None) - Device on which inference is run, e.g. ``"cuda"``. When None, the - device the model already lives on is used. + Device on which inference is run, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. verbose : int (default: 0) When positive, prints progress information. diff --git a/tslearn/foundation/_forecasting.py b/tslearn/foundation/_forecasting.py index 5342cfa9..2697f7bb 100644 --- a/tslearn/foundation/_forecasting.py +++ b/tslearn/foundation/_forecasting.py @@ -269,8 +269,7 @@ class ZeroShotForecaster(_BaseFoundationForecaster): input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") How the context is laid out when handed over to the model. ``"univariate"`` forecasts every channel of every series - independently and feeds a ``(n_ts * d, sz)`` array, which is what - most time series foundation models expect. The other two layouts feed + independently. The other two layouts feed a ``(n_ts, sz, d)`` or ``(n_ts, d, sz)`` array respectively, for models that are natively multivariate. context_length : int or None (default: None) @@ -278,11 +277,8 @@ class ZeroShotForecaster(_BaseFoundationForecaster): are used as context. horizon_axis : int or "auto" (default: "auto") Which axis of the array returned by the model holds the forecast - horizon, axis 0 being the series. There is no shared convention here: - Chronos-2 returns ``(n_series, n_quantiles, horizon)`` while its own - ``predict_quantiles`` returns ``(n_series, horizon, n_quantiles)``, - so setting this explicitly, e.g. ``horizon_axis=-1``, is the reliable - option. When ``"auto"``, the axis is inferred from the returned + horizon, axis 0 being the series. + When ``"auto"``, the axis is inferred from the returned shape, preferring an axis whose size is exactly the horizon and, among those, the last one. A warning is raised if the shape leaves the choice ambiguous, which happens when the model returns as many @@ -293,26 +289,16 @@ class ZeroShotForecaster(_BaseFoundationForecaster): Probabilistic models return several quantiles or sample paths; this is the quantile used to derive a point forecast from them. When None, the mean is used instead. Note that the quantile is taken over the - values the model returned, without assuming which level each of them - corresponds to, since that information is not exposed in any - standard way. For a model returning evenly spread quantile levels, - the default therefore recovers the median forecast, up to any - quantile crossing. + values the model returned. model_kwargs : dict or None (default: None) Extra keyword arguments passed to the model's forecasting method. - warn_on_fit : bool (default: True) - Whether calling :meth:`fit` should emit a warning reminding that - nothing is being learnt. Attributes ---------- - n_features_in_ : int - Number of features (channels) of the series seen during fit. context_format_ : str How the batch of contexts ended up being presented to the model, among ``"2d"``, ``"3d"`` and ``"list"``. Set after the first call to - :meth:`predict` with the ``"univariate"`` layout, as implementations - disagree on this and the accepted one is discovered by trial. + :meth:`predict` with the ``"univariate"`` layout. Notes ----- @@ -331,11 +317,6 @@ class ZeroShotForecaster(_BaseFoundationForecaster): >>> model.predict(X, n=24).shape # doctest: +SKIP (10, 24, 1) - References - ---------- - .. [1] A. F. Ansari, L. Stella, C. Turkmen, et al. Chronos: Learning the - Language of Time Series. Transactions on Machine Learning Research, 2024. - """ def __init__( @@ -347,7 +328,6 @@ def __init__( horizon_axis="auto", quantile=0.5, model_kwargs=None, - warn_on_fit=True, ): self.model = model self.predict_fn = predict_fn @@ -356,7 +336,6 @@ def __init__( self.horizon_axis = horizon_axis self.quantile = quantile self.model_kwargs = model_kwargs - self.warn_on_fit = warn_on_fit def _resolve_predict_fn(self): """Build a ``(context, horizon) -> forecast`` callable from the model.""" @@ -444,17 +423,12 @@ def fit(self, X, y=None): """ self._check_layout() X = self._check_input(X) - if self.warn_on_fit: - warnings.warn( - "ZeroShotForecaster.fit does not train anything: the wrapped " - "model is used as-is for zero-shot forecasting. Use " - "LinearProbeForecaster to fit a head on top of it, or pass " - "warn_on_fit=False to silence this warning.", - UserWarning, - stacklevel=2, - ) - self.n_features_in_ = X.shape[2] - self._is_fitted_ = True + warnings.warn( + "ZeroShotForecaster.fit does not train anything: the wrapped " + "model is used as-is for zero-shot forecasting.", + UserWarning, + stacklevel=2, + ) return self def predict(self, X, n=1): @@ -482,11 +456,6 @@ def predict(self, X, n=1): if n < 1: raise ValueError(f"`n` must be a positive integer, got {n}.") X = self._check_input(X) - if getattr(self, "n_features_in_", X.shape[2]) != X.shape[2]: - raise ValueError( - f"Series with {self.n_features_in_} features were expected, got " - f"{X.shape[2]}." - ) if self.context_length is not None: X = X[:, -self.context_length :] @@ -523,13 +492,13 @@ def predict(self, X, n=1): class LinearProbeForecaster(_BaseFoundationForecaster): - """Forecast by fitting a linear head on a frozen pre-trained model. + """Forecast by fitting a head on a frozen pre-trained model. Linear probing keeps the pre-trained model entirely frozen and only fits a linear map from its representations to the quantity of interest. Compared to zero-shot use, it adapts the model to the data at hand at a very small computational cost, and compared to fine-tuning, it leaves the pre-trained - weights untouched, which makes it both cheap and hard to overfit. + weights untouched, which makes it much cheaper. Training pairs are cut out of the provided series with a sliding window: the ``context_length`` timestamps preceding a cut point are embedded, and @@ -543,17 +512,15 @@ class LinearProbeForecaster(_BaseFoundationForecaster): The head fitted on top of the frozen representations. When None, a :class:`sklearn.linear_model.RidgeCV` is used, which selects its regularization strength by cross-validation and admits a closed-form - solution. Any multi-output capable regressor can be passed instead; - note that using a non-linear one makes the name "probe" a misnomer, - and the resulting scores no longer measure linear separability of the - representations. Frozen representations sometimes have a very small + solution. Any multi-output capable regressor can be passed instead. + Frozen representations sometimes have a very small variance, in which case prepending a - :class:`sklearn.preprocessing.StandardScaler` to the head helps. + :class:`sklearn.preprocessing.StandardScaler` to the head might help. context_length : int (default: 128) Number of timestamps fed to the pre-trained model. horizon : int (default: 1) - Number of timestamps forecasted at once. This is fixed at fit time, - as one linear head is fitted for the whole horizon. + Number of timestamps forecasted at once. This cannot be changed after + fit time, as one linear head is fitted for the whole horizon. stride : int (default: 1) Step between two consecutive training windows. Larger values yield fewer, less redundant training pairs and a faster fit. @@ -565,9 +532,8 @@ class LinearProbeForecaster(_BaseFoundationForecaster): :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. pooling : {"mean", "max", "cls", "last", "flatten"} (default: "mean") How token representations are aggregated, see - :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. Unlike - that class, a probe needs a flat feature matrix, so ``pooling=None`` - is not accepted here. + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` and + note that ``pooling=None`` is not accepted here. tokens : slice, (start, stop) pair or None (default: None) Which tokens to keep before pooling, see :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. Restricting @@ -578,14 +544,20 @@ class LinearProbeForecaster(_BaseFoundationForecaster): Index of the token selected when ``pooling="cls"``. input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") How the context is laid out when handed over to the model. + ``"univariate"`` forecasts every channel of every series + independently. The other two layouts feed + a ``(n_ts, sz, d)`` or ``(n_ts, d, sz)`` array respectively, for + models that are natively multivariate. input_name : str or None (default: None) Name of the ``forward`` argument receiving the context values. model_kwargs : dict or None (default: None) Extra keyword arguments passed to every ``forward`` call. batch_size : int (default: 32) - Number of series embedded at once. + Number of series embedded at once, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. device : str or None (default: None) - Device on which inference is run. + Device on which inference is run, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. verbose : int (default: 0) When positive, prints progress information. @@ -594,7 +566,7 @@ class LinearProbeForecaster(_BaseFoundationForecaster): embedder_ : TimeSeriesFoundationEmbedder The frozen feature extractor. probe_ : sklearn estimator - The fitted linear head. + The fitted head. n_features_in_ : int Number of features (channels) of the series seen during fit. n_windows_ : int @@ -611,6 +583,7 @@ class LinearProbeForecaster(_BaseFoundationForecaster): >>> pipeline = Chronos2Pipeline.from_pretrained("amazon/chronos-2") # doctest: +SKIP >>> model = LinearProbeForecaster(pipeline.model, horizon=24) # doctest: +SKIP >>> model.fit(X_train) # doctest: +SKIP + ... >>> model.predict(X_test).shape # doctest: +SKIP (10, 24, 1) >>> model.predict(X_test, n=12).shape # doctest: +SKIP @@ -618,7 +591,8 @@ class LinearProbeForecaster(_BaseFoundationForecaster): >>> model.predict(X_test, n=36) # doctest: +SKIP Traceback (most recent call last): ... - ValueError: This estimator was fitted for a horizon of 24, so `n` must lie in [1, 24], got 36. Refit with a larger `horizon` to forecast further ahead. + ValueError: This estimator was fitted for a horizon of 24, so `n` must lie in [1, 24], got 36. + Refit with a larger `horizon` to forecast further ahead. References ---------- @@ -721,7 +695,7 @@ def _validate_params_(self): ) def fit(self, X, y=None): - """Fit a linear head on top of the frozen pre-trained model. + """Fit the probe on top of the frozen pre-trained model. Parameters ---------- @@ -751,7 +725,8 @@ def fit(self, X, y=None): return self def predict(self, X, n=None): - """Forecast ``horizon`` timestamps ahead of the given series. + """Forecast ``horizon`` (or ``n``, if set) timestamps ahead of the + given series. Parameters ---------- From c6801ce113010fcb641786cb17b71b7708cd7431 Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Thu, 13 Aug 2026 12:39:16 +0200 Subject: [PATCH 07/18] generic token pooling --- .../plot_foundation_linear_probe.py | 4 +-- .../plot_foundation_forecasting.py | 18 ++++++---- tests/test_foundation.py | 22 ++++++------ tslearn/foundation/_classification.py | 19 +++++----- tslearn/foundation/_embedding.py | 36 ++++++++++--------- tslearn/foundation/_forecasting.py | 13 +++---- 6 files changed, 60 insertions(+), 52 deletions(-) diff --git a/docs/examples/classification/plot_foundation_linear_probe.py b/docs/examples/classification/plot_foundation_linear_probe.py index 3c732832..f2e70d33 100644 --- a/docs/examples/classification/plot_foundation_linear_probe.py +++ b/docs/examples/classification/plot_foundation_linear_probe.py @@ -98,7 +98,7 @@ # token plays in a text encoder. n_layers = len(pipeline.model.encoder.block) -poolings = ["mean", "max", "cls"] +poolings = ["mean", "max", "token"] # Every other layer is enough to see the trend, and halves the build time layers = sorted({*range(0, n_layers, 2), n_layers - 1}) @@ -110,7 +110,7 @@ layer=layer, pooling=pooling, # Chronos-2 places its register token after the context tokens - cls_index=-2 if pooling == "cls" else 0, + token_index=-2 if pooling == "token" else 0, tokens=(0, -2), ).fit(X_train, y_train) accuracies[layer, pooling] = model.score(X_test, y_test) diff --git a/docs/examples/forecasting/plot_foundation_forecasting.py b/docs/examples/forecasting/plot_foundation_forecasting.py index de82ffb0..fcb6648d 100644 --- a/docs/examples/forecasting/plot_foundation_forecasting.py +++ b/docs/examples/forecasting/plot_foundation_forecasting.py @@ -120,12 +120,16 @@ # only, so an explicit layer is required here. # * ``pooling`` says how the token representations are aggregated into a single # vector. Chronos-2 emits one token per patch of the context, plus a register -# token and a forecast token, so both averaging (``pooling="mean"``) and -# picking that forecast token (``pooling="cls"``) are sensible. +# token and a forecast token, so both averaging (``pooling="mean"``) and +# picking a single token (``pooling="token"``) are sensible. +# Chronos-2's forecast token is a natural choice +# here: ``token_index=-1`` reads it directly, rather than averaging over +# representations that were never meant to summarize the series for +# forecasting. # * ``tokens`` restricts which tokens take part in that aggregation. Chronos-2 # appends a register token and a forecast token after its context tokens, and -# neither represents the input series, so ``tokens=(0, -2)`` keeps the average -# clean. +# neither represents the input series, so ``tokens=(0, -2)`` focues on the +# slice 0:-2 and keeps the average clean. # # .. image:: /_static/img/foundation_tokens.svg # :width: 700 @@ -223,7 +227,7 @@ results = {} for layer in [-1, -2, -4]: - for pooling in ["mean", "cls"]: + for pooling in ["mean", "token"]: model = LinearProbeForecaster( pipeline.model, context_length=context_length, @@ -231,8 +235,8 @@ stride=48, layer=layer, pooling=pooling, - # Chronos-2 appends its register token after the context tokens - cls_index=-2, + # Chronos-2's forecast token is its last one + token_index=-1, tokens=(0, -2), ).fit(X_train) results[(layer, pooling)] = mae(X_test, model.predict(X_train)) diff --git a/tests/test_foundation.py b/tests/test_foundation.py index ec0fd8a6..2be9cee6 100644 --- a/tests/test_foundation.py +++ b/tests/test_foundation.py @@ -65,7 +65,7 @@ class _DummyBackbone(torch.nn.Module): """A patch-based encoder taking univariate series, like Chronos-2. A register token is appended after the context tokens, so that the - ``pooling="cls"`` code path can be exercised on a non-zero ``cls_index``. + ``pooling="token"`` code path can be exercised on a non-zero ``token_index``. """ def __init__(self, n_layers=3, d_model=D_MODEL, patch_size=PATCH_SIZE, seed=0): @@ -166,7 +166,7 @@ def test_embedder_channel_stacking_is_order_preserving(): ) -@pytest.mark.parametrize("pooling", ["mean", "max", "cls", "last", "flatten"]) +@pytest.mark.parametrize("pooling", ["mean", "max", "token", "last", "flatten"]) def test_embedder_poolings(pooling): X = _dataset(n_ts=4, sz=32, d=1) embedder = TimeSeriesFoundationEmbedder(_DummyBackbone(), pooling=pooling) @@ -235,13 +235,13 @@ def test_embedder_token_selection(): ) assert not np.allclose(without_register, over_everything) - # ... but not to "cls", which is meant to reach an excluded token + # ... but not to "token", which is meant to reach an excluded token np.testing.assert_allclose( TimeSeriesFoundationEmbedder( - backbone, pooling="cls", cls_index=-1, tokens=(0, -1) + backbone, pooling="token", token_index=-1, tokens=(0, -1) ).fit_transform(X), TimeSeriesFoundationEmbedder( - backbone, pooling="cls", cls_index=-1 + backbone, pooling="token", token_index=-1 ).fit_transform(X), ) @@ -305,12 +305,12 @@ def test_probes_accept_token_selection(): assert model.predict(X_fc).shape == (6, 2, 1) -def test_embedder_cls_index_selects_the_register_token(): +def test_embedder_token_index_selects_the_register_token(): X = _dataset(n_ts=3, sz=32, d=1) # The register token is appended last and does not depend on the input, so # selecting it must yield identical embeddings for all series. embeddings = TimeSeriesFoundationEmbedder( - _DummyBackbone(), pooling="cls", cls_index=-1, layer=0 + _DummyBackbone(), pooling="token", token_index=-1, layer=0 ).fit_transform(X) np.testing.assert_allclose( embeddings, np.repeat(embeddings[:1], 3, axis=0), rtol=1e-5, atol=1e-6 @@ -318,7 +318,7 @@ def test_embedder_cls_index_selects_the_register_token(): # Whereas a context token does depend on the input embeddings = TimeSeriesFoundationEmbedder( - _DummyBackbone(), pooling="cls", cls_index=0, layer=0 + _DummyBackbone(), pooling="token", token_index=0, layer=0 ).fit_transform(X) assert not np.allclose(embeddings, embeddings[:1]) @@ -390,9 +390,9 @@ def test_embedder_errors(): TimeSeriesFoundationEmbedder(_DummyPipeline()).fit(X) with pytest.raises(ValueError, match="out of range"): TimeSeriesFoundationEmbedder(_DummyBackbone(n_layers=3), layer=7).fit(X) - with pytest.raises(ValueError, match="cls_index"): + with pytest.raises(ValueError, match="token_index"): TimeSeriesFoundationEmbedder( - _DummyBackbone(), pooling="cls", cls_index=999 + _DummyBackbone(), pooling="token", token_index=999 ).fit(X) with pytest.raises(ValueError, match="features"): TimeSeriesFoundationEmbedder(_DummyBackbone()).fit(X).transform( @@ -1213,7 +1213,7 @@ def test_linear_probe_classifier_no_decision_function(): def test_linear_probe_classifier_layer_and_pooling(): X, y = _classification_dataset() for layer in (0, 1, None): - for pooling in ("mean", "max", "cls"): + for pooling in ("mean", "max", "token"): model = LinearProbeClassifier( _DummyBackbone(), layer=layer, pooling=pooling ).fit(X, y) diff --git a/tslearn/foundation/_classification.py b/tslearn/foundation/_classification.py index c4c75bfb..0774dfab 100644 --- a/tslearn/foundation/_classification.py +++ b/tslearn/foundation/_classification.py @@ -48,21 +48,22 @@ class LinearProbeClassifier(TimeSeriesMixin, ClassifierMixin, BaseEstimator): layers_path : str or None (default: None) Dotted path to the stack of layers, see :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. - pooling : {"mean", "max", "cls", "last", "flatten"} (default: "mean") + pooling : {"mean", "max", "token", "last", "flatten"} (default: "mean") How token representations are aggregated, see :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` and note that ``pooling=None`` is not accepted here. - Default is mean pooling, but if the pre-trained model outputs a CLS - token, it should be better to set ``cls_index`` and use ``"cls"`` - as a pooling strategy. + Default is mean pooling, but if the pre-trained model outputs a + class token, it should be better to set ``token_index`` + to this token's index and use ``"token"`` as a pooling strategy. tokens : slice, (start, stop) pair or None (default: None) Which tokens to keep before pooling, see :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. Restricting the average to the tokens that actually represent the context, e.g. ``tokens=(0, -2)`` for Chronos-2, avoids diluting it with class or forecast tokens. - cls_index : int (default: 0) - Index of the token selected when ``pooling="cls"``. + token_index : int (default: 0) + Index of the token selected when ``pooling="token"``, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") How the context is laid out when handed over to the model. ``"univariate"`` forecasts every channel of every series @@ -116,7 +117,7 @@ def __init__( layers_path=None, pooling="mean", tokens=None, - cls_index=0, + token_index=0, input_layout="univariate", input_name=None, model_kwargs=None, @@ -131,7 +132,7 @@ def __init__( self.layers_path = layers_path self.pooling = pooling self.tokens = tokens - self.cls_index = cls_index + self.token_index = token_index self.input_layout = input_layout self.input_name = input_name self.model_kwargs = model_kwargs @@ -146,7 +147,7 @@ def _make_embedder(self): layers_path=self.layers_path, pooling=self.pooling, tokens=self.tokens, - cls_index=self.cls_index, + token_index=self.token_index, input_layout=self.input_layout, context_length=self.context_length, input_name=self.input_name, diff --git a/tslearn/foundation/_embedding.py b/tslearn/foundation/_embedding.py index a4454506..21059ba8 100644 --- a/tslearn/foundation/_embedding.py +++ b/tslearn/foundation/_embedding.py @@ -40,7 +40,7 @@ #: Supported ways of aggregating token representations. ``None``, for which #: ``"none"`` is accepted as an alias, keeps them all and turns the estimator #: into a time series to time series transform. -POOLINGS = ("mean", "max", "cls", "last", "flatten", "none", None) +POOLINGS = ("mean", "max", "token", "last", "flatten", "none", None) #: Supported layouts for the array handed over to the wrapped model. LAYOUTS = ("univariate", "channels_last", "channels_first") @@ -217,11 +217,13 @@ class TimeSeriesFoundationEmbedder(TimeSeriesMixin, TransformerMixin, BaseEstima ``"encoder.block"``. Only used when ``layer`` is an integer. When None, the stack is auto-detected as the longest list of identical sub-modules. - pooling : {"mean", "max", "cls", "last", "flatten", None} (default: "mean") + pooling : {"mean", "max", "token", "last", "flatten", None} (default: "mean") How to aggregate the ``n_tokens`` representations of a series into a - single vector. ``"cls"`` selects the single token at index - ``cls_index``, which is how models exposing a class (or register) - token are usually probed. ``"last"`` selects the last token and + single vector. ``"token"`` selects the single token at index + ``token_index``, which is how one reaches a token that carries a + meaning of its own, such as a class token for + classification, or a forecast token for forecasting. + ``"last"`` selects the last token and ``"flatten"`` concatenates them all, which yields a much larger, and context-length dependent, feature vector. @@ -240,9 +242,9 @@ class TimeSeriesFoundationEmbedder(TimeSeriesMixin, TransformerMixin, BaseEstima tokens; excluding them is usually what one wants, both to build clean per-timestep representations and to avoid polluting an average. When None, all tokens are kept. This parameter does not affect - ``pooling="cls"``. - cls_index : int (default: 0) - Index of the token to select when ``pooling="cls"``. + ``pooling="token"``. + token_index : int (default: 0) + Index of the token to select when ``pooling="token"``. input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") How the context is laid out when handed over to the model. ``"univariate"`` forecasts every channel of every series @@ -311,7 +313,7 @@ def __init__( layers_path=None, pooling="mean", tokens=None, - cls_index=0, + token_index=0, input_layout="univariate", context_length=None, input_name=None, @@ -325,7 +327,7 @@ def __init__( self.layers_path = layers_path self.pooling = pooling self.tokens = tokens - self.cls_index = cls_index + self.token_index = token_index self.input_layout = input_layout self.context_length = context_length self.input_name = input_name @@ -443,16 +445,16 @@ def _pool(self, hidden_states): ) pooling = _normalize_pooling(self.pooling) - # The class token is deliberately read before any token selection, as - # it is usually one of the tokens `tokens` is meant to filter out. - if pooling == "cls": + # The selected token is deliberately read before any token selection, + # as it is usually one of the tokens `tokens` is meant to filter out. + if pooling == "token": n_tokens = hidden_states.shape[1] - if not -n_tokens <= self.cls_index < n_tokens: + if not -n_tokens <= self.token_index < n_tokens: raise ValueError( - f"`cls_index={self.cls_index}` is out of range: the model " - f"produced {n_tokens} tokens." + f"`token_index={self.token_index}` is out of range: the " + f"model produced {n_tokens} tokens." ) - return hidden_states[:, self.cls_index] + return hidden_states[:, self.token_index] hidden_states = hidden_states[:, _token_slice(self.tokens)] if hidden_states.shape[1] == 0: diff --git a/tslearn/foundation/_forecasting.py b/tslearn/foundation/_forecasting.py index 2697f7bb..669d2621 100644 --- a/tslearn/foundation/_forecasting.py +++ b/tslearn/foundation/_forecasting.py @@ -530,7 +530,7 @@ class LinearProbeForecaster(_BaseFoundationForecaster): layers_path : str or None (default: None) Dotted path to the stack of layers, see :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. - pooling : {"mean", "max", "cls", "last", "flatten"} (default: "mean") + pooling : {"mean", "max", "token", "last", "flatten"} (default: "mean") How token representations are aggregated, see :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` and note that ``pooling=None`` is not accepted here. @@ -540,8 +540,9 @@ class LinearProbeForecaster(_BaseFoundationForecaster): the average to the tokens that actually represent the context, e.g. ``tokens=(0, -2)`` for Chronos-2, avoids diluting it with class or forecast tokens. - cls_index : int (default: 0) - Index of the token selected when ``pooling="cls"``. + token_index : int (default: 0) + Index of the token selected when ``pooling="token"``, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") How the context is laid out when handed over to the model. ``"univariate"`` forecasts every channel of every series @@ -612,7 +613,7 @@ def __init__( layers_path=None, pooling="mean", tokens=None, - cls_index=0, + token_index=0, input_layout="univariate", input_name=None, model_kwargs=None, @@ -629,7 +630,7 @@ def __init__( self.layers_path = layers_path self.pooling = pooling self.tokens = tokens - self.cls_index = cls_index + self.token_index = token_index self.input_layout = input_layout self.input_name = input_name self.model_kwargs = model_kwargs @@ -644,7 +645,7 @@ def _make_embedder(self): layers_path=self.layers_path, pooling=self.pooling, tokens=self.tokens, - cls_index=self.cls_index, + token_index=self.token_index, input_layout=self.input_layout, context_length=None, input_name=self.input_name, From 6dfdf5d7e2e750f019d0955dc3cc86478f844e05 Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Thu, 13 Aug 2026 12:47:51 +0200 Subject: [PATCH 08/18] better doc notebooks --- .../classification/plot_foundation_linear_probe.py | 9 ++++----- docs/examples/forecasting/plot_foundation_forecasting.py | 3 +-- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/docs/examples/classification/plot_foundation_linear_probe.py b/docs/examples/classification/plot_foundation_linear_probe.py index f2e70d33..006a2105 100644 --- a/docs/examples/classification/plot_foundation_linear_probe.py +++ b/docs/examples/classification/plot_foundation_linear_probe.py @@ -5,16 +5,15 @@ Time series foundation models are pre-trained for forecasting, yet the representations they build along the way carry a lot of information about the shape of a series, which makes them useful for other tasks as well. Linear -probing [1]_ is the standard way to measure that: the pre-trained model is kept -frozen and used as a feature extractor, and a plain linear classifier is fitted +probing [1]_ is a cheap way to implement that principle: the pre-trained model +is kept frozen and used as a feature extractor, and a plain classifier is fitted on top of its representations. Because the head is linear and the backbone is never updated, the accuracy reached tells us how linearly separable the classes already are in the representation space. This example applies :class:`~tslearn.foundation.LinearProbeClassifier` to a -UCR dataset, using Chronos-2 [2]_ as the frozen backbone, and compares it to -tslearn's classical baselines. Running it requires the ``chronos-forecasting`` -package:: +UCR dataset, using Chronos-2 [2]_ as the frozen backbone. +Running it requires the ``chronos-forecasting`` package:: pip install "chronos-forecasting>=2.0" diff --git a/docs/examples/forecasting/plot_foundation_forecasting.py b/docs/examples/forecasting/plot_foundation_forecasting.py index fcb6648d..5cbdabf7 100644 --- a/docs/examples/forecasting/plot_foundation_forecasting.py +++ b/docs/examples/forecasting/plot_foundation_forecasting.py @@ -158,8 +158,7 @@ # The head fitted on top of the frozen representations defaults to a # :class:`sklearn.linear_model.RidgeCV`, which picks its regularization # strength by cross-validation. Any scikit-learn regressor can be passed -# instead through the ``probe`` parameter, for instance to control the -# regularization path explicitly. +# instead through the ``probe`` parameter. ############################################################################## # Comparison From 45c0ea729cd879e7f32406a6b540e9c550944b1c Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Thu, 13 Aug 2026 13:00:47 +0200 Subject: [PATCH 09/18] typo --- docs/examples/forecasting/plot_foundation_forecasting.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/examples/forecasting/plot_foundation_forecasting.py b/docs/examples/forecasting/plot_foundation_forecasting.py index 5cbdabf7..9fb24cb6 100644 --- a/docs/examples/forecasting/plot_foundation_forecasting.py +++ b/docs/examples/forecasting/plot_foundation_forecasting.py @@ -128,7 +128,7 @@ # forecasting. # * ``tokens`` restricts which tokens take part in that aggregation. Chronos-2 # appends a register token and a forecast token after its context tokens, and -# neither represents the input series, so ``tokens=(0, -2)`` focues on the +# neither represents the input series, so ``tokens=(0, -2)`` focuses on the # slice 0:-2 and keeps the average clean. # # .. image:: /_static/img/foundation_tokens.svg From 96faf8369c389b18e3a7769ba15f1a948b6217ac Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Thu, 13 Aug 2026 14:01:24 +0200 Subject: [PATCH 10/18] presentation --- docs/examples/forecasting/plot_foundation_forecasting.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/examples/forecasting/plot_foundation_forecasting.py b/docs/examples/forecasting/plot_foundation_forecasting.py index 9fb24cb6..6ede9ee8 100644 --- a/docs/examples/forecasting/plot_foundation_forecasting.py +++ b/docs/examples/forecasting/plot_foundation_forecasting.py @@ -241,4 +241,4 @@ results[(layer, pooling)] = mae(X_test, model.predict(X_train)) for (layer, pooling), score in sorted(results.items(), key=lambda kv: kv[1]): - print(f"layer={str(layer):>4}, pooling={pooling:>4}: MAE = {score:.4f}") + print(f"layer={str(layer):>5}, pooling={pooling:>4}: MAE = {score:.4f}") From 173d1870c0382a5c9fc78c2a1e831cfd6ddd657c Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Thu, 20 Aug 2026 16:17:59 +0200 Subject: [PATCH 11/18] test multiple foundation models and document their usage --- .github/workflows/test_foundation_zoo.yml | 63 +++ .../forecasting/plot_foundation_model_zoo.py | 412 ++++++++++++++++++ tests/test_foundation_zoo.py | 220 ++++++++++ 3 files changed, 695 insertions(+) create mode 100644 .github/workflows/test_foundation_zoo.yml create mode 100644 docs/examples/forecasting/plot_foundation_model_zoo.py create mode 100644 tests/test_foundation_zoo.py diff --git a/.github/workflows/test_foundation_zoo.yml b/.github/workflows/test_foundation_zoo.yml new file mode 100644 index 00000000..f168dcbe --- /dev/null +++ b/.github/workflows/test_foundation_zoo.yml @@ -0,0 +1,63 @@ +name: Test foundation model zoo + +# Exercises tslearn.foundation against real pre-trained checkpoints from the +# Hugging Face Hub (see tests/test_foundation_zoo.py and the +# plot_foundation_model_zoo.py gallery example it backs). Each model pulls in +# its own package, sometimes with conflicting version pins (torch, +# transformers...), so every model runs as its own isolated job rather than +# being part of the regular test suite. + +on: + workflow_dispatch: + schedule: + - cron: "30 3 * * 0" # weekly, downloads are too heavy to run nightly + +jobs: + test_foundation_zoo: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + - name: chronos_bolt + deps: | + python -m pip install "chronos-forecasting>=2.0" + - name: timesfm + deps: | + python -m pip install "timesfm[torch]" + - name: moirai + deps: | + python -m pip install "uni2ts @ git+https://github.com/SalesforceAIResearch/uni2ts.git" + - name: ttm + deps: | + python -m pip install "granite-tsfm[notebooks] @ git+https://github.com/ibm-granite/granite-tsfm.git@v0.2.22" + python -m pip install "transformers==4.44.2" + - name: moment + deps: | + python -m pip install momentfm + - name: time_moe + deps: | + python -m pip install "transformers==4.40.0" + env: + TSLEARN_RUN_FOUNDATION_ZOO: "1" + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Install tslearn + run: | + python -m pip install --upgrade pip + python -m pip install --extra-index-url https://download.pytorch.org/whl/cpu .[tests] torch + shell: bash + + - name: Install ${{ matrix.name }} dependencies + run: ${{ matrix.deps }} + shell: bash + + - name: Test ${{ matrix.name }} + run: python -m pytest -v tests/test_foundation_zoo.py -k "${{ matrix.name }}" + shell: bash diff --git a/docs/examples/forecasting/plot_foundation_model_zoo.py b/docs/examples/forecasting/plot_foundation_model_zoo.py new file mode 100644 index 00000000..d787b233 --- /dev/null +++ b/docs/examples/forecasting/plot_foundation_model_zoo.py @@ -0,0 +1,412 @@ +""" +A tour of the pre-trained forecasting model zoo +================================================= + +:mod:`tslearn.foundation` does not target a specific pre-trained model: it +only assumes that a model exposes a forecasting method, or a PyTorch +``forward``, following one of a handful of widespread conventions (see +:class:`~tslearn.foundation.ZeroShotForecaster` and +:class:`~tslearn.foundation.LinearProbeForecaster`). This example surveys +seven models currently considered state of the art for time series +forecasting, and shows, for each of them, what it takes to plug it in as a +:class:`~tslearn.foundation.ZeroShotForecaster`, a +:class:`~tslearn.foundation.LinearProbeForecaster`, or both. The smallest +checkpoint published for every model is used throughout, to keep downloads +light. + +For a line-by-line walkthrough of the module's parameters, see the +:doc:`Chronos-2 example `, which this one +complements rather than repeats. + +================= ======================= ========= ============ +Model Package Zero-shot Linear probe +================= ======================= ========= ============ +Chronos-2 [1]_ ``chronos-forecasting`` yes yes +Chronos-Bolt [2]_ ``chronos-forecasting`` yes yes +TimesFM [3]_ ``timesfm`` yes -- +Moirai [4]_ ``uni2ts`` yes -- +TTM [5]_ ``granite-tsfm`` yes -- +MOMENT [6]_ ``momentfm`` -- yes +Time-MoE [7]_ ``transformers`` yes yes +================= ======================= ========= ============ + +The right-hand column is not a limitation of tslearn: it reflects how each +model exposes its internals. Models built around a clean +``forward(context) -> hidden_states`` (the two Chronos variants, Time-MoE) +probe as easily as they forecast zero-shot. Models whose low-level module +expects an already-patchified input, an explicit padding mask, or a packed +multi-series format (TimesFM, Moirai, TTM) are naturally used zero-shot, and +would need a bespoke wrapper to expose clean per-token representations for +probing. MOMENT is the mirror case: only its reconstruction head is +pre-trained, so its zero-shot *forecasts* would come out of a head that has +never been trained, while its embeddings, precisely what linear probing +needs, are excellent. + +References +---------- +.. [1] A. F. Ansari, O. Shchur, J. Küken, et al. Chronos-2: From Univariate + to Universal Forecasting. arXiv:2510.15821, 2025. +.. [2] A. F. Ansari, L. Stella, C. Turkmen, et al. Chronos: Learning the + Language of Time Series. TMLR, 2024. +.. [3] A. Das, W. Kong, R. Sen, Y. Zhou. A decoder-only foundation model for + time-series forecasting. ICML, 2024. +.. [4] G. Woo, C. Liu, A. Kumar, et al. Unified Training of Universal Time + Series Forecasting Transformers. ICML, 2024. +.. [5] V. Ekambaram, A. Jati, P. Dayama, et al. Tiny Time Mixers (TTMs): Fast + Pre-trained Models for Enhanced Zero/Few-Shot Forecasting of Multivariate + Time Series. NeurIPS, 2024. +.. [6] M. Goswami, K. Szafer, A. Choudhry, et al. MOMENT: A Family of Open + Time-series Foundation Models. ICML, 2024. +.. [7] X. Shi, S. Wang, Y. Nie, et al. Time-MoE: Billion-Scale Time Series + Foundation Models with Mixture of Experts. ICLR, 2025. +""" + +############################################################################## +# Data +# ---- +# +# The same small set of noisy sine waves used by the Chronos-2 example is +# reused here, so that Chronos-Bolt's forecasts below can be read against +# the same yardstick. + +import numpy as np + +from tslearn.preprocessing import TimeSeriesScalerMeanVariance + +rng = np.random.RandomState(0) + +n_ts, sz, horizon, context_length = 16, 256, 24, 96 +t = np.arange(sz + horizon) + +periods = rng.uniform(30, 50, size=n_ts) +phases = rng.uniform(0, 2 * np.pi, size=n_ts) + +full_series = np.sin( + 2 * np.pi * t[None, :] / periods[:, None] + phases[:, None] +) + 0.1 * rng.randn(n_ts, sz + horizon) +full_series = TimeSeriesScalerMeanVariance().fit_transform(full_series) + +X_train, X_test = full_series[:, :sz], full_series[:, sz:] + +############################################################################## +# Chronos-2 +# --------- +# +# Chronos-2 [1]_ is covered in depth by the :doc:`plot_foundation_forecasting` +# example: zero-shot forecasting with :class:`~tslearn.foundation.ZeroShotForecaster` +# needs nothing beyond the pipeline itself, and linear probing with +# :class:`~tslearn.foundation.LinearProbeForecaster` reads its forecast token +# (``layer=-2, pooling="token", token_index=-1, tokens=(0, -2)``). It is not +# repeated here. + +############################################################################## +# Chronos-Bolt +# ------------ +# +# Chronos-Bolt [2]_ is a faster, encoder-decoder T5 model, distilled to +# predict all quantiles in a single forward pass rather than +# autoregressively, from the same package as Chronos-2:: +# +# pip install "chronos-forecasting>=2.0" +# +# Two details set it apart from Chronos-2: +# +# * its ``predict`` method requires a :class:`torch.Tensor`, not a NumPy +# array, so the auto-detected calling convention of +# :class:`~tslearn.foundation.ZeroShotForecaster` does not apply and an +# explicit ``predict_fn`` is needed; +# * its ``forward`` takes a raw ``context`` and exposes an encoder made of a +# plain stack of T5 blocks, at ``encoder.block``, with no register or +# forecast token to filter out, so probing needs no ``tokens`` argument. + +import torch +from chronos import BaseChronosPipeline + +from tslearn.foundation import LinearProbeForecaster, ZeroShotForecaster + +pipeline = BaseChronosPipeline.from_pretrained( + "amazon/chronos-bolt-small", device_map="cpu" +) + +zero_shot = ZeroShotForecaster( + pipeline, + predict_fn=lambda model, context, horizon: model.predict( + torch.as_tensor(context, dtype=torch.float32), prediction_length=horizon + ), + horizon_axis=-1, +) +y_zero_shot = zero_shot.predict(X_train, n=horizon) + +probe = LinearProbeForecaster( + pipeline.model, + context_length=context_length, + horizon=horizon, + stride=8, + layer=-1, + layers_path="encoder.block", + pooling="mean", +) +probe.fit(X_train) +y_probe = probe.predict(X_train) + +from tslearn.metrics.performance import mae + +print(f"{'Chronos-Bolt (zero-shot)':>28}: MAE = {mae(X_test, y_zero_shot):.4f}") +print(f"{'Chronos-Bolt (linear probe)':>28}: MAE = {mae(X_test, y_probe):.4f}") + +############################################################################## +import matplotlib.pyplot as plt + +fig, ax = plt.subplots(figsize=(11, 3), layout="constrained") +context = 80 +ax.plot(np.arange(sz - context, sz), X_train[0, -context:, 0], color="0.4", label="context") +ax.plot(np.arange(sz, sz + horizon), X_test[0, :, 0], color="k", label="ground truth") +ax.plot(np.arange(sz, sz + horizon), y_zero_shot[0, :, 0], label="zero-shot", alpha=0.8) +ax.plot(np.arange(sz, sz + horizon), y_probe[0, :, 0], label="linear probe", alpha=0.8) +ax.axvline(sz, color="0.8", linestyle="--") +ax.legend(loc="upper left", ncol=4, fontsize="small") +ax.set_title("Chronos-Bolt") +plt.show() + +############################################################################## +# TimesFM +# ------- +# +# TimesFM [3]_ is a decoder-only model, pre-trained to predict a whole patch +# of future values at once rather than one step at a time:: +# +# pip install "timesfm[torch]" +# +# Its ``forecast`` method takes the horizon *before* the context +# (``forecast(horizon, inputs)``), which conflicts with the positional +# calling convention :class:`~tslearn.foundation.ZeroShotForecaster` uses +# when auto-detecting a method, so ``predict_fn`` is required here too. The +# horizon and context length are fixed once and for all through a +# :class:`~timesfm.ForecastConfig` passed to ``compile``, which builds the +# model's static computation graph. +# +# .. code-block:: python +# +# import timesfm +# +# from tslearn.foundation import ZeroShotForecaster +# +# model = timesfm.TimesFM_2p5_200M_torch.from_pretrained( +# "google/timesfm-2.5-200m-pytorch" +# ) +# model.compile(timesfm.ForecastConfig( +# max_context=context_length, max_horizon=horizon, normalize_inputs=True +# )) +# +# zero_shot = ZeroShotForecaster( +# model, +# predict_fn=lambda model, context, horizon: model.forecast( +# horizon=horizon, inputs=list(context) +# )[0], +# context_length=context_length, +# ) +# y_zero_shot = zero_shot.predict(X_train, n=horizon) +# +# TimesFM's own module (``model.model``) does not lend itself to linear +# probing out of the box: its ``forward`` expects the context already split +# into patches and concatenated with an explicit padding mask, a +# transformation that :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` +# does not perform, so probing it would require wrapping that patching logic +# in a small ``torch.nn.Module`` of one's own. + +############################################################################## +# Moirai +# ------ +# +# Moirai [4]_ is a masked-encoder model, natively multivariate and trained +# across many frequencies at once:: +# +# pip install "uni2ts @ git+https://github.com/SalesforceAIResearch/uni2ts.git" +# +# The forecasting horizon, and the number of sample paths drawn from the +# predictive distribution, are baked into the model at construction time +# rather than passed to a ``predict``-like method, so, again, a +# ``predict_fn`` builds the ``past_target`` / ``past_observed_target`` / +# ``past_is_pad`` triplet that :class:`~uni2ts.model.moirai.MoiraiForecast` +# expects. +# +# .. code-block:: python +# +# import torch +# +# from uni2ts.model.moirai import MoiraiForecast, MoiraiModule +# +# from tslearn.foundation import ZeroShotForecaster +# +# module = MoiraiModule.from_pretrained("Salesforce/moirai-1.1-R-small") +# forecast_model = MoiraiForecast( +# module=module, +# prediction_length=horizon, +# context_length=context_length, +# patch_size=32, +# target_dim=1, +# feat_dynamic_real_dim=0, +# past_feat_dynamic_real_dim=0, +# ) +# +# def predict_fn(model, context, horizon): +# past_target = torch.as_tensor(context, dtype=torch.float32).unsqueeze(-1) +# past_observed = torch.ones_like(past_target, dtype=torch.bool) +# past_is_pad = torch.zeros(past_target.shape[:2], dtype=torch.bool) +# return model(past_target, past_observed, past_is_pad, num_samples=20) +# +# zero_shot = ZeroShotForecaster( +# forecast_model, predict_fn=predict_fn, context_length=context_length +# ) +# y_zero_shot = zero_shot.predict(X_train, n=horizon) +# +# Its backbone (``MoiraiModule``) consumes a packed representation of the +# whole batch, with explicit ``sample_id``, ``variate_id`` and ``time_id`` +# tensors describing which timestep and which series each row belongs to. +# Reproducing that packing outside of ``uni2ts`` itself is enough work that +# zero-shot use is, in practice, the natural way to reach for Moirai. + +############################################################################## +# TTM (Tiny Time Mixers) +# ----------------------- +# +# TTM [5]_ departs from the transformer architecture used by every other +# model in this list: it is a light-weight MLP-Mixer, which is also why it +# is small enough to not need a dedicated "small" checkpoint:: +# +# pip install "granite-tsfm[notebooks] @ git+https://github.com/ibm-granite/granite-tsfm.git@v0.2.22" +# +# Its context length and horizon are fixed by the checkpoint (512 and 96 for +# ``granite-timeseries-ttm-r2``) rather than adjustable at call time, and, +# since the raw model is called directly, ``predict_fn`` again does the +# work that a ``predict`` method would otherwise do. +# +# .. code-block:: python +# +# import torch +# +# from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction +# +# from tslearn.foundation import ZeroShotForecaster +# +# model = TinyTimeMixerForPrediction.from_pretrained( +# "ibm-granite/granite-timeseries-ttm-r2", num_input_channels=1 +# ) +# +# def predict_fn(model, context, horizon): +# past_values = torch.as_tensor(context, dtype=torch.float32).unsqueeze(-1) +# return model(past_values=past_values).prediction_outputs +# +# zero_shot = ZeroShotForecaster( +# model, predict_fn=predict_fn, context_length=model.config.context_length +# ) +# y_zero_shot = zero_shot.predict(X_train, n=model.config.prediction_length) +# +# TTM's mixer blocks operate on a ``(batch, channels, patches, dim)`` tensor +# rather than the ``(batch, tokens, dim)`` shape +# :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` expects, so +# probing it directly would need a small reshaping wrapper too. + +############################################################################## +# MOMENT +# ------ +# +# MOMENT [6]_ is pre-trained as a masked autoencoder: reconstructing masked +# patches is what gives it strong general-purpose embeddings, but it means +# its forecasting head, unlike its embeddings, has never actually been +# trained:: +# +# pip install momentfm +# +# Loaded with ``task_name="embedding"``, a forward pass already returns a +# single pooled vector per series (``output.embeddings``), so this is the +# rare model that needs no ``layer``/``pooling`` gymnastics to reach a +# feature matrix, only a hook on its encoder to read the token-level +# representations that :class:`~tslearn.foundation.LinearProbeForecaster` +# expects. +# +# .. code-block:: python +# +# from momentfm import MOMENTPipeline +# +# from tslearn.foundation import LinearProbeForecaster +# +# model = MOMENTPipeline.from_pretrained( +# "AutonLab/MOMENT-1-small", model_kwargs={"task_name": "embedding"} +# ) +# model.init() +# +# probe = LinearProbeForecaster( +# model, +# context_length=512, +# horizon=horizon, +# stride=32, +# layer=-1, +# layers_path="encoder.block", +# pooling="mean", +# input_layout="channels_first", +# ) +# probe.fit(X_train) +# y_probe = probe.predict(X_train) +# +# ``input_layout="channels_first"`` matters here: MOMENT's ``forward`` +# expects an explicit channel axis (``x_enc`` of shape +# ``(batch, channels, sz)``), unlike most other models in this list which +# take a plain ``(batch, sz)`` batch of univariate contexts. Calling +# :meth:`~tslearn.foundation.LinearProbeForecaster.predict` (a zero-shot +# forecast) would technically run, but on a randomly initialized head, so it +# is deliberately left out. + +############################################################################## +# Time-MoE +# -------- +# +# Time-MoE [7]_ is a decoder-only, GPT-style model trained with next-value +# prediction and a sparse mixture-of-experts feed-forward block, distributed +# as ``transformers`` "remote code" rather than through its own package:: +# +# pip install transformers +# +# .. code-block:: python +# +# from transformers import AutoModelForCausalLM +# +# from tslearn.foundation import LinearProbeForecaster, ZeroShotForecaster +# +# model = AutoModelForCausalLM.from_pretrained( +# "Maple728/TimeMoE-50M", trust_remote_code=True +# ) +# +# Forecasting several steps ahead means generating token by token, through +# ``generate`` rather than a single forward pass; ``max_new_tokens`` is not +# one of the horizon argument names +# :class:`~tslearn.foundation.ZeroShotForecaster` looks for, and the output +# needs slicing to drop the echoed context, so a ``predict_fn`` is used once +# more: +# +# .. code-block:: python +# +# def predict_fn(model, context, horizon): +# context = torch.as_tensor(context, dtype=torch.float32) +# out = model.generate(input_ids=context, max_new_tokens=horizon) +# return out[:, -horizon:] +# +# zero_shot = ZeroShotForecaster(model, predict_fn=predict_fn) +# y_zero_shot = zero_shot.predict(X_train, n=horizon) +# +# Being decoder-only, its natural pooling is ``"last"``, the representation +# of the final context token, which has attended to every earlier one: +# +# .. code-block:: python +# +# probe = LinearProbeForecaster( +# model, +# context_length=context_length, +# horizon=horizon, +# stride=8, +# layer=-1, +# layers_path="model.layers", +# pooling="last", +# ) +# probe.fit(X_train) +# y_probe = probe.predict(X_train) diff --git a/tests/test_foundation_zoo.py b/tests/test_foundation_zoo.py new file mode 100644 index 00000000..fc7b390e --- /dev/null +++ b/tests/test_foundation_zoo.py @@ -0,0 +1,220 @@ +"""Tests exercising :mod:`tslearn.foundation` against real pre-trained models. + +Unlike :mod:`test_foundation`, these tests download actual checkpoints from +the Hugging Face Hub and require model-specific packages (``timesfm``, +``uni2ts``, ``granite-tsfm``, ``momentfm``, a recent ``transformers``...) +that are *not* part of tslearn's own dependencies and that, in some cases, +pin conflicting versions of ``torch`` or ``transformers`` against one +another. Each test is therefore independent, skipped unless both its own +package is importable and the ``TSLEARN_RUN_FOUNDATION_ZOO`` environment +variable is set, and meant to be run in its own dedicated environment (see +``.github/workflows/test_foundation_zoo.yml``) rather than as part of the +regular test suite. + +The code below mirrors the snippets given in the +``plot_foundation_model_zoo.py`` gallery example: this file is what backs +the claim, made there, that each recipe actually works. +""" + +import os + +import numpy as np + +import pytest + +from tslearn.generators import random_walks + +pytestmark = pytest.mark.skipif( + not os.environ.get("TSLEARN_RUN_FOUNDATION_ZOO"), + reason="Set TSLEARN_RUN_FOUNDATION_ZOO=1 to run tests that download " + "real pre-trained models from the Hugging Face Hub.", +) + +N_TS, SZ, CONTEXT_LENGTH, HORIZON = 5, 200, 64, 12 + + +def _data(): + return random_walks(n_ts=N_TS, sz=SZ, random_state=0).astype(np.float32) + + +def test_chronos_bolt(): + torch = pytest.importorskip("torch") + chronos = pytest.importorskip("chronos") + + from tslearn.foundation import LinearProbeForecaster, ZeroShotForecaster + + pipeline = chronos.BaseChronosPipeline.from_pretrained( + "amazon/chronos-bolt-small", device_map="cpu" + ) + + zero_shot = ZeroShotForecaster( + pipeline, + predict_fn=lambda model, context, horizon: model.predict( + torch.as_tensor(context, dtype=torch.float32), prediction_length=horizon + ), + horizon_axis=-1, + ) + X = _data() + y_zero_shot = zero_shot.predict(X, n=HORIZON) + assert y_zero_shot.shape == (N_TS, HORIZON, 1) + + probe = LinearProbeForecaster( + pipeline.model, + context_length=CONTEXT_LENGTH, + horizon=HORIZON, + stride=8, + layer=-1, + layers_path="encoder.block", + pooling="mean", + ) + probe.fit(X) + y_probe = probe.predict(X) + assert y_probe.shape == (N_TS, HORIZON, 1) + + +def test_timesfm(): + torch = pytest.importorskip("torch") + timesfm = pytest.importorskip("timesfm") + + from tslearn.foundation import ZeroShotForecaster + + model = timesfm.TimesFM_2p5_200M_torch.from_pretrained( + "google/timesfm-2.5-200m-pytorch" + ) + model.compile( + timesfm.ForecastConfig( + max_context=CONTEXT_LENGTH, max_horizon=HORIZON, normalize_inputs=True + ) + ) + + zero_shot = ZeroShotForecaster( + model, + predict_fn=lambda model, context, horizon: model.forecast( + horizon=horizon, inputs=list(context) + )[0], + context_length=CONTEXT_LENGTH, + ) + X = _data() + y_zero_shot = zero_shot.predict(X, n=HORIZON) + assert y_zero_shot.shape == (N_TS, HORIZON, 1) + + +def test_moirai(): + torch = pytest.importorskip("torch") + pytest.importorskip("uni2ts") + from uni2ts.model.moirai import MoiraiForecast, MoiraiModule + + from tslearn.foundation import ZeroShotForecaster + + module = MoiraiModule.from_pretrained("Salesforce/moirai-1.1-R-small") + forecast_model = MoiraiForecast( + module=module, + prediction_length=HORIZON, + context_length=CONTEXT_LENGTH, + patch_size=32, + target_dim=1, + feat_dynamic_real_dim=0, + past_feat_dynamic_real_dim=0, + ) + + def predict_fn(model, context, horizon): + past_target = torch.as_tensor(context, dtype=torch.float32).unsqueeze(-1) + past_observed = torch.ones_like(past_target, dtype=torch.bool) + past_is_pad = torch.zeros(past_target.shape[:2], dtype=torch.bool) + with torch.no_grad(): + return model(past_target, past_observed, past_is_pad, num_samples=20) + + zero_shot = ZeroShotForecaster( + forecast_model, predict_fn=predict_fn, context_length=CONTEXT_LENGTH + ) + X = _data() + y_zero_shot = zero_shot.predict(X, n=HORIZON) + assert y_zero_shot.shape == (N_TS, HORIZON, 1) + + +def test_ttm(): + torch = pytest.importorskip("torch") + pytest.importorskip("tsfm_public") + from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction + + from tslearn.foundation import ZeroShotForecaster + + model = TinyTimeMixerForPrediction.from_pretrained( + "ibm-granite/granite-timeseries-ttm-r2", num_input_channels=1, revision="main" + ) + + def predict_fn(model, context, horizon): + past_values = torch.as_tensor(context, dtype=torch.float32).unsqueeze(-1) + with torch.no_grad(): + return model(past_values=past_values).prediction_outputs + + zero_shot = ZeroShotForecaster( + model, predict_fn=predict_fn, context_length=model.config.context_length + ) + X = random_walks( + n_ts=N_TS, sz=model.config.context_length + 50, random_state=0 + ).astype(np.float32) + y_zero_shot = zero_shot.predict(X, n=model.config.prediction_length) + assert y_zero_shot.shape == (N_TS, model.config.prediction_length, 1) + + +def test_moment(): + pytest.importorskip("torch") + pytest.importorskip("momentfm") + from momentfm import MOMENTPipeline + + from tslearn.foundation import LinearProbeForecaster + + model = MOMENTPipeline.from_pretrained( + "AutonLab/MOMENT-1-small", model_kwargs={"task_name": "embedding"} + ) + model.init() + + probe = LinearProbeForecaster( + model, + context_length=512, + horizon=HORIZON, + stride=32, + layer=-1, + layers_path="encoder.block", + pooling="mean", + input_layout="channels_first", + ) + X = random_walks(n_ts=N_TS, sz=700, random_state=0).astype(np.float32) + probe.fit(X) + y_probe = probe.predict(X) + assert y_probe.shape == (N_TS, HORIZON, 1) + + +def test_time_moe(): + torch = pytest.importorskip("torch") + transformers = pytest.importorskip("transformers") + + from tslearn.foundation import LinearProbeForecaster, ZeroShotForecaster + + model = transformers.AutoModelForCausalLM.from_pretrained( + "Maple728/TimeMoE-50M", trust_remote_code=True + ) + + def predict_fn(model, context, horizon): + context = torch.as_tensor(context, dtype=torch.float32) + out = model.generate(input_ids=context, max_new_tokens=horizon) + return out[:, -horizon:] + + zero_shot = ZeroShotForecaster(model, predict_fn=predict_fn) + X = _data() + y_zero_shot = zero_shot.predict(X, n=HORIZON) + assert y_zero_shot.shape == (N_TS, HORIZON, 1) + + probe = LinearProbeForecaster( + model, + context_length=CONTEXT_LENGTH, + horizon=HORIZON, + stride=8, + layer=-1, + layers_path="model.layers", + pooling="last", + ) + probe.fit(X) + y_probe = probe.predict(X) + assert y_probe.shape == (N_TS, HORIZON, 1) From 944bc7c4b0039ef68abd315d775abb9129539e84 Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Fri, 21 Aug 2026 16:31:46 +0200 Subject: [PATCH 12/18] scale only for linear probe models --- .../plot_foundation_forecasting.py | 70 +++++++++++++------ .../forecasting/plot_foundation_model_zoo.py | 51 +++++++++++--- 2 files changed, 88 insertions(+), 33 deletions(-) diff --git a/docs/examples/forecasting/plot_foundation_forecasting.py b/docs/examples/forecasting/plot_foundation_forecasting.py index 6ede9ee8..a4ac76ab 100644 --- a/docs/examples/forecasting/plot_foundation_forecasting.py +++ b/docs/examples/forecasting/plot_foundation_forecasting.py @@ -38,7 +38,7 @@ import numpy as np -from tslearn.preprocessing import TimeSeriesScalerMeanVariance +from tslearn.utils import to_time_series_dataset rng = np.random.RandomState(0) @@ -54,7 +54,7 @@ + trends[:, None] * t[None, :] + 0.1 * rng.randn(n_ts, sz + horizon) ) -full_series = TimeSeriesScalerMeanVariance().fit_transform(full_series) +full_series = to_time_series_dataset(full_series) X_train, X_test = full_series[:, :sz], full_series[:, sz:] print(f"{X_train.shape=}, {X_test.shape=}") @@ -69,9 +69,14 @@ # tslearn estimator, whose ``predict`` method returns an array of shape # ``(n_ts, n, d)`` like every other forecaster of the library. # -# Calling ``fit`` is not needed here since the purpose of a -# :class:`~tslearn.foundation.ZeroShotForecaster` is to reuse a pre-trained +# Calling ``fit`` is not needed here since the purpose of a +# :class:`~tslearn.foundation.ZeroShotForecaster` is to reuse a pre-trained # model without any training. +# +# Note that the raw series are passed as they are, with no prior scaling: +# :class:`~tslearn.foundation.ZeroShotForecaster` wraps Chronos-2's full +# inference pipeline, which normalizes every context internally before +# forecasting. from chronos import Chronos2Pipeline @@ -108,8 +113,14 @@ # matters because every window has to go through the pre-trained model once. # # The probing estimator needs the model itself rather than the inference -# pipeline, since it reads hidden states rather than forecasts. Three options -# drive which representations are used: +# pipeline, since it reads hidden states rather than forecasts. This means it +# bypasses the normalization the pipeline applies internally, unlike +# :class:`~tslearn.foundation.ZeroShotForecaster` above, so it has to be +# restored explicitly, by placing a +# :class:`~tslearn.preprocessing.TimeSeriesScalerMeanVariance` ahead of it in +# a :class:`~sklearn.pipeline.Pipeline`. +# +# Three options drive which representations are used: # # * ``layer`` selects the block to probe: an integer places a forward hook on # the corresponding block, so ``layer=-2`` probes the penultimate one. The @@ -137,9 +148,12 @@ # :alt: Chronos-2 appends a register token and a forecast token after its # context tokens; tokens=(0, -2) keeps only the context tokens. +from sklearn.pipeline import Pipeline + from tslearn.foundation import LinearProbeForecaster +from tslearn.preprocessing import TimeSeriesScalerMeanVariance -probe = LinearProbeForecaster( +forecaster = LinearProbeForecaster( pipeline.model, context_length=context_length, horizon=horizon, @@ -148,17 +162,25 @@ pooling="mean", tokens=(0, -2), ) +scaler = TimeSeriesScalerMeanVariance(per_timeseries=False) +probe = Pipeline([("scale", scaler), ("probe", forecaster)]) probe.fit(X_train) -y_probe = probe.predict(X_train) -print(f"{probe.n_windows_} training windows, " - f"{probe.embedder_.embedding_size_}-dimensional embeddings") +print(f"{forecaster.n_windows_} training windows, " + f"{forecaster.embedder_.embedding_size_}-dimensional embeddings") ############################################################################## # The head fitted on top of the frozen representations defaults to a # :class:`sklearn.linear_model.RidgeCV`, which picks its regularization # strength by cross-validation. Any scikit-learn regressor can be passed # instead through the ``probe`` parameter. +# +# The scaler also normalizes the training targets, since they are cut from +# the same series as the context, so the forecasts come out on that +# normalized scale. They are put back in the original units using the +# ``mean_`` and ``std_`` the scaler computed when it was fitted. + +y_probe = probe.predict(X_train) * scaler.std_ + scaler.mean_ ############################################################################## # Comparison @@ -227,18 +249,22 @@ results = {} for layer in [-1, -2, -4]: for pooling in ["mean", "token"]: - model = LinearProbeForecaster( - pipeline.model, - context_length=context_length, - horizon=horizon, - stride=48, - layer=layer, - pooling=pooling, - # Chronos-2's forecast token is its last one - token_index=-1, - tokens=(0, -2), - ).fit(X_train) - results[(layer, pooling)] = mae(X_test, model.predict(X_train)) + model = Pipeline([ + ("scale", scaler), + ("probe", LinearProbeForecaster( + pipeline.model, + context_length=context_length, + horizon=horizon, + stride=48, + layer=layer, + pooling=pooling, + # Chronos-2's forecast token is its last one + token_index=-1, + tokens=(0, -2), + )), + ]).fit(X_train) + y = model.predict(X_train) * scaler.std_ + scaler.mean_ + results[(layer, pooling)] = mae(X_test, y) for (layer, pooling), score in sorted(results.items(), key=lambda kv: kv[1]): print(f"layer={str(layer):>5}, pooling={pooling:>4}: MAE = {score:.4f}") diff --git a/docs/examples/forecasting/plot_foundation_model_zoo.py b/docs/examples/forecasting/plot_foundation_model_zoo.py index d787b233..86f276ee 100644 --- a/docs/examples/forecasting/plot_foundation_model_zoo.py +++ b/docs/examples/forecasting/plot_foundation_model_zoo.py @@ -71,7 +71,7 @@ import numpy as np -from tslearn.preprocessing import TimeSeriesScalerMeanVariance +from tslearn.utils import to_time_series_dataset rng = np.random.RandomState(0) @@ -84,7 +84,7 @@ full_series = np.sin( 2 * np.pi * t[None, :] / periods[:, None] + phases[:, None] ) + 0.1 * rng.randn(n_ts, sz + horizon) -full_series = TimeSeriesScalerMeanVariance().fit_transform(full_series) +full_series = to_time_series_dataset(full_series) X_train, X_test = full_series[:, :sz], full_series[:, sz:] @@ -118,11 +118,20 @@ # * its ``forward`` takes a raw ``context`` and exposes an encoder made of a # plain stack of T5 blocks, at ``encoder.block``, with no register or # forecast token to filter out, so probing needs no ``tokens`` argument. +# +# As with Chronos-2, the raw series are passed to +# :class:`~tslearn.foundation.ZeroShotForecaster` unscaled, since it wraps +# the full pipeline. The linear probe, on the other hand, reads +# ``pipeline.model`` directly and so needs scaling restored explicitly through +# a :class:`~sklearn.pipeline.Pipeline`, as detailed in the +# :doc:`plot_foundation_forecasting` example. import torch from chronos import BaseChronosPipeline +from sklearn.pipeline import Pipeline from tslearn.foundation import LinearProbeForecaster, ZeroShotForecaster +from tslearn.preprocessing import TimeSeriesScalerMeanVariance pipeline = BaseChronosPipeline.from_pretrained( "amazon/chronos-bolt-small", device_map="cpu" @@ -137,7 +146,7 @@ ) y_zero_shot = zero_shot.predict(X_train, n=horizon) -probe = LinearProbeForecaster( +forecaster = LinearProbeForecaster( pipeline.model, context_length=context_length, horizon=horizon, @@ -146,8 +155,14 @@ layers_path="encoder.block", pooling="mean", ) +scaler = TimeSeriesScalerMeanVariance(per_timeseries=False) +probe = Pipeline([("scale", scaler), ("probe", forecaster)]) probe.fit(X_train) -y_probe = probe.predict(X_train) + +# The scaler also normalizes the training targets, so forecasts come out on +# that scale and are put back in the original units, using the ``mean_`` and +# ``std_`` the scaler computed when it was fitted, before comparison. +y_probe = probe.predict(X_train) * scaler.std_ + scaler.mean_ from tslearn.metrics.performance import mae @@ -325,18 +340,25 @@ # representations that :class:`~tslearn.foundation.LinearProbeForecaster` # expects. # +# As with the linear probes above, ``model`` is the bare backbone rather than +# a pipeline, so scaling has to be restored the same way, through a +# :class:`~sklearn.pipeline.Pipeline`. +# # .. code-block:: python # # from momentfm import MOMENTPipeline +# from sklearn.pipeline import Pipeline # # from tslearn.foundation import LinearProbeForecaster +# from tslearn.preprocessing import TimeSeriesScalerMeanVariance # # model = MOMENTPipeline.from_pretrained( # "AutonLab/MOMENT-1-small", model_kwargs={"task_name": "embedding"} # ) # model.init() # -# probe = LinearProbeForecaster( +# scaler = TimeSeriesScalerMeanVariance(per_timeseries=False) +# probe = Pipeline([("scale", scaler), ("probe", LinearProbeForecaster( # model, # context_length=512, # horizon=horizon, @@ -345,9 +367,9 @@ # layers_path="encoder.block", # pooling="mean", # input_layout="channels_first", -# ) +# ))]) # probe.fit(X_train) -# y_probe = probe.predict(X_train) +# y_probe = probe.predict(X_train) * scaler.std_ + scaler.mean_ # # ``input_layout="channels_first"`` matters here: MOMENT's ``forward`` # expects an explicit channel axis (``x_enc`` of shape @@ -395,11 +417,18 @@ # y_zero_shot = zero_shot.predict(X_train, n=horizon) # # Being decoder-only, its natural pooling is ``"last"``, the representation -# of the final context token, which has attended to every earlier one: +# of the final context token, which has attended to every earlier one. Here +# too, ``model`` is the bare backbone rather than a normalizing pipeline, so +# scaling is restored through a :class:`~sklearn.pipeline.Pipeline`: # # .. code-block:: python # -# probe = LinearProbeForecaster( +# from sklearn.pipeline import Pipeline +# +# from tslearn.preprocessing import TimeSeriesScalerMeanVariance +# +# scaler = TimeSeriesScalerMeanVariance(per_timeseries=False) +# probe = Pipeline([("scale", scaler), ("probe", LinearProbeForecaster( # model, # context_length=context_length, # horizon=horizon, @@ -407,6 +436,6 @@ # layer=-1, # layers_path="model.layers", # pooling="last", -# ) +# ))]) # probe.fit(X_train) -# y_probe = probe.predict(X_train) +# y_probe = probe.predict(X_train) * scaler.std_ + scaler.mean_ From 933fbc12674a8509c3e897e2ff880b99382a65c0 Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Sat, 22 Aug 2026 16:25:02 +0200 Subject: [PATCH 13/18] make it time-proof --- docs/examples/forecasting/plot_foundation_model_zoo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/examples/forecasting/plot_foundation_model_zoo.py b/docs/examples/forecasting/plot_foundation_model_zoo.py index 86f276ee..c8f33ff5 100644 --- a/docs/examples/forecasting/plot_foundation_model_zoo.py +++ b/docs/examples/forecasting/plot_foundation_model_zoo.py @@ -7,7 +7,7 @@ ``forward``, following one of a handful of widespread conventions (see :class:`~tslearn.foundation.ZeroShotForecaster` and :class:`~tslearn.foundation.LinearProbeForecaster`). This example surveys -seven models currently considered state of the art for time series +seven models for time series forecasting, and shows, for each of them, what it takes to plug it in as a :class:`~tslearn.foundation.ZeroShotForecaster`, a :class:`~tslearn.foundation.LinearProbeForecaster`, or both. The smallest From 85f7722fdffffaaa0c44103f50ac4cfcd0e2fbd3 Mon Sep 17 00:00:00 2001 From: Romain Tavenard Date: Mon, 24 Aug 2026 09:23:11 +0200 Subject: [PATCH 14/18] removing useless LinearProbeClassifier --- CHANGELOG.md | 8 +- .../plot_foundation_linear_probe.py | 68 ++-- docs/gen_modules/tslearn.foundation.rst | 1 - tests/test_foundation.py | 93 +----- tslearn/foundation/__init__.py | 16 +- tslearn/foundation/_classification.py | 300 ------------------ tslearn/foundation/_embedding.py | 6 +- tslearn/foundation/_forecasting.py | 4 +- 8 files changed, 59 insertions(+), 437 deletions(-) delete mode 100644 tslearn/foundation/_classification.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ab7a20c9..99f88c19 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,11 +15,13 @@ Changelogs for this project are recorded in this file since v0.2.0. * New `tslearn.foundation` module to re-use pre-trained time series models, such as the ones published on the Hugging Face Hub, behind the usual tslearn API. It provides - `ZeroShotForecaster`, `LinearProbeForecaster`, `LinearProbeClassifier` and the - underlying `TimeSeriesFoundationEmbedder` feature extractor, which allows one to choose + `ZeroShotForecaster`, `LinearProbeForecaster` and the underlying + `TimeSeriesFoundationEmbedder` feature extractor, which allows one to choose the probed layer, the tokens taken into account and the pooling applied to their representations. With `pooling=None`, the latter acts as a time series to time series - transform, returning one representation per token. Requires PyTorch. + transform, returning one representation per token, and being a regular scikit-learn + transformer, it composes with a `sklearn.pipeline.Pipeline` for tasks such as linear + probing for classification. Requires PyTorch. * soft-dtw related tools now support Itakura and Sakoe-Chiba global constraints. ([#189](https://github.com/tslearn-team/tslearn/issues/189)) * Multithreading support added to `cdist_soft_dtw` and `softdtw_barycenter`. ([#310](https://github.com/tslearn-team/tslearn/issues/310)) diff --git a/docs/examples/classification/plot_foundation_linear_probe.py b/docs/examples/classification/plot_foundation_linear_probe.py index 006a2105..fe0a98a4 100644 --- a/docs/examples/classification/plot_foundation_linear_probe.py +++ b/docs/examples/classification/plot_foundation_linear_probe.py @@ -11,9 +11,10 @@ never updated, the accuracy reached tells us how linearly separable the classes already are in the representation space. -This example applies :class:`~tslearn.foundation.LinearProbeClassifier` to a -UCR dataset, using Chronos-2 [2]_ as the frozen backbone. -Running it requires the ``chronos-forecasting`` package:: +This example applies that idea to a UCR dataset, composing +:class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` with a classifier +inside a :class:`~sklearn.pipeline.Pipeline`, and using Chronos-2 [2]_ as the +frozen backbone. Running it requires the ``chronos-forecasting`` package:: pip install "chronos-forecasting>=2.0" @@ -54,22 +55,27 @@ # Probing the pre-trained model # ----------------------------- # -# The classifier only needs the pre-trained model and, since Chronos-2 returns +# The embedder only needs the pre-trained model and, since Chronos-2 returns # forecasts rather than hidden states, an explicit layer to read # representations from. A forward hook is placed on the selected block, so no -# modification of the model is needed. +# modification of the model is needed. Being a regular scikit-learn +# transformer, it composes with a classifier inside a +# :class:`~sklearn.pipeline.Pipeline`, which is what "linear probing" amounts +# to here. # # Fitting is fast: each series goes through the model exactly once, and the # only thing actually trained is a logistic regression over a few hundred # features. from chronos import Chronos2Pipeline +from sklearn.linear_model import LogisticRegression +from sklearn.pipeline import Pipeline -from tslearn.foundation import LinearProbeClassifier +from tslearn.foundation import TimeSeriesFoundationEmbedder pipeline = Chronos2Pipeline.from_pretrained("autogluon/chronos-2-small") -clf = LinearProbeClassifier( +embedder = TimeSeriesFoundationEmbedder( pipeline.model, layer=-2, pooling="mean", @@ -77,9 +83,10 @@ # tokens; neither represents the series, so they are left out of the average tokens=(0, -2), ) +clf = Pipeline([("embed", embedder), ("classify", LogisticRegression(max_iter=1000))]) clf.fit(X_train, y_train) -print(f"Embedding size: {clf.embedder_.embedding_size_}") +print(f"Embedding size: {embedder.embedding_size_}") print(f"Test accuracy: {clf.score(X_test, y_test):.3f}") ############################################################################## @@ -104,14 +111,17 @@ accuracies = {} for layer in layers: for pooling in poolings: - model = LinearProbeClassifier( - pipeline.model, - layer=layer, - pooling=pooling, - # Chronos-2 places its register token after the context tokens - token_index=-2 if pooling == "token" else 0, - tokens=(0, -2), - ).fit(X_train, y_train) + model = Pipeline([ + ("embed", TimeSeriesFoundationEmbedder( + pipeline.model, + layer=layer, + pooling=pooling, + # Chronos-2 places its register token after the context tokens + token_index=-2 if pooling == "token" else 0, + tokens=(0, -2), + )), + ("classify", LogisticRegression(max_iter=1000)), + ]).fit(X_train, y_train) accuracies[layer, pooling] = model.score(X_test, y_test) header = "layer " + " ".join(f"{pooling:>6}" for pooling in poolings) @@ -142,22 +152,23 @@ # Using the representations elsewhere # ----------------------------------- # -# The feature extractor can also be used on its own, through -# :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. Being a regular -# scikit-learn transformer, it composes with the rest of the ecosystem: here we -# project the frozen representations of the test set onto two dimensions to see -# whether the classes separate. +# :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` is a regular +# scikit-learn transformer, so it composes with the rest of the ecosystem the +# same way it did with the classifier above: here it is chained with a +# :class:`~sklearn.decomposition.PCA` inside a +# :class:`~sklearn.pipeline.Pipeline`, to project the frozen representations +# of the test set onto two dimensions and see whether the classes separate. from sklearn.decomposition import PCA -from tslearn.foundation import TimeSeriesFoundationEmbedder - best_layer, best_pooling = max(accuracies, key=accuracies.get) -embedder = TimeSeriesFoundationEmbedder( - pipeline.model, layer=best_layer, tokens=(0, -2) -) -embeddings = embedder.fit_transform(X_test) -projected = PCA(n_components=2).fit_transform(embeddings) +projector = Pipeline([ + ("embed", TimeSeriesFoundationEmbedder( + pipeline.model, layer=best_layer, tokens=(0, -2) + )), + ("pca", PCA(n_components=2)), +]) +projected = projector.fit_transform(X_test) fig, ax = plt.subplots(figsize=(6, 5), layout="constrained") for label in np.unique(y_test): @@ -187,7 +198,6 @@ from tslearn.clustering import TimeSeriesKMeans -from sklearn.pipeline import Pipeline from sklearn.metrics import adjusted_rand_score embedder = TimeSeriesFoundationEmbedder( diff --git a/docs/gen_modules/tslearn.foundation.rst b/docs/gen_modules/tslearn.foundation.rst index a3787ed9..4a81b52f 100644 --- a/docs/gen_modules/tslearn.foundation.rst +++ b/docs/gen_modules/tslearn.foundation.rst @@ -13,7 +13,6 @@ tslearn.foundation ZeroShotForecaster LinearProbeForecaster - LinearProbeClassifier TimeSeriesFoundationEmbedder \ No newline at end of file diff --git a/tests/test_foundation.py b/tests/test_foundation.py index 2be9cee6..81704b41 100644 --- a/tests/test_foundation.py +++ b/tests/test_foundation.py @@ -11,7 +11,6 @@ from sklearn.base import clone from sklearn.linear_model import Ridge -from sklearn.naive_bayes import GaussianNB from sklearn.pipeline import make_pipeline from sklearn.svm import LinearSVC @@ -20,7 +19,6 @@ torch = pytest.importorskip("torch") from tslearn.foundation import ( # noqa: E402 - LinearProbeClassifier, LinearProbeForecaster, TimeSeriesFoundationEmbedder, ZeroShotForecaster, @@ -283,10 +281,8 @@ def test_embedder_token_selection_errors(): def test_probes_reject_unpooled_representations(): - X, y = _classification_dataset() + X = _dataset(n_ts=4, sz=32, d=1) for pooling in (None, "none"): - with pytest.raises(ValueError, match="pooling=None"): - LinearProbeClassifier(_DummyBackbone(), pooling=pooling).fit(X, y) with pytest.raises(ValueError, match="pooling=None"): LinearProbeForecaster( _DummyBackbone(), pooling=pooling, context_length=16, horizon=2 @@ -294,10 +290,6 @@ def test_probes_reject_unpooled_representations(): def test_probes_accept_token_selection(): - X, y = _classification_dataset() - model = LinearProbeClassifier(_DummyBackbone(), tokens=(0, -1)).fit(X, y) - assert model.predict(X).shape == (len(y),) - X_fc = _dataset(n_ts=6, sz=48, d=1) model = LinearProbeForecaster( _DummyBackbone(), context_length=16, horizon=2, stride=8, tokens=(0, -1) @@ -1156,94 +1148,11 @@ def test_linear_probe_forecaster_fit_predict(): np.testing.assert_allclose(model.fit_predict(X, n=2), model.predict(X, n=2)) -# --------------------------------------------------------------------------- -# LinearProbeClassifier -# --------------------------------------------------------------------------- - - -def _classification_dataset(n_per_class=15, sz=32, seed=0): - rng = np.random.RandomState(seed) - t = np.linspace(0, 4 * np.pi, sz) - sines = np.sin(t)[None, :] + 0.1 * rng.randn(n_per_class, sz) - lines = np.linspace(-1, 1, sz)[None, :] + 0.1 * rng.randn(n_per_class, sz) - X = np.concatenate([sines, lines])[:, :, None] - y = np.array(["sine"] * n_per_class + ["line"] * n_per_class) - return X, y - - -def test_linear_probe_classifier(): - X, y = _classification_dataset() - model = LinearProbeClassifier(_DummyBackbone()).fit(X, y) - - assert set(model.classes_) == {"sine", "line"} - assert model.predict(X).shape == (len(y),) - assert model.predict_proba(X).shape == (len(y), 2) - assert model.decision_function(X).shape == (len(y),) - np.testing.assert_allclose(model.predict_proba(X).sum(axis=1), 1.0) - # These two classes are easily separable - assert model.score(X, y) > 0.9 - - -def test_linear_probe_classifier_multivariate(): - X, y = _classification_dataset() - X = np.concatenate([X, X[::-1]], axis=2) - model = LinearProbeClassifier(_DummyBackbone()).fit(X, y) - assert model.embedder_.embedding_size_ == 2 * D_MODEL - assert model.predict(X).shape == (len(y),) - - -def test_linear_probe_classifier_accepts_any_probe(): - X, y = _classification_dataset() - model = LinearProbeClassifier(_DummyBackbone(), probe=LinearSVC()).fit(X, y) - assert isinstance(model.probe_, LinearSVC) - with pytest.raises(AttributeError, match="predict_proba"): - model.predict_proba(X) - assert model.decision_function(X).shape == (len(y),) - - -def test_linear_probe_classifier_no_decision_function(): - X, y = _classification_dataset() - model = LinearProbeClassifier(_DummyBackbone(), probe=GaussianNB()).fit(X, y) - assert isinstance(model.probe_, GaussianNB) - with pytest.raises(AttributeError, match="decision_function"): - model.decision_function(X) - assert model.predict_proba(X).shape == (len(y), 2) - - -def test_linear_probe_classifier_layer_and_pooling(): - X, y = _classification_dataset() - for layer in (0, 1, None): - for pooling in ("mean", "max", "token"): - model = LinearProbeClassifier( - _DummyBackbone(), layer=layer, pooling=pooling - ).fit(X, y) - assert model.predict(X).shape == (len(y),) - - -def test_linear_probe_classifier_transform(): - X, y = _classification_dataset() - model = LinearProbeClassifier(_DummyBackbone()).fit(X, y) - assert model.transform(X).shape == (len(y), D_MODEL) - - -def test_linear_probe_classifier_errors(): - X, y = _classification_dataset() - with pytest.raises(ValueError, match="inconsistent numbers of samples"): - LinearProbeClassifier(_DummyBackbone()).fit(X, y[:-1]) - with pytest.raises(ValueError, match="input_layout"): - LinearProbeClassifier(_DummyBackbone(), input_layout="bogus").fit(X, y) - - model = LinearProbeClassifier(_DummyBackbone()).fit(X, y) - with pytest.raises(ValueError, match="features"): - model.predict(np.concatenate([X, X], axis=2)) - - def test_estimators_are_clonable(): for estimator in ( TimeSeriesFoundationEmbedder(_DummyBackbone()), ZeroShotForecaster(_DummyPipeline()), LinearProbeForecaster(_DummyBackbone()), - LinearProbeClassifier(_DummyBackbone()), ): cloned = clone(estimator) assert cloned is not estimator diff --git a/tslearn/foundation/__init__.py b/tslearn/foundation/__init__.py index ad336514..dbf7f78b 100644 --- a/tslearn/foundation/__init__.py +++ b/tslearn/foundation/__init__.py @@ -3,16 +3,18 @@ time series models, such as the ones published on the Hugging Face Hub, behind the usual tslearn API. -Three adaptation strategies are covered: +Two adaptation strategies are covered: * zero-shot forecasting, with :class:`ZeroShotForecaster`, which uses a pre-trained forecaster as-is; -* linear probing for forecasting, with :class:`LinearProbeForecaster`; -* linear probing for classification, with :class:`LinearProbeClassifier`. +* linear probing for forecasting, with :class:`LinearProbeForecaster`. -The last two rely on :class:`TimeSeriesFoundationEmbedder`, which turns any -frozen PyTorch model into a feature extractor and lets one choose which layer -to read representations from and how to pool them. +Linear probing relies on :class:`TimeSeriesFoundationEmbedder`, which turns any frozen +PyTorch model into a feature extractor and lets one choose which layer to +read representations from and how to pool them. Being a regular +scikit-learn transformer, it also covers linear probing for classification, +or any other downstream task, by composing with a +:class:`sklearn.pipeline.Pipeline`. These estimators are deliberately agnostic to any specific model implementation: they duck-type the wrapped object and offer escape hatches @@ -25,12 +27,10 @@ """ -from ._classification import LinearProbeClassifier from ._embedding import TimeSeriesFoundationEmbedder from ._forecasting import LinearProbeForecaster, ZeroShotForecaster __all__ = [ - "LinearProbeClassifier", "LinearProbeForecaster", "TimeSeriesFoundationEmbedder", "ZeroShotForecaster", diff --git a/tslearn/foundation/_classification.py b/tslearn/foundation/_classification.py deleted file mode 100644 index 0774dfab..00000000 --- a/tslearn/foundation/_classification.py +++ /dev/null @@ -1,300 +0,0 @@ -"""Re-use of pre-trained time series models for classification.""" - -import numpy as np - -from sklearn.base import BaseEstimator, ClassifierMixin, clone -from sklearn.linear_model import LogisticRegression -from sklearn.utils.multiclass import check_classification_targets -from sklearn.utils.validation import check_is_fitted - -from tslearn.bases import TimeSeriesMixin -from tslearn.utils import check_array, to_time_series_dataset - -from ._embedding import ( - LAYOUTS, - TimeSeriesFoundationEmbedder, - _check_probe_pooling, - _require_torch, -) - - -class LinearProbeClassifier(TimeSeriesMixin, ClassifierMixin, BaseEstimator): - """Classify time series with a head on a frozen pre-trained model. - - The pre-trained model is used as a frozen feature extractor and a - classifier is fitted on the resulting representations. - - Linear probing keeps the pre-trained model entirely frozen and only fits a - map from its representations to the class predictions. Compared to fine-tuning, - it leaves the pre-trained weights untouched, which makes it much cheaper. - - Parameters - ---------- - model : torch.nn.Module - A pre-trained model used as a frozen feature extractor. - probe : sklearn estimator or None (default: None) - The head fitted on top of the frozen representations. When None, a - :class:`sklearn.linear_model.LogisticRegression` is used. Any - scikit-learn classifier can be passed instead. - Frozen representations sometimes have a very small - variance, in which case prepending a - :class:`sklearn.preprocessing.StandardScaler` to the head might help. - context_length : int or None (default: None) - When set, only the last ``context_length`` timestamps of each series - are fed to the model. - layer : int or None (default: None) - Layer to probe, see - :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. - layers_path : str or None (default: None) - Dotted path to the stack of layers, see - :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. - pooling : {"mean", "max", "token", "last", "flatten"} (default: "mean") - How token representations are aggregated, see - :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` and - note that ``pooling=None`` is not accepted here. - Default is mean pooling, but if the pre-trained model outputs a - class token, it should be better to set ``token_index`` - to this token's index and use ``"token"`` as a pooling strategy. - tokens : slice, (start, stop) pair or None (default: None) - Which tokens to keep before pooling, see - :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. Restricting - the average to the tokens that actually represent the context, e.g. - ``tokens=(0, -2)`` for Chronos-2, avoids diluting it with class or - forecast tokens. - token_index : int (default: 0) - Index of the token selected when ``pooling="token"``, see - :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. - input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") - How the context is laid out when handed over to the model. - ``"univariate"`` forecasts every channel of every series - independently. The other two layouts feed - a ``(n_ts, sz, d)`` or ``(n_ts, d, sz)`` array respectively, for - models that are natively multivariate. - input_name : str or None (default: None) - Name of the ``forward`` argument receiving the context values. - model_kwargs : dict or None (default: None) - Extra keyword arguments passed to every ``forward`` call. - batch_size : int (default: 32) - Number of series embedded at once, see - :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. - device : str or None (default: None) - Device on which inference is run, see - :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. - verbose : int (default: 0) - When positive, prints progress information. - - Attributes - ---------- - embedder_ : TimeSeriesFoundationEmbedder - The frozen feature extractor. - probe_ : sklearn estimator - The fitted head. - classes_ : array of shape=(n_classes,) - Class labels known to the classifier. - n_features_in_ : int - Number of features (channels) of the series seen during fit. - - See Also - -------- - LinearProbeForecaster: Linear probing for time series forecasting. - TimeSeriesFoundationEmbedder: The underlying feature extractor. - - Examples - -------- - >>> from chronos import Chronos2Pipeline # doctest: +SKIP - >>> pipeline = Chronos2Pipeline.from_pretrained("amazon/chronos-2") # doctest: +SKIP - >>> model = LinearProbeClassifier(pipeline.model, layer=-2) # doctest: +SKIP - >>> model.fit(X_train, y_train).score(X_test, y_test) # doctest: +SKIP - 0.93 - """ - - def __init__( - self, - model, - probe=None, - context_length=None, - layer=None, - layers_path=None, - pooling="mean", - tokens=None, - token_index=0, - input_layout="univariate", - input_name=None, - model_kwargs=None, - batch_size=32, - device=None, - verbose=0, - ): - self.model = model - self.probe = probe - self.context_length = context_length - self.layer = layer - self.layers_path = layers_path - self.pooling = pooling - self.tokens = tokens - self.token_index = token_index - self.input_layout = input_layout - self.input_name = input_name - self.model_kwargs = model_kwargs - self.batch_size = batch_size - self.device = device - self.verbose = verbose - - def _make_embedder(self): - return TimeSeriesFoundationEmbedder( - model=self.model, - layer=self.layer, - layers_path=self.layers_path, - pooling=self.pooling, - tokens=self.tokens, - token_index=self.token_index, - input_layout=self.input_layout, - context_length=self.context_length, - input_name=self.input_name, - model_kwargs=self.model_kwargs, - batch_size=self.batch_size, - device=self.device, - verbose=self.verbose, - ) - - def _make_probe(self): - if self.probe is None: - return LogisticRegression(max_iter=1000) - return clone(self.probe) - - def _check_input(self, X): - if self.input_layout not in LAYOUTS: - raise ValueError( - f"`input_layout` must be one of {LAYOUTS}, got " - f"'{self.input_layout}'." - ) - _check_probe_pooling(self.pooling) - X = check_array(X, allow_nd=True, force_all_finite=True) - return to_time_series_dataset(X) - - def fit(self, X, y): - """Fit a classifier on top of the frozen pre-trained model. - - Parameters - ---------- - X : array-like of shape=(n_ts, sz, d) - Time series dataset. - y : array-like of shape=(n_ts,) - Class labels. - - Returns - ------- - self - The fitted estimator - - """ - _require_torch() - X = self._check_input(X) - y = np.asarray(y) - check_classification_targets(y) - if len(y) != X.shape[0]: - raise ValueError( - f"X and y have inconsistent numbers of samples: {X.shape[0]} " - f"and {len(y)}." - ) - - self.embedder_ = self._make_embedder() - embeddings = self.embedder_.fit_transform(X) - - self.probe_ = self._make_probe() - self.probe_.fit(embeddings, y) - - self.classes_ = np.asarray(self.probe_.classes_) - self.n_features_in_ = X.shape[2] - return self - - def _transform(self, X): - check_is_fitted(self, "probe_") - X = self._check_input(X) - if X.shape[2] != self.n_features_in_: - raise ValueError( - f"Series with {self.n_features_in_} features were expected, got " - f"{X.shape[2]}." - ) - return self.embedder_.transform(X) - - def transform(self, X): - """Return the frozen representations the classifier operates on. - - Parameters - ---------- - X : array-like of shape=(n_ts, sz, d) - Time series dataset. - - Returns - ------- - array of shape=(n_ts, embedding_size) - Frozen representations of the input series. - - """ - return self._transform(X) - - def predict(self, X): - """Predict the class of each time series. - - Parameters - ---------- - X : array-like of shape=(n_ts, sz, d) - Time series dataset. - - Returns - ------- - array of shape=(n_ts,) - Predicted class labels. - - """ - return self.probe_.predict(self._transform(X)) - - def predict_proba(self, X): - """Predict class probabilities for each time series. - - Parameters - ---------- - X : array-like of shape=(n_ts, sz, d) - Time series dataset. - - Returns - ------- - array of shape=(n_ts, n_classes) - Predicted class probabilities, ordered as ``classes_``. - - """ - check_is_fitted(self, "probe_") - if not hasattr(self.probe_, "predict_proba"): - raise AttributeError( - f"The probe {type(self.probe_).__name__} does not expose " - "`predict_proba`." - ) - return self.probe_.predict_proba(self._transform(X)) - - def decision_function(self, X): - """Return the decision function of the linear head. - - Parameters - ---------- - X : array-like of shape=(n_ts, sz, d) - Time series dataset. - - Returns - ------- - array of shape=(n_ts,) or (n_ts, n_classes) - Confidence scores. - - """ - check_is_fitted(self, "probe_") - if not hasattr(self.probe_, "decision_function"): - raise AttributeError( - f"The probe {type(self.probe_).__name__} does not expose " - "`decision_function`." - ) - return self.probe_.decision_function(self._transform(X)) - - def __sklearn_tags__(self): - tags = super().__sklearn_tags__() - tags.input_tags.allow_nan = False - return tags diff --git a/tslearn/foundation/_embedding.py b/tslearn/foundation/_embedding.py index 21059ba8..a0b022d2 100644 --- a/tslearn/foundation/_embedding.py +++ b/tslearn/foundation/_embedding.py @@ -184,9 +184,9 @@ class TimeSeriesFoundationEmbedder(TimeSeriesMixin, TransformerMixin, BaseEstima This transformer runs a (typically pre-trained) PyTorch model in inference mode and pools its hidden states into a single vector per time series. It is the building block used by - :class:`~tslearn.foundation.LinearProbeForecaster` and - :class:`~tslearn.foundation.LinearProbeClassifier`, and can also be used on - its own, for instance inside a :class:`sklearn.pipeline.Pipeline`. + :class:`~tslearn.foundation.LinearProbeForecaster`, and can also be used on + its own, for instance inside a :class:`sklearn.pipeline.Pipeline` ahead of + a classifier, to implement linear probing for classification. The transformer is deliberately agnostic to any particular model implementation. It only assumes that: diff --git a/tslearn/foundation/_forecasting.py b/tslearn/foundation/_forecasting.py index 669d2621..53d8217b 100644 --- a/tslearn/foundation/_forecasting.py +++ b/tslearn/foundation/_forecasting.py @@ -576,7 +576,9 @@ class LinearProbeForecaster(_BaseFoundationForecaster): See Also -------- ZeroShotForecaster: Use a pre-trained model without any training. - LinearProbeClassifier: Linear probing for time series classification. + TimeSeriesFoundationEmbedder: The underlying feature extractor, which + can be composed with a classifier through a + :class:`sklearn.pipeline.Pipeline` to build a linear probe model for classification, if needed. Examples -------- From af47965698725bb14355f7a45e85185594108e01 Mon Sep 17 00:00:00 2001 From: charavelg Date: Tue, 1 Sep 2026 10:39:15 +0200 Subject: [PATCH 15/18] Try to enable action from branch --- .github/workflows/test_foundation_zoo.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/test_foundation_zoo.yml b/.github/workflows/test_foundation_zoo.yml index f168dcbe..9d3553b4 100644 --- a/.github/workflows/test_foundation_zoo.yml +++ b/.github/workflows/test_foundation_zoo.yml @@ -11,6 +11,9 @@ on: workflow_dispatch: schedule: - cron: "30 3 * * 0" # weekly, downloads are too heavy to run nightly + push: + branches: + - "foundation-module" jobs: test_foundation_zoo: From 0e4f8bc1531844c7fc832c78f30e7d644919ff77 Mon Sep 17 00:00:00 2001 From: charavelg Date: Wed, 2 Sep 2026 11:44:25 +0200 Subject: [PATCH 16/18] Code cleanup --- .../forecasting/plot_foundation_model_zoo.py | 8 +- tests/test_foundation.py | 273 ++++++++++++------ tests/test_foundation_zoo.py | 15 +- tslearn/foundation/_embedding.py | 123 +++----- tslearn/foundation/_forecasting.py | 30 +- 5 files changed, 250 insertions(+), 199 deletions(-) diff --git a/docs/examples/forecasting/plot_foundation_model_zoo.py b/docs/examples/forecasting/plot_foundation_model_zoo.py index c8f33ff5..5a99e09f 100644 --- a/docs/examples/forecasting/plot_foundation_model_zoo.py +++ b/docs/examples/forecasting/plot_foundation_model_zoo.py @@ -137,13 +137,7 @@ "amazon/chronos-bolt-small", device_map="cpu" ) -zero_shot = ZeroShotForecaster( - pipeline, - predict_fn=lambda model, context, horizon: model.predict( - torch.as_tensor(context, dtype=torch.float32), prediction_length=horizon - ), - horizon_axis=-1, -) +zero_shot = ZeroShotForecaster(pipeline) y_zero_shot = zero_shot.predict(X_train, n=horizon) forecaster = LinearProbeForecaster( diff --git a/tests/test_foundation.py b/tests/test_foundation.py index 81704b41..99547595 100644 --- a/tests/test_foundation.py +++ b/tests/test_foundation.py @@ -14,10 +14,12 @@ from sklearn.pipeline import make_pipeline from sklearn.svm import LinearSVC -from tslearn.generators import random_walks - -torch = pytest.importorskip("torch") +try: + import torch +except ImportError: + torch = None +from tslearn.generators import random_walks from tslearn.foundation import ( # noqa: E402 LinearProbeForecaster, TimeSeriesFoundationEmbedder, @@ -29,107 +31,118 @@ D_MODEL = 8 -class _Block(torch.nn.Module): - """A single, deliberately simple, encoder block.""" +@pytest.mark.skipif(torch is not None, reason="PyTorch is installed") +def test_no_torch(): + with pytest.raises(ValueError, match="torch is not installed"): + TimeSeriesFoundationEmbedder("whatever") + with pytest.raises(ValueError, match="torch is not installed"): + LinearProbeForecaster("whatever") + with pytest.raises(ValueError, match="torch is not installed"): + ZeroShotForecaster("whatever") - def __init__(self, d_model=D_MODEL): - super().__init__() - self.linear = torch.nn.Linear(d_model, d_model) - def forward(self, hidden_states): - return hidden_states + torch.tanh(self.linear(hidden_states)) +if torch: + class _Block(torch.nn.Module): + """A single, deliberately simple, encoder block.""" + def __init__(self, d_model=D_MODEL): + super().__init__() + self.linear = torch.nn.Linear(d_model, d_model) -class _Encoder(torch.nn.Module): - def __init__(self, n_layers=3, d_model=D_MODEL): - super().__init__() - self.block = torch.nn.ModuleList([_Block(d_model) for _ in range(n_layers)]) - self.final_layer_norm = torch.nn.LayerNorm(d_model) + def forward(self, hidden_states): + return hidden_states + torch.tanh(self.linear(hidden_states)) - def forward(self, hidden_states): - for block in self.block: - hidden_states = block(hidden_states) - return self.final_layer_norm(hidden_states) + class _Encoder(torch.nn.Module): + def __init__(self, n_layers=3, d_model=D_MODEL): + super().__init__() + self.block = torch.nn.ModuleList([_Block(d_model) for _ in range(n_layers)]) + self.final_layer_norm = torch.nn.LayerNorm(d_model) -class _Output: - """Stands in for a ``transformers`` ``ModelOutput``.""" + def forward(self, hidden_states): + for block in self.block: + hidden_states = block(hidden_states) + return self.final_layer_norm(hidden_states) - def __init__(self, last_hidden_state): - self.last_hidden_state = last_hidden_state + class _Output: + """Stands in for a ``transformers`` ``ModelOutput``.""" -class _DummyBackbone(torch.nn.Module): - """A patch-based encoder taking univariate series, like Chronos-2. + def __init__(self, last_hidden_state): + self.last_hidden_state = last_hidden_state - A register token is appended after the context tokens, so that the - ``pooling="token"`` code path can be exercised on a non-zero ``token_index``. - """ - def __init__(self, n_layers=3, d_model=D_MODEL, patch_size=PATCH_SIZE, seed=0): - super().__init__() - # Seeded so that results do not depend on which test ran before - torch.manual_seed(seed) - self.patch_size = patch_size - self.embedding = torch.nn.Linear(patch_size, d_model) - self.register_token = torch.nn.Parameter(torch.randn(1, 1, d_model)) - self.encoder = _Encoder(n_layers, d_model) + class _DummyBackbone(torch.nn.Module): + """A patch-based encoder taking univariate series, like Chronos-2. - def forward(self, context): - batch_size, sz = context.shape - n_patches = sz // self.patch_size - patches = context[:, : n_patches * self.patch_size] - patches = patches.reshape(batch_size, n_patches, self.patch_size) - hidden_states = self.embedding(patches) - register = self.register_token.expand(batch_size, -1, -1) - hidden_states = torch.cat([hidden_states, register], dim=1) - return _Output(self.encoder(hidden_states)) + A register token is appended after the context tokens, so that the + ``pooling="token"`` code path can be exercised on a non-zero ``token_index``. + """ + def __init__(self, n_layers=3, d_model=D_MODEL, patch_size=PATCH_SIZE, seed=0): + super().__init__() + # Seeded so that results do not depend on which test ran before + torch.manual_seed(seed) + self.patch_size = patch_size + self.embedding = torch.nn.Linear(patch_size, d_model) + self.register_token = torch.nn.Parameter(torch.randn(1, 1, d_model)) + self.encoder = _Encoder(n_layers, d_model) -class _MultivariateBackbone(torch.nn.Module): - """An encoder taking (batch, sz, d) arrays, like ``transformers`` models.""" + def forward(self, context): + batch_size, sz = context.shape + n_patches = sz // self.patch_size + patches = context[:, : n_patches * self.patch_size] + patches = patches.reshape(batch_size, n_patches, self.patch_size) + hidden_states = self.embedding(patches) + register = self.register_token.expand(batch_size, -1, -1) + hidden_states = torch.cat([hidden_states, register], dim=1) + return _Output(self.encoder(hidden_states)) - def __init__(self, n_channels, d_model=D_MODEL, seed=0): - super().__init__() - torch.manual_seed(seed) - self.embedding = torch.nn.Linear(n_channels, d_model) - self.encoder = _Encoder(2, d_model) - def forward(self, past_values): - return _Output(self.encoder(self.embedding(past_values))) + class _MultivariateBackbone(torch.nn.Module): + """An encoder taking (batch, sz, d) arrays, like ``transformers`` models.""" + def __init__(self, n_channels, d_model=D_MODEL, seed=0): + super().__init__() + torch.manual_seed(seed) + self.embedding = torch.nn.Linear(n_channels, d_model) + self.encoder = _Encoder(2, d_model) -class _DummyPipeline: - """A zero-shot forecaster mimicking ``Chronos2Pipeline``. + def forward(self, past_values): + return _Output(self.encoder(self.embedding(past_values))) - It returns one tensor per series, of shape - ``(n_variates, prediction_length, n_quantiles)``. - """ - n_quantiles = 9 + class _DummyPipeline: + """A zero-shot forecaster mimicking ``Chronos2Pipeline``. - def predict(self, inputs, prediction_length=1, batch_size=256): - inputs = np.asarray(inputs, dtype=np.float64) - # Naive forecast: repeat the last observed value - last = inputs[:, -1] - quantile_offsets = np.linspace(-1.0, 1.0, self.n_quantiles) - forecast = ( - last[:, None, None] - + np.zeros((1, prediction_length, 1)) - + quantile_offsets[None, None, :] - ) - return [torch.as_tensor(row[None]) for row in forecast] + It returns one tensor per series, of shape + ``(n_variates, prediction_length, n_quantiles)``. + """ + n_quantiles = 9 -def _dataset(n_ts=6, sz=32, d=1, seed=0): - return random_walks(n_ts=n_ts, sz=sz, d=d, random_state=seed) + def predict(self, inputs, prediction_length=1, batch_size=256): + inputs = np.asarray(inputs, dtype=np.float64) + # Naive forecast: repeat the last observed value + last = inputs[:, -1] + quantile_offsets = np.linspace(-1.0, 1.0, self.n_quantiles) + forecast = ( + last[:, None, None] + + np.zeros((1, prediction_length, 1)) + + quantile_offsets[None, None, :] + ) + return [torch.as_tensor(row[None]) for row in forecast] + + + def _dataset(n_ts=6, sz=32, d=1, seed=0): + return random_walks(n_ts=n_ts, sz=sz, d=d, random_state=seed) # --------------------------------------------------------------------------- # TimeSeriesFoundationEmbedder # --------------------------------------------------------------------------- - +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_shapes_and_layouts(): X = _dataset(n_ts=5, sz=32, d=1) embedder = TimeSeriesFoundationEmbedder(_DummyBackbone()) @@ -149,6 +162,7 @@ def test_embedder_shapes_and_layouts(): assert embedder.fit_transform(X_multi).shape == (5, D_MODEL) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_channel_stacking_is_order_preserving(): # The embedding of a multivariate series must be the concatenation of the # embeddings of its channels, in channel order. @@ -164,6 +178,7 @@ def test_embedder_channel_stacking_is_order_preserving(): ) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") @pytest.mark.parametrize("pooling", ["mean", "max", "token", "last", "flatten"]) def test_embedder_poolings(pooling): X = _dataset(n_ts=4, sz=32, d=1) @@ -175,6 +190,7 @@ def test_embedder_poolings(pooling): assert embedder.n_tokens_ is None +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") @pytest.mark.parametrize("pooling", [None, "none"]) def test_embedder_without_pooling_is_a_series_to_series_transform(pooling): X = _dataset(n_ts=4, sz=32, d=1) @@ -195,6 +211,7 @@ def test_embedder_without_pooling_is_a_series_to_series_transform(pooling): assert labels.shape == (4,) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_token_selection(): X = _dataset(n_ts=4, sz=32, d=1) n_context = 32 // PATCH_SIZE # tokens that represent the series @@ -244,6 +261,7 @@ def test_embedder_token_selection(): ) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_token_selection_multivariate(): X = _dataset(n_ts=4, sz=32, d=3) n_context = 32 // PATCH_SIZE @@ -270,6 +288,7 @@ def test_embedder_token_selection_multivariate(): ) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_token_selection_errors(): X = _dataset(n_ts=4, sz=32, d=1) with pytest.raises(ValueError, match="`tokens` must be"): @@ -280,6 +299,7 @@ def test_embedder_token_selection_errors(): TimeSeriesFoundationEmbedder(_DummyBackbone(), tokens=(0, 0)).fit(X) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_probes_reject_unpooled_representations(): X = _dataset(n_ts=4, sz=32, d=1) for pooling in (None, "none"): @@ -289,6 +309,7 @@ def test_probes_reject_unpooled_representations(): ).fit(X) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_probes_accept_token_selection(): X_fc = _dataset(n_ts=6, sz=48, d=1) model = LinearProbeForecaster( @@ -297,6 +318,7 @@ def test_probes_accept_token_selection(): assert model.predict(X_fc).shape == (6, 2, 1) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_token_index_selects_the_register_token(): X = _dataset(n_ts=3, sz=32, d=1) # The register token is appended last and does not depend on the input, so @@ -315,6 +337,7 @@ def test_embedder_token_index_selects_the_register_token(): assert not np.allclose(embeddings, embeddings[:1]) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_layer_selection(): X = _dataset(n_ts=4, sz=32, d=1) backbone = _DummyBackbone(n_layers=3) @@ -346,6 +369,7 @@ def test_embedder_layer_selection(): ) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_leaves_the_model_frozen(): X = _dataset(n_ts=4, sz=32, d=1) backbone = _DummyBackbone() @@ -356,22 +380,7 @@ def test_embedder_leaves_the_model_frozen(): torch.testing.assert_close(parameter.detach(), reference) -def test_embedder_context_length_truncation(): - X = _dataset(n_ts=4, sz=64, d=1) - embedder = TimeSeriesFoundationEmbedder(_DummyBackbone(), context_length=32) - np.testing.assert_allclose( - embedder.fit_transform(X), - TimeSeriesFoundationEmbedder(embedder.model).fit_transform(X[:, -32:]), - rtol=1e-5, - atol=1e-6, - ) - - with pytest.raises(ValueError, match="context_length"): - TimeSeriesFoundationEmbedder( - _DummyBackbone(), context_length=128 - ).fit_transform(X) - - +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_errors(): X = _dataset(n_ts=4, sz=32, d=1) with pytest.raises(ValueError, match="pooling"): @@ -390,8 +399,24 @@ def test_embedder_errors(): TimeSeriesFoundationEmbedder(_DummyBackbone()).fit(X).transform( _dataset(n_ts=4, sz=32, d=2) ) - - + with pytest.raises(RuntimeError, match="Could not determine which argument"): + class _InvalidForwardBackbone(torch.nn.Module): + def forward(self, *args, **kwargs): pass + TimeSeriesFoundationEmbedder(_InvalidForwardBackbone()).fit(X) + with pytest.raises(RuntimeError, match="forward hook was never triggered"): + class _InvalidLayerBackbone(torch.nn.Module): + def __init__(self): + super().__init__() + self.linears = torch.nn.ModuleList([ + torch.nn.Linear(32, 32), + torch.nn.Linear(D_MODEL, D_MODEL) + ]) + def forward(self, data): + self.linears[0](data) + TimeSeriesFoundationEmbedder(_InvalidLayerBackbone(), layer=1).fit(X) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_in_a_sklearn_pipeline(): X = _dataset(n_ts=20, sz=32, d=1) y = np.arange(20) % 2 @@ -402,6 +427,7 @@ def test_embedder_in_a_sklearn_pipeline(): assert clone(pipeline) is not pipeline +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_channels_first_layout(): X = _dataset(n_ts=4, sz=32, d=3) seen = {} @@ -425,6 +451,7 @@ def forward(self, series): assert embeddings.shape == (4, D_MODEL) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_layers_path_with_numeric_segment(): X = _dataset(n_ts=3, sz=32, d=1) @@ -462,6 +489,7 @@ def forward(self, context): ).fit(X) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_autodetect_layers_failure(): X = _dataset(n_ts=3, sz=32, d=1) @@ -470,13 +498,26 @@ def __init__(self): super().__init__() self.linear = torch.nn.Linear(32, D_MODEL) - def forward(self, x): - return _Output(self.linear(x).unsqueeze(1)) + def forward(self, x): pass with pytest.raises(ValueError, match="Could not automatically locate"): TimeSeriesFoundationEmbedder(_NoLayerStack(), layer=0).fit(X) + class _HeterogeneousLayersBackbone(torch.nn.Module): + def __init__(self, d_model=D_MODEL): + super().__init__() + self.bloc = torch.nn.ModuleList([ + torch.nn.Linear(d_model, d_model), + torch.nn.LayerNorm(d_model) + ]) + + def forward(self, x): pass + + with pytest.raises(ValueError, match="Could not automatically locate"): + TimeSeriesFoundationEmbedder(_HeterogeneousLayersBackbone(), layer=0).fit(X) + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_accepts_dict_model_outputs(): X = _dataset(n_ts=3, sz=32, d=1) @@ -494,6 +535,7 @@ def forward(self, x): assert embeddings.shape == (3, D_MODEL) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_accepts_hidden_states_tuple_attribute(): X = _dataset(n_ts=3, sz=32, d=1) @@ -522,6 +564,7 @@ def forward(self, x): np.testing.assert_allclose(embeddings, 2 * raw.numpy(), rtol=1e-5, atol=1e-6) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_accepts_plain_tuple_outputs(): X = _dataset(n_ts=3, sz=32, d=1) @@ -542,6 +585,7 @@ def forward(self, x): assert embeddings.shape == (3, D_MODEL) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_plain_tuple_falls_back_to_any_tensor(): X = _dataset(n_ts=3, sz=32, d=1) @@ -559,6 +603,7 @@ def forward(self, x): assert embeddings.shape == (3, D_MODEL) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_rejects_unrecognized_model_output(): X = _dataset(n_ts=3, sz=32, d=1) @@ -570,17 +615,19 @@ def forward(self, x): TimeSeriesFoundationEmbedder(_BogusOutputBackbone()).fit(X) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_rejects_invalid_hidden_state_rank(): X = _dataset(n_ts=3, sz=32, d=1) class _WeirdRankBackbone(torch.nn.Module): def forward(self, x): - return x[:, None, None, :] # 4d, not a valid hidden-state rank + return torch.tensor(x[:, None, None, :]) # 4d, not a valid hidden-state rank with pytest.raises(ValueError, match="Expected hidden states of shape"): TimeSeriesFoundationEmbedder(_WeirdRankBackbone()).fit(X) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_explicit_device(): X = _dataset(n_ts=3, sz=32, d=1) embeddings = TimeSeriesFoundationEmbedder( @@ -589,6 +636,7 @@ def test_embedder_explicit_device(): assert embeddings.shape == (3, D_MODEL) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_explicit_input_name(): X = _dataset(n_ts=3, sz=32, d=1) @@ -607,6 +655,7 @@ def forward(self, my_custom_arg): assert embeddings.shape == (3, D_MODEL) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_resolves_input_name_by_position_when_unrecognized(): X = _dataset(n_ts=3, sz=32, d=1) @@ -622,10 +671,11 @@ def forward(self, data): embedder = TimeSeriesFoundationEmbedder(_PositionalArgBackbone()) embeddings = embedder.fit_transform(X) - assert embedder._input_name_ == "data" + assert embedder._input_name == "data" assert embeddings.shape == (3, D_MODEL) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_verbose_prints_progress(capsys): X = _dataset(n_ts=5, sz=32, d=1) TimeSeriesFoundationEmbedder( @@ -635,6 +685,7 @@ def test_embedder_verbose_prints_progress(capsys): assert "Embedded 5/5 series" in captured.out +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_embedder_flatten_pooling_warns_on_length_mismatch(): X_fit = _dataset(n_ts=3, sz=32, d=1) X_other = _dataset(n_ts=3, sz=48, d=1, seed=1) @@ -650,6 +701,7 @@ def test_embedder_flatten_pooling_warns_on_length_mismatch(): # --------------------------------------------------------------------------- +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster(): X = _dataset(n_ts=5, sz=32, d=1) model = ZeroShotForecaster(_DummyPipeline()) @@ -666,6 +718,7 @@ def test_zero_shot_forecaster(): ) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_multivariate(): X = _dataset(n_ts=5, sz=32, d=3) predicted = ZeroShotForecaster(_DummyPipeline()).predict(X, n=4) @@ -676,6 +729,7 @@ def test_zero_shot_forecaster_multivariate(): ) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_quantile_selection(): X = _dataset(n_ts=3, sz=32, d=1) low = ZeroShotForecaster(_DummyPipeline(), quantile=0.1).predict(X, n=3) @@ -683,6 +737,7 @@ def test_zero_shot_forecaster_quantile_selection(): assert np.all(low < high) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_custom_predict_fn(): X = _dataset(n_ts=4, sz=32, d=1) @@ -694,6 +749,7 @@ def predict_fn(model, context, horizon): np.testing.assert_array_equal(predicted, np.zeros((4, 5, 1))) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_context_length(): X = _dataset(n_ts=4, sz=64, d=1) seen = {} @@ -708,6 +764,7 @@ def predict_fn(model, context, horizon): assert seen["shape"] == (4, 16) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_negotiates_the_context_format(): X = _dataset(n_ts=4, sz=32, d=1) @@ -747,6 +804,7 @@ def predict(self, inputs, prediction_length=1): ZeroShotForecaster(_RejectsEverything()).predict(X, n=3) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_explicit_horizon_axis(): X = _dataset(n_ts=4, sz=32, d=1) @@ -788,6 +846,7 @@ def predict(self, inputs, prediction_length=1): ZeroShotForecaster(_HorizonLast(), horizon_axis="last").predict(X, n=3) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_explicit_horizon_axis_lifts_ambiguity(): X = _dataset(n_ts=4, sz=32, d=1) @@ -808,6 +867,7 @@ def predict(self, inputs, prediction_length=1): ) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_warns_on_ambiguous_output_shape(): X = _dataset(n_ts=4, sz=32, d=1) @@ -829,6 +889,7 @@ def predict(self, inputs, prediction_length=1): assert ZeroShotForecaster(_Ambiguous()).predict(X, n=1).shape == (4, 1, 1) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_reduces_quantiles_to_a_point_forecast(): X = _dataset(n_ts=4, sz=32, d=1) @@ -853,6 +914,7 @@ def predict(self, inputs, prediction_length=1): ) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_errors(): X = _dataset(n_ts=4, sz=32, d=1) with pytest.raises(ValueError, match="`X` is required"): @@ -865,6 +927,7 @@ def test_zero_shot_forecaster_errors(): ZeroShotForecaster(_DummyPipeline(), input_layout="bogus").predict(X) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_skips_methods_without_a_horizon_argument(): X = _dataset(n_ts=4, sz=32, d=1) calls = [] @@ -887,6 +950,7 @@ def forecast(self, inputs, prediction_length=1): assert calls == ["forecast"] +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_accepts_object_attribute_output(): X = _dataset(n_ts=4, sz=32, d=1) @@ -905,6 +969,7 @@ def predict(self, inputs, prediction_length=1): assert predicted.shape == (4, 3, 1) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_univariate_row_mismatch_error(): X = _dataset(n_ts=4, sz=32, d=1) @@ -918,6 +983,7 @@ def predict(self, inputs, prediction_length=1): ZeroShotForecaster(_WrongRowCountModel()).predict(X, n=2) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_univariate_forecast_too_short_error(): X = _dataset(n_ts=4, sz=32, d=1) @@ -930,6 +996,7 @@ def predict(self, inputs, prediction_length=1): ZeroShotForecaster(_TooShortModel()).predict(X, n=3) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_accepts_dict_output(): X = _dataset(n_ts=4, sz=32, d=1) @@ -944,6 +1011,7 @@ def predict(self, inputs, prediction_length=1): assert predicted.shape == (4, 3, 1) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_empty_output_error(): X = _dataset(n_ts=4, sz=32, d=1) @@ -955,6 +1023,7 @@ def predict(self, inputs, prediction_length=1): ZeroShotForecaster(_EmptyOutputModel()).predict(X, n=3) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_heterogeneous_output_keeps_first_entry(): X = _dataset(n_ts=4, sz=32, d=1) @@ -971,6 +1040,7 @@ def predict(self, inputs, prediction_length=1): assert predicted.shape == (4, 3, 1) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_horizon_axis_prefers_smallest_longer_axis(): X = _dataset(n_ts=4, sz=32, d=1) @@ -988,6 +1058,7 @@ def predict(self, inputs, prediction_length=1): np.testing.assert_allclose(predicted, np.zeros((4, 4, 1))) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_horizon_axis_cannot_be_identified(): X = _dataset(n_ts=4, sz=32, d=1) @@ -1002,6 +1073,7 @@ def predict(self, inputs, prediction_length=1): ZeroShotForecaster(_TooShortEverywhereModel()).predict(X, n=4) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_natively_multivariate_channels_last(): X = _dataset(n_ts=4, sz=32, d=3) @@ -1020,6 +1092,7 @@ def predict(self, inputs, prediction_length=1): np.testing.assert_allclose(predicted, np.repeat(X[:, -1:], 4, axis=1)) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_natively_multivariate_channels_first(): X = _dataset(n_ts=4, sz=32, d=3) @@ -1038,6 +1111,7 @@ def predict(self, inputs, prediction_length=1): np.testing.assert_allclose(predicted, np.repeat(X[:, -1:], 4, axis=1)) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_zero_shot_forecaster_natively_multivariate_errors(): X = _dataset(n_ts=4, sz=32, d=3) @@ -1068,6 +1142,7 @@ def predict(self, inputs, prediction_length=1): # --------------------------------------------------------------------------- +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_linear_probe_forecaster(): X = _dataset(n_ts=8, sz=64, d=1) model = LinearProbeForecaster( @@ -1082,6 +1157,7 @@ def test_linear_probe_forecaster(): np.testing.assert_allclose(model.predict(X, n=2), predicted[:, :2]) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_linear_probe_forecaster_multivariate(): X = _dataset(n_ts=6, sz=48, d=2) model = LinearProbeForecaster( @@ -1091,6 +1167,7 @@ def test_linear_probe_forecaster_multivariate(): assert model.embedder_.embedding_size_ == 2 * D_MODEL +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_linear_probe_forecaster_learns_something(): # On a dataset of pure sine waves with varying phases, a probe on top of a # frozen encoder should do clearly better than forecasting the last value. @@ -1105,6 +1182,7 @@ def test_linear_probe_forecaster_learns_something(): assert model.score(X[30:]) > 0.5 +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_linear_probe_forecaster_accepts_any_probe(): X = _dataset(n_ts=6, sz=48, d=1) model = LinearProbeForecaster( @@ -1120,6 +1198,7 @@ def test_linear_probe_forecaster_accepts_any_probe(): assert not hasattr(model.probe, "coef_") +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_linear_probe_forecaster_errors(): X = _dataset(n_ts=4, sz=48, d=1) with pytest.raises(ValueError, match="context_length \\+ horizon"): @@ -1140,6 +1219,7 @@ def test_linear_probe_forecaster_errors(): model.predict(_dataset(n_ts=4, sz=48, d=3)) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_linear_probe_forecaster_fit_predict(): X = _dataset(n_ts=6, sz=48, d=1) model = LinearProbeForecaster( @@ -1148,6 +1228,7 @@ def test_linear_probe_forecaster_fit_predict(): np.testing.assert_allclose(model.fit_predict(X, n=2), model.predict(X, n=2)) +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") def test_estimators_are_clonable(): for estimator in ( TimeSeriesFoundationEmbedder(_DummyBackbone()), diff --git a/tests/test_foundation_zoo.py b/tests/test_foundation_zoo.py index fc7b390e..2eb8a575 100644 --- a/tests/test_foundation_zoo.py +++ b/tests/test_foundation_zoo.py @@ -38,7 +38,6 @@ def _data(): def test_chronos_bolt(): - torch = pytest.importorskip("torch") chronos = pytest.importorskip("chronos") from tslearn.foundation import LinearProbeForecaster, ZeroShotForecaster @@ -47,13 +46,7 @@ def test_chronos_bolt(): "amazon/chronos-bolt-small", device_map="cpu" ) - zero_shot = ZeroShotForecaster( - pipeline, - predict_fn=lambda model, context, horizon: model.predict( - torch.as_tensor(context, dtype=torch.float32), prediction_length=horizon - ), - horizon_axis=-1, - ) + zero_shot = ZeroShotForecaster(pipeline) X = _data() y_zero_shot = zero_shot.predict(X, n=HORIZON) assert y_zero_shot.shape == (N_TS, HORIZON, 1) @@ -73,7 +66,7 @@ def test_chronos_bolt(): def test_timesfm(): - torch = pytest.importorskip("torch") + pytest.importorskip("torch") timesfm = pytest.importorskip("timesfm") from tslearn.foundation import ZeroShotForecaster @@ -100,7 +93,6 @@ def test_timesfm(): def test_moirai(): - torch = pytest.importorskip("torch") pytest.importorskip("uni2ts") from uni2ts.model.moirai import MoiraiForecast, MoiraiModule @@ -133,7 +125,6 @@ def predict_fn(model, context, horizon): def test_ttm(): - torch = pytest.importorskip("torch") pytest.importorskip("tsfm_public") from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction @@ -159,7 +150,6 @@ def predict_fn(model, context, horizon): def test_moment(): - pytest.importorskip("torch") pytest.importorskip("momentfm") from momentfm import MOMENTPipeline @@ -197,7 +187,6 @@ def test_time_moe(): ) def predict_fn(model, context, horizon): - context = torch.as_tensor(context, dtype=torch.float32) out = model.generate(input_ids=context, max_new_tokens=horizon) return out[:, -horizon:] diff --git a/tslearn/foundation/_embedding.py b/tslearn/foundation/_embedding.py index a0b022d2..8e011d86 100644 --- a/tslearn/foundation/_embedding.py +++ b/tslearn/foundation/_embedding.py @@ -13,7 +13,7 @@ try: import torch -except ImportError: # pragma: no cover +except ImportError: torch = None @@ -50,17 +50,6 @@ def _normalize_pooling(pooling): return None if pooling == "none" else pooling -def _check_probe_pooling(pooling): - """Reject the poolings that do not yield a flat feature matrix.""" - if _normalize_pooling(pooling) is None: - raise ValueError( - "A linear probe needs one feature vector per series, so " - "`pooling=None` is not supported here. Use another pooling, or " - "TimeSeriesFoundationEmbedder directly if you want to keep one " - "representation per token." - ) - - def _token_slice(tokens): """Turn the ``tokens`` parameter into a slice over the token axis.""" if tokens is None: @@ -87,21 +76,12 @@ def _layout_to_model_input(X, layout): """ n_ts, sz, d = X.shape if layout == "univariate": - return np.ascontiguousarray(np.swapaxes(X, 1, 2)).reshape(n_ts * d, sz) + return torch.swapaxes(X, 1, 2).reshape(n_ts * d, sz).contiguous() if layout == "channels_first": - return np.ascontiguousarray(np.swapaxes(X, 1, 2)) + return torch.swapaxes(X, 1, 2).contiguous() if layout == "channels_last": - return np.ascontiguousarray(X) - raise ValueError(f"`input_layout` must be one of {LAYOUTS}, got '{layout}'.") - - -def _require_torch(): - if torch is None: # pragma: no cover - raise ImportError( - "PyTorch is required by the tslearn.foundation module. " - "Install it with `pip install torch`." - ) - + return X.contiguous() + raise ValueError(f"`input_layout` must be one of {LAYOUTS}, got '{layout}'.") # pragma: no cover def _resolve_attribute_path(model, path): """Return the sub-module of ``model`` designated by a dotted ``path``.""" @@ -127,28 +107,26 @@ def _autodetect_layers(model): encoder/decoder blocks are declared in essentially every implementation published on the Hugging Face Hub. """ - _require_torch() - best_path, best_modules = None, None - for path, module in model.named_modules(): + best_modules = None + for module in model.modules(): if not isinstance(module, torch.nn.ModuleList) or len(module) < 2: continue types = {type(block) for block in module} if len(types) > 1: continue if best_modules is None or len(module) > len(best_modules): - best_path, best_modules = path, module + best_modules = module if best_modules is None: raise ValueError( "Could not automatically locate the stack of layers of the provided " "model. Pass an explicit `layers_path` (e.g. 'encoder.block') to " "select the layers to probe." ) - return best_path, best_modules + return best_modules def _as_tensor(output): """Extract the hidden state tensor out of an arbitrary module output.""" - _require_torch() if isinstance(output, torch.Tensor): return output for name in CANDIDATE_HIDDEN_STATE_NAMES: @@ -251,19 +229,12 @@ class TimeSeriesFoundationEmbedder(TimeSeriesMixin, TransformerMixin, BaseEstima independently. The other two layouts feed a ``(n_ts, sz, d)`` or ``(n_ts, d, sz)`` array respectively, for models that are natively multivariate. - context_length : int or None (default: None) - When set, only the last ``context_length`` timestamps of each series - are fed to the model. input_name : str or None (default: None) Name of the ``forward`` argument receiving the context values. model_kwargs : dict or None (default: None) Extra keyword arguments passed to every ``forward`` call. batch_size : int (default: 32) - Number of series embedded at once, see - :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. - device : str or None (default: None) - Device on which inference is run, see - :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. + Number of series embedded at once. verbose : int (default: 0) When positive, prints progress information. @@ -286,9 +257,9 @@ class TimeSeriesFoundationEmbedder(TimeSeriesMixin, TransformerMixin, BaseEstima Examples -------- - >>> from chronos import Chronos2Pipeline # doctest: +SKIP - >>> pipeline = Chronos2Pipeline.from_pretrained("amazon/chronos-2") # doctest: +SKIP - >>> embedder = TimeSeriesFoundationEmbedder(pipeline.model, layer=-2) # doctest: +SKIP + >>> from chronos import Chronos2Model # doctest: +SKIP + >>> model = Chronos2Model.from_pretrained("amazon/chronos-2") # doctest: +SKIP + >>> embedder = TimeSeriesFoundationEmbedder(model, layer=-2) # doctest: +SKIP >>> embedder.fit_transform(X).shape # doctest: +SKIP (10, 512) @@ -315,13 +286,16 @@ def __init__( tokens=None, token_index=0, input_layout="univariate", - context_length=None, input_name=None, model_kwargs=None, batch_size=32, device=None, verbose=0, ): + if torch is None: + raise ValueError( + "Could not use TimeSeriesFoundationEmbedder since torch is not installed" + ) self.model = model self.layer = layer self.layers_path = layers_path @@ -329,7 +303,6 @@ def __init__( self.tokens = tokens self.token_index = token_index self.input_layout = input_layout - self.context_length = context_length self.input_name = input_name self.model_kwargs = model_kwargs self.batch_size = batch_size @@ -341,13 +314,11 @@ def _validate_params_(self): raise ValueError( f"`pooling` must be one of {POOLINGS}, got '{self.pooling}'." ) - _token_slice(self.tokens) if self.input_layout not in LAYOUTS: raise ValueError( f"`input_layout` must be one of {LAYOUTS}, got " f"'{self.input_layout}'." ) - _require_torch() if not isinstance(self.model, torch.nn.Module): raise TypeError( "`model` must be a torch.nn.Module, got " @@ -357,40 +328,43 @@ def _validate_params_(self): @property def _device(self): - if self.device is not None: - return torch.device(self.device) try: return next(self.model.parameters()).device - except StopIteration: # pragma: no cover + except StopIteration: return torch.device("cpu") + @property + def _dtype(self): + try: + return next(self.model.parameters()).dtype + except StopIteration: + return torch.get_default_dtype() + def _resolve_input_name(self): if self.input_name is not None: return self.input_name - try: - signature = inspect.signature(self.model.forward) - except (TypeError, ValueError): # pragma: no cover - return CANDIDATE_INPUT_NAMES[0] - parameters = signature.parameters + + parameters = inspect.signature(self.model.forward).parameters + # First found among candidates for name in CANDIDATE_INPUT_NAMES: if name in parameters: return name + # First found among signature parameters for name, parameter in parameters.items(): - if name == "self": - continue if parameter.kind in ( parameter.POSITIONAL_ONLY, parameter.POSITIONAL_OR_KEYWORD, ): return name - raise ValueError( # pragma: no cover + + raise RuntimeError( "Could not determine which argument of the model's `forward` method " "should receive the time series. Pass an explicit `input_name`." ) def _resolve_layers(self): if self.layers_path is not None: - return self.layers_path, _resolve_attribute_path( + return _resolve_attribute_path( self.model, self.layers_path ) return _autodetect_layers(self.model) @@ -399,14 +373,15 @@ def _forward_hidden_states(self, batch): """Run the model on a batch and return hidden states of shape (batch, n_tokens, dim).""" kwargs = dict(self.model_kwargs or {}) - kwargs[self._input_name_] = batch + if self._input_name: + kwargs[self._input_name] = batch if self.layer is None: with torch.no_grad(): output = self.model(**kwargs) return _as_tensor(output) - _, layers = self._resolve_layers() + layers = self._resolve_layers() try: layer_module = layers[self.layer] except IndexError: @@ -427,7 +402,7 @@ def hook(_module, _inputs, output): finally: handle.remove() - if "hidden_states" not in captured: # pragma: no cover + if "hidden_states" not in captured: raise RuntimeError( "The forward hook was never triggered; the selected layer does " "not seem to take part in the model's forward pass." @@ -474,14 +449,7 @@ def _pool(self, hidden_states): def _check_input(self, X): X = check_array(X, allow_nd=True, force_all_finite=True) - X = to_time_series_dataset(X) - if self.context_length is not None: - if X.shape[1] < self.context_length: - raise ValueError( - f"Series of length at least context_length=" - f"{self.context_length} are required, got {X.shape[1]}." - ) - X = X[:, -self.context_length :] + X = to_time_series_dataset(X, dtype=self._dtype, be="torch") return X def _embed(self, X): @@ -494,18 +462,16 @@ def _embed(self, X): flat = _layout_to_model_input(X, self.input_layout) embeddings = [] - device = self._device for start in range(0, flat.shape[0], self.batch_size): - chunk = flat[start : start + self.batch_size] - batch = torch.as_tensor(np.asarray(chunk, dtype=np.float32), device=device) + batch = flat[start : start + self.batch_size].to(self._device) hidden_states = self._forward_hidden_states(batch) - embeddings.append(self._pool(hidden_states).to(torch.float32).cpu().numpy()) + embeddings.append(self._pool(hidden_states).to(torch.float32)) if self.verbose: print( f"Embedded {min(start + self.batch_size, flat.shape[0])}" f"/{flat.shape[0]} series" ) - embeddings = np.concatenate(embeddings, axis=0) + embeddings = torch.concatenate(embeddings, axis=0) if self.input_layout == "univariate": if embeddings.ndim == 3: @@ -514,7 +480,7 @@ def _embed(self, X): # side by side, as in any other tslearn time series dataset. n_tokens, dim = embeddings.shape[1:] embeddings = embeddings.reshape(n_ts, d, n_tokens, dim) - embeddings = np.swapaxes(embeddings, 1, 2) + embeddings = torch.swapaxes(embeddings, 1, 2) embeddings = embeddings.reshape(n_ts, n_tokens, d * dim) else: # Concatenate the per-channel embeddings of a same series @@ -538,14 +504,18 @@ def fit(self, X, y=None): """ self._validate_params_() X = self._check_input(X) + self.model.eval() - self._input_name_ = self._resolve_input_name() + self._input_name = self._resolve_input_name() + self.n_features_in_ = X.shape[2] self._sz_fit_ = X.shape[1] + # A single series is enough to know the shape of the representations embedded = self._embed(X[:1]) self.embedding_size_ = embedded.shape[-1] self.n_tokens_ = embedded.shape[1] if embedded.ndim == 3 else None + return self def transform(self, X): @@ -599,6 +569,5 @@ def fit_transform(self, X, y=None, **fit_params): def __sklearn_tags__(self): tags = super().__sklearn_tags__() tags.target_tags.required = False - tags.non_deterministic = False tags.input_tags.allow_nan = False return tags diff --git a/tslearn/foundation/_forecasting.py b/tslearn/foundation/_forecasting.py index 53d8217b..5419f092 100644 --- a/tslearn/foundation/_forecasting.py +++ b/tslearn/foundation/_forecasting.py @@ -15,14 +15,13 @@ from ._embedding import ( LAYOUTS, TimeSeriesFoundationEmbedder, - _check_probe_pooling, _layout_to_model_input, - _require_torch, + _normalize_pooling, ) try: import torch -except ImportError: # pragma: no cover +except ImportError: torch = None @@ -111,6 +110,17 @@ def _unwrap_forecast(output): return _to_numpy(output) +def _check_probe_pooling(pooling): + """Reject the poolings that do not yield a flat feature matrix.""" + if _normalize_pooling(pooling) is None: + raise ValueError( + "A linear probe needs one feature vector per series, so " + "`pooling=None` is not supported here. Use another pooling, or " + "TimeSeriesFoundationEmbedder directly if you want to keep one " + "representation per token." + ) + + def _resolve_horizon_axis(shape, horizon, horizon_axis): """Locate the axis of a forecast array that holds the forecast horizon.""" ndim = len(shape) @@ -197,7 +207,7 @@ class _BaseFoundationForecaster(TimeSeriesMixin, BaseEstimator): def _check_input(self, X): X = check_array(X, allow_nd=True, force_all_finite=True) - return to_time_series_dataset(X) + return to_time_series_dataset(X, be="torch") def _check_layout(self): if self.input_layout not in LAYOUTS: @@ -329,6 +339,11 @@ def __init__( quantile=0.5, model_kwargs=None, ): + if torch is None: + raise ValueError( + "Could not use ZeroShotForecaster since torch is not installed" + ) + self.model = model self.predict_fn = predict_fn self.input_layout = input_layout @@ -623,6 +638,11 @@ def __init__( device=None, verbose=0, ): + if torch is None: + raise ValueError( + "Could not use LinearProbeForecaster since torch is not installed" + ) + self.model = model self.probe = probe self.context_length = context_length @@ -649,7 +669,6 @@ def _make_embedder(self): tokens=self.tokens, token_index=self.token_index, input_layout=self.input_layout, - context_length=None, input_name=self.input_name, model_kwargs=self.model_kwargs, batch_size=self.batch_size, @@ -713,7 +732,6 @@ def fit(self, X, y=None): """ self._validate_params_() - _require_torch() X = self._check_input(X) contexts, targets = self._make_windows(X) From df012808aa314dcbb5e14affb38b444f12ccee45 Mon Sep 17 00:00:00 2001 From: charavelg Date: Thu, 3 Sep 2026 10:06:25 +0200 Subject: [PATCH 17/18] Coverage --- pyproject.toml | 3 +++ tests/test_foundation.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 70642669..af9508da 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,3 +81,6 @@ testpaths = [ "tests", "tslearn" ] + +[tool.coverage.run] +omit = ["tests/test_foundation_zoo.py"] \ No newline at end of file diff --git a/tests/test_foundation.py b/tests/test_foundation.py index 99547595..84cdffd3 100644 --- a/tests/test_foundation.py +++ b/tests/test_foundation.py @@ -936,7 +936,7 @@ class _PartialModel: """Exposes a `predict` with no forecast-horizon argument, which must be skipped in favor of `forecast`.""" - def predict(self, inputs): + def predict(self, inputs): # pragma: no cover calls.append("predict") raise AssertionError("predict should never be called") From 38175181f550feb1d8fbdff4a0b7d29dbb13a6d4 Mon Sep 17 00:00:00 2001 From: charavelg Date: Thu, 3 Sep 2026 11:32:59 +0200 Subject: [PATCH 18/18] Foundation zoo updates (dtype friendly) --- .github/workflows/test_foundation_zoo.yml | 7 +- .../forecasting/plot_foundation_model_zoo.py | 20 ++--- tests/test_foundation_zoo.py | 76 +++++++++++-------- tslearn/foundation/_forecasting.py | 18 ++++- 4 files changed, 70 insertions(+), 51 deletions(-) diff --git a/.github/workflows/test_foundation_zoo.yml b/.github/workflows/test_foundation_zoo.yml index 9d3553b4..fbadf8cb 100644 --- a/.github/workflows/test_foundation_zoo.yml +++ b/.github/workflows/test_foundation_zoo.yml @@ -30,17 +30,16 @@ jobs: python -m pip install "timesfm[torch]" - name: moirai deps: | - python -m pip install "uni2ts @ git+https://github.com/SalesforceAIResearch/uni2ts.git" + python -m pip install uni2ts - name: ttm deps: | - python -m pip install "granite-tsfm[notebooks] @ git+https://github.com/ibm-granite/granite-tsfm.git@v0.2.22" - python -m pip install "transformers==4.44.2" + python -m pip install granite-tsfm - name: moment deps: | python -m pip install momentfm - name: time_moe deps: | - python -m pip install "transformers==4.40.0" + python -m pip install "transformers==4.40.1" env: TSLEARN_RUN_FOUNDATION_ZOO: "1" steps: diff --git a/docs/examples/forecasting/plot_foundation_model_zoo.py b/docs/examples/forecasting/plot_foundation_model_zoo.py index 5a99e09f..dc4e1b73 100644 --- a/docs/examples/forecasting/plot_foundation_model_zoo.py +++ b/docs/examples/forecasting/plot_foundation_model_zoo.py @@ -109,15 +109,10 @@ # # pip install "chronos-forecasting>=2.0" # -# Two details set it apart from Chronos-2: -# -# * its ``predict`` method requires a :class:`torch.Tensor`, not a NumPy -# array, so the auto-detected calling convention of -# :class:`~tslearn.foundation.ZeroShotForecaster` does not apply and an -# explicit ``predict_fn`` is needed; -# * its ``forward`` takes a raw ``context`` and exposes an encoder made of a -# plain stack of T5 blocks, at ``encoder.block``, with no register or -# forecast token to filter out, so probing needs no ``tokens`` argument. +# A details set it apart from Chronos-2: its ``forward`` takes a raw +# ``context`` and exposes an encoder made of a +# plain stack of T5 blocks, at ``encoder.block``, with no register or +# forecast token to filter out, so probing needs no ``tokens`` argument. # # As with Chronos-2, the raw series are passed to # :class:`~tslearn.foundation.ZeroShotForecaster` unscaled, since it wraps @@ -259,7 +254,7 @@ # ) # # def predict_fn(model, context, horizon): -# past_target = torch.as_tensor(context, dtype=torch.float32).unsqueeze(-1) +# past_target = context.unsqueeze(-1) # past_observed = torch.ones_like(past_target, dtype=torch.bool) # past_is_pad = torch.zeros(past_target.shape[:2], dtype=torch.bool) # return model(past_target, past_observed, past_is_pad, num_samples=20) @@ -303,7 +298,7 @@ # ) # # def predict_fn(model, context, horizon): -# past_values = torch.as_tensor(context, dtype=torch.float32).unsqueeze(-1) +# past_values = context.unsqueeze(-1) # return model(past_values=past_values).prediction_outputs # # zero_shot = ZeroShotForecaster( @@ -354,7 +349,7 @@ # scaler = TimeSeriesScalerMeanVariance(per_timeseries=False) # probe = Pipeline([("scale", scaler), ("probe", LinearProbeForecaster( # model, -# context_length=512, +# context_length=context_length, # horizon=horizon, # stride=32, # layer=-1, @@ -403,7 +398,6 @@ # .. code-block:: python # # def predict_fn(model, context, horizon): -# context = torch.as_tensor(context, dtype=torch.float32) # out = model.generate(input_ids=context, max_new_tokens=horizon) # return out[:, -horizon:] # diff --git a/tests/test_foundation_zoo.py b/tests/test_foundation_zoo.py index 2eb8a575..3098c6ee 100644 --- a/tests/test_foundation_zoo.py +++ b/tests/test_foundation_zoo.py @@ -29,15 +29,16 @@ reason="Set TSLEARN_RUN_FOUNDATION_ZOO=1 to run tests that download " "real pre-trained models from the Hugging Face Hub.", ) +torch = pytest.importorskip("torch", reason="torch not installed") -N_TS, SZ, CONTEXT_LENGTH, HORIZON = 5, 200, 64, 12 +N_TS, SZ, D, CONTEXT_LENGTH, HORIZON = 5, 200, 1, 64, 12 -def _data(): - return random_walks(n_ts=N_TS, sz=SZ, random_state=0).astype(np.float32) - - -def test_chronos_bolt(): +@pytest.mark.parametrize("data", [ + random_walks(n_ts=N_TS, sz=SZ, random_state=0).astype(np.float64), + torch.rand(N_TS, SZ, D, dtype=torch.float64), +]) +def test_chronos_bolt(data): chronos = pytest.importorskip("chronos") from tslearn.foundation import LinearProbeForecaster, ZeroShotForecaster @@ -47,8 +48,7 @@ def test_chronos_bolt(): ) zero_shot = ZeroShotForecaster(pipeline) - X = _data() - y_zero_shot = zero_shot.predict(X, n=HORIZON) + y_zero_shot = zero_shot.predict(data, n=HORIZON) assert y_zero_shot.shape == (N_TS, HORIZON, 1) probe = LinearProbeForecaster( @@ -60,13 +60,16 @@ def test_chronos_bolt(): layers_path="encoder.block", pooling="mean", ) - probe.fit(X) - y_probe = probe.predict(X) + probe.fit(data) + y_probe = probe.predict(data) assert y_probe.shape == (N_TS, HORIZON, 1) -def test_timesfm(): - pytest.importorskip("torch") +@pytest.mark.parametrize("data", [ + random_walks(n_ts=N_TS, sz=SZ, random_state=0), + torch.rand(N_TS, SZ, D, dtype=torch.float64), +]) +def test_timesfm(data): timesfm = pytest.importorskip("timesfm") from tslearn.foundation import ZeroShotForecaster @@ -87,12 +90,15 @@ def test_timesfm(): )[0], context_length=CONTEXT_LENGTH, ) - X = _data() - y_zero_shot = zero_shot.predict(X, n=HORIZON) + y_zero_shot = zero_shot.predict(data, n=HORIZON) assert y_zero_shot.shape == (N_TS, HORIZON, 1) -def test_moirai(): +@pytest.mark.parametrize("data", [ + random_walks(n_ts=N_TS, sz=SZ, random_state=0), + torch.rand(N_TS, SZ, D, dtype=torch.float64), +]) +def test_moirai(data): pytest.importorskip("uni2ts") from uni2ts.model.moirai import MoiraiForecast, MoiraiModule @@ -110,7 +116,7 @@ def test_moirai(): ) def predict_fn(model, context, horizon): - past_target = torch.as_tensor(context, dtype=torch.float32).unsqueeze(-1) + past_target = context.unsqueeze(-1) past_observed = torch.ones_like(past_target, dtype=torch.bool) past_is_pad = torch.zeros(past_target.shape[:2], dtype=torch.bool) with torch.no_grad(): @@ -119,12 +125,15 @@ def predict_fn(model, context, horizon): zero_shot = ZeroShotForecaster( forecast_model, predict_fn=predict_fn, context_length=CONTEXT_LENGTH ) - X = _data() - y_zero_shot = zero_shot.predict(X, n=HORIZON) + y_zero_shot = zero_shot.predict(data, n=HORIZON) assert y_zero_shot.shape == (N_TS, HORIZON, 1) -def test_ttm(): +@pytest.mark.parametrize("data", [ + random_walks(n_ts=N_TS, sz=SZ, random_state=0), + torch.rand(N_TS, SZ, D, dtype=torch.float64), +]) +def test_ttm(data): pytest.importorskip("tsfm_public") from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction @@ -135,7 +144,7 @@ def test_ttm(): ) def predict_fn(model, context, horizon): - past_values = torch.as_tensor(context, dtype=torch.float32).unsqueeze(-1) + past_values = context.unsqueeze(-1) with torch.no_grad(): return model(past_values=past_values).prediction_outputs @@ -149,7 +158,11 @@ def predict_fn(model, context, horizon): assert y_zero_shot.shape == (N_TS, model.config.prediction_length, 1) -def test_moment(): +@pytest.mark.parametrize("data", [ + random_walks(n_ts=N_TS, sz=SZ, random_state=0), + torch.rand(N_TS, SZ, D, dtype=torch.float64), +]) +def test_moment(data): pytest.importorskip("momentfm") from momentfm import MOMENTPipeline @@ -162,7 +175,7 @@ def test_moment(): probe = LinearProbeForecaster( model, - context_length=512, + context_length=CONTEXT_LENGTH, horizon=HORIZON, stride=32, layer=-1, @@ -170,14 +183,16 @@ def test_moment(): pooling="mean", input_layout="channels_first", ) - X = random_walks(n_ts=N_TS, sz=700, random_state=0).astype(np.float32) - probe.fit(X) - y_probe = probe.predict(X) + probe.fit(data) + y_probe = probe.predict(data) assert y_probe.shape == (N_TS, HORIZON, 1) -def test_time_moe(): - torch = pytest.importorskip("torch") +@pytest.mark.parametrize("data", [ + random_walks(n_ts=N_TS, sz=SZ, random_state=0), + torch.rand(N_TS, SZ, D, dtype=torch.float64), +]) +def test_time_moe(data): transformers = pytest.importorskip("transformers") from tslearn.foundation import LinearProbeForecaster, ZeroShotForecaster @@ -191,8 +206,7 @@ def predict_fn(model, context, horizon): return out[:, -horizon:] zero_shot = ZeroShotForecaster(model, predict_fn=predict_fn) - X = _data() - y_zero_shot = zero_shot.predict(X, n=HORIZON) + y_zero_shot = zero_shot.predict(data, n=HORIZON) assert y_zero_shot.shape == (N_TS, HORIZON, 1) probe = LinearProbeForecaster( @@ -204,6 +218,6 @@ def predict_fn(model, context, horizon): layers_path="model.layers", pooling="last", ) - probe.fit(X) - y_probe = probe.predict(X) + probe.fit(data) + y_probe = probe.predict(data) assert y_probe.shape == (N_TS, HORIZON, 1) diff --git a/tslearn/foundation/_forecasting.py b/tslearn/foundation/_forecasting.py index 5419f092..7267639e 100644 --- a/tslearn/foundation/_forecasting.py +++ b/tslearn/foundation/_forecasting.py @@ -207,7 +207,7 @@ class _BaseFoundationForecaster(TimeSeriesMixin, BaseEstimator): def _check_input(self, X): X = check_array(X, allow_nd=True, force_all_finite=True) - return to_time_series_dataset(X, be="torch") + return to_time_series_dataset(X, self._dtype, be="torch") def _check_layout(self): if self.input_layout not in LAYOUTS: @@ -352,6 +352,13 @@ def __init__( self.quantile = quantile self.model_kwargs = model_kwargs + @property + def _dtype(self): + try: + return next(self.model.parameters()).dtype + except (StopIteration, AttributeError): + return torch.get_default_dtype() + def _resolve_predict_fn(self): """Build a ``(context, horizon) -> forecast`` callable from the model.""" if self.predict_fn is not None: @@ -436,8 +443,6 @@ def fit(self, X, y=None): The estimator, ready to be used for prediction """ - self._check_layout() - X = self._check_input(X) warnings.warn( "ZeroShotForecaster.fit does not train anything: the wrapped " "model is used as-is for zero-shot forecasting.", @@ -660,6 +665,13 @@ def __init__( self.device = device self.verbose = verbose + @property + def _dtype(self): + try: + return next(self.model.parameters()).dtype + except StopIteration: + return torch.get_default_dtype() + def _make_embedder(self): return TimeSeriesFoundationEmbedder( model=self.model,