diff --git a/benchmark_utils/base_solver.py b/benchmark_utils/base_solver.py index cd82ebd..10755db 100644 --- a/benchmark_utils/base_solver.py +++ b/benchmark_utils/base_solver.py @@ -6,12 +6,238 @@ import numpy as np import torch from benchopt import BaseSolver +from sklearn.linear_model import LogisticRegression, Ridge +from benchmark_utils.adapters.base import BaseTSFMAdapter +from benchmark_utils.adapters.forecast_residual import ForecastResidualAdapter +from benchmark_utils.adapters.linear_probe import LinearProbeAdapter from benchmark_utils.covariates import Covariates from benchmark_utils.inputs import ForecastInput from benchmark_utils.outputs import ForecastOutput -TaskType = Literal["forecasting", "classification", "anomaly_detection"] +TaskType = Literal[ + "forecasting", "classification", "anomaly_detection", "event_detection" +] + + +# --------------------------------------------------------------------------- +# Private adapter / encoder helpers used by build_adapter defaults +# --------------------------------------------------------------------------- + + +class _SolverForecastAdapter(BaseTSFMAdapter): + """Wraps BaseTSFMSolver.forecast() as a BaseTSFMAdapter.""" + + def __init__(self, solver: "BaseTSFMSolver") -> None: + self.solver = solver + + def predict( + self, x: ForecastInput, prediction_length: int | None = None + ) -> ForecastOutput: + horizon = prediction_length or self.solver.meta.get("prediction_length", 1) + return self.solver.forecast(x, horizon, self.solver.quantile_levels) + + +class _SolverEmbedEncoder: + """Wraps BaseTSFMSolver.embed() as a flat encoder for LinearProbeAdapter.""" + + def __init__(self, solver: "BaseTSFMSolver") -> None: + self.solver = solver + + def encode(self, X: np.ndarray | list) -> np.ndarray: + # X: list of (T_i, C) / (T_i,), or array (B, T, C) / (T, C) / (T,) + if isinstance(X, list): + return self.solver.embed(X) + X = np.asarray(X, dtype=np.float32) + if X.ndim <= 2: + return self.solver.embed([X]) # single series + return self.solver.embed(list(X)) # batch + + +class _SolverTimeEmbedPooledEncoder: + """Wraps BaseTSFMSolver.time_embed() with mean pooling for LinearProbeAdapter.""" + + def __init__(self, solver: "BaseTSFMSolver") -> None: + self.solver = solver + + def encode(self, X: np.ndarray | list) -> np.ndarray: + # X: list of (T_i, C) / (T_i,), or array (B, T, C) / (T, C) / (T,) + if isinstance(X, list): + series_list = X + else: + X = np.asarray(X, dtype=np.float32) + series_list = [X] if X.ndim <= 2 else list(X) + time_embs = self.solver.time_embed(series_list) # list of (T'_i, D) + return np.stack([emb.mean(axis=0) for emb in time_embs], axis=0) # (B, D) + + +class _WindowedForecastAdapter(BaseTSFMAdapter): + """Point forecast via embed on sliding windows + ridge regression. + + Builds (window_embedding → next_H_values) training pairs from the + training series, then at inference embeds the last ``window_size`` + timesteps before each cutoff and predicts the next ``prediction_length`` + values. Always outputs a single quantile at 0.5 (point forecast). + """ + + def __init__( + self, solver: "BaseTSFMSolver", window_size: int, prediction_length: int + ) -> None: + self.solver = solver + self.window_size = window_size + self.prediction_length = prediction_length + self._head: Ridge | None = None + + def fit( + self, X_train: Sequence[np.ndarray], y_train: Any = None + ) -> "_WindowedForecastAdapter": + windows, targets = [], [] + for series in X_train: + series = np.asarray(series, dtype=np.float32) + if series.ndim == 1: + series = series[:, None] + T, C = series.shape + for t in range(self.window_size, T - self.prediction_length + 1): + windows.append(series[t - self.window_size: t]) + targets.append(series[t: t + self.prediction_length].flatten()) + + if not windows: + return self + embs = self.solver.embed(windows) # (N, D) + self._head = Ridge().fit(embs, np.stack(targets)) # targets: (N, H*C) + return self + + def predict( + self, x: ForecastInput, prediction_length: int | None = None + ) -> ForecastOutput: + # Ignore prediction_length override — trained for a fixed horizon. + H = self.prediction_length + windows, layout, per_series_shape = [], [], [] + + for series_idx, (series, cutoffs) in enumerate(zip(x.x, x.cutoff_indexes)): + series = np.asarray(series, dtype=np.float32) + if series.ndim == 1: + series = series[:, None] + _, C = series.shape + per_series_shape.append((C, len(cutoffs))) + for cutoff_idx, cutoff in enumerate(cutoffs): + hist = series[:cutoff] + if len(hist) >= self.window_size: + window = hist[-self.window_size:] + else: + pad = np.zeros( + (self.window_size - len(hist), hist.shape[1]), dtype=np.float32 + ) + window = np.concatenate([pad, hist], axis=0) + windows.append(window) + layout.append((series_idx, cutoff_idx)) + + if not windows or self._head is None: + return ForecastOutput(quantiles=[], quantile_levels=(0.5,)) + + embs = self.solver.embed(windows) # (N, D) + preds = self._head.predict(embs) # (N, H*C) + + per_series = [ + np.empty((n_cutoffs, H, C, 1), dtype=np.float32) + for C, n_cutoffs in per_series_shape + ] + for i, (series_idx, cutoff_idx) in enumerate(layout): + C = per_series_shape[series_idx][0] + per_series[series_idx][cutoff_idx, :, :, 0] = preds[i].reshape(H, C) + + return ForecastOutput(quantiles=per_series, quantile_levels=(0.5,)) + + +class _TimeEmbedEventAdapter(BaseTSFMAdapter): + """Event detection via temporal embeddings + per-position logistic regression. + + The temporal embedding may be at a coarser stride than the original + series; it is resampled back to the original length T via nearest- + neighbour indexing before fitting and scoring. + """ + + def __init__(self, solver: "BaseTSFMSolver") -> None: + self.solver = solver + + def _align(self, emb: np.ndarray, T: int) -> np.ndarray: + """Resample emb (T', D) to length T via nearest-neighbour indices.""" + T_prime = emb.shape[0] + if T_prime == T: + return emb + idx = np.round(np.linspace(0, T_prime - 1, T)).astype(int) + return emb[idx] + + def fit( + self, X_train: Sequence[np.ndarray], y_train: Sequence[np.ndarray] + ) -> "_TimeEmbedEventAdapter": + series_list = [np.asarray(s, dtype=np.float32) for s in X_train] + time_embs = self.solver.time_embed(series_list) # list of (T'_i, D) + embs_all, labels_all = [], [] + for emb, labels in zip(time_embs, y_train): + T = len(labels) + embs_all.append(self._align(emb, T)) + labels_all.append(np.asarray(labels)) + X = np.concatenate(embs_all, axis=0) # (sum T_i, D) + y = np.concatenate(labels_all, axis=0) # (sum T_i,) + self._head = LogisticRegression(max_iter=1000).fit(X, y) + return self + + def predict(self, x: np.ndarray) -> np.ndarray: + # x: (T, C) → scores: (T,) probabilities of the positive class + x = np.asarray(x, dtype=np.float32) + T = x.shape[0] + emb = self.solver.time_embed([x])[0] # (T', D) + return self._head.predict_proba(self._align(emb, T))[:, 1] + + +class _WindowedEventAdapter(BaseTSFMAdapter): + """Event detection via causal windowed embedding + per-position logistic regression. + + At each timestep, a window of size ``window_size`` ending at that + position is embedded. Positions near the start of the series are zero- + padded so that every timestep receives a score. + """ + + def __init__(self, solver: "BaseTSFMSolver", window_size: int) -> None: + self.solver = solver + self.window_size = window_size + + def _causal_windows(self, series: np.ndarray) -> list[np.ndarray]: + """One zero-padded causal window per timestep.""" + T, C = series.shape + padded = np.zeros((self.window_size - 1 + T, C), dtype=np.float32) + padded[self.window_size - 1:] = series + return [padded[t: t + self.window_size] for t in range(T)] + + def fit( + self, X_train: Sequence[np.ndarray], y_train: Sequence[np.ndarray] + ) -> "_WindowedEventAdapter": + all_windows, all_labels = [], [] + for series, labels in zip(X_train, y_train): + series = np.asarray(series, dtype=np.float32) + if series.ndim == 1: + series = series[:, None] + for window, label in zip(self._causal_windows(series), labels): + all_windows.append(window) + all_labels.append(label) + + embs = self.solver.embed(all_windows) # (N, D) + self._head = LogisticRegression(max_iter=1000).fit(embs, np.array(all_labels)) + return self + + def predict(self, x: np.ndarray) -> np.ndarray: + # x: (T, C) → scores: (T,) probabilities of the positive class + x = np.asarray(x, dtype=np.float32) + if x.ndim == 1: + x = x[:, None] + embs = self.solver.embed(self._causal_windows(x)) # (T, D) + return self._head.predict_proba(embs)[:, 1] # (T,) + + +# --------------------------------------------------------------------------- +# Base solver +# --------------------------------------------------------------------------- class BaseTSFMSolver(BaseSolver): @@ -26,14 +252,14 @@ class BaseTSFMSolver(BaseSolver): Subclasses only need to implement: - supported_tasks: set of task names the model supports + - model_id: unique string identifying the current model variant - load_model(): load/initialize the model for the given device - - build_adapter(): create task-specific adapters + - At least one of forecast_batch / embed_batch / time_embed_batch Attributes ---------- supported_tasks - Subset of {"forecasting", "classification", "anomaly_detection"} - that this solver supports. Must be set by subclass. + Subset of TaskType that this solver supports. Must be set by subclass. task Current task being solved (set in set_objective). @@ -45,7 +271,7 @@ class BaseTSFMSolver(BaseSolver): Task metadata like prediction_length, n_classes, etc. model - The loaded TSFM model (cached across multiple set_objective calls). + The loaded TSFM model (cached by model_id across set_objective calls). device "cuda" or "cpu", automatically selected in set_objective. @@ -67,36 +293,28 @@ class BaseTSFMSolver(BaseSolver): dtype: str | torch.dtype def __init__(self, **kwargs: Any) -> None: - """Initialize solver with model-specific setup. - - Subclasses can override this method to perform model-specific - initialization. If overriding, call super().__init__(**kwargs). - """ super().__init__() - - # Initialize cached model state - self._loaded_model = None - self.model = None + self._loaded_model_id: str | None = None + self.model: Any = None for key, value in kwargs.items(): setattr(self, key, value) @property @abstractmethod def supported_tasks(self) -> set[TaskType]: - """Return a set of supported task names. + """Return a set of supported task names.""" - Returns - ------- - set of str - Subset of {"forecasting", "classification", "anomaly_detection"} - """ + @property + @abstractmethod + def model_id(self) -> str: + """A unique string identifying the current model variant.""" @abstractmethod def load_model(self, device: str | torch.device, dtype: torch.dtype) -> Any: """Load and return the TSFM model. - Called once per model variant (cached by subclass if needed). - This method is called inside set_objective and is NOT timed. + Called once per model_id (cached by base class). This method is + called inside set_objective and is NOT timed. Parameters ---------- @@ -111,12 +329,28 @@ def load_model(self, device: str | torch.device, dtype: torch.dtype) -> Any: The loaded TSFM model """ - @abstractmethod - def build_adapter(self, task: TaskType, model: Any) -> Any: + def build_adapter(self, task: TaskType, model: Any) -> BaseTSFMAdapter: """Create and optionally fit a task-specific adapter. - Called from run() for the current task and model. - If the adapter requires fitting (e.g., LinearProbe), do it here. + Default strategies — the first capability the solver implements is used: + + forecasting + 1. forecast_batch (zero-shot, via _SolverForecastAdapter) + 2. embed_batch (windowed ridge regression, via _WindowedForecastAdapter) + classification + 1. embed_batch (flat embedding + LinearProbeAdapter) + 2. time_embed_batch (mean-pooled temporal embedding + LinearProbeAdapter) + anomaly_detection + 1. embed_batch (distance-from-mean score, via LinearProbeAdapter) + 2. forecast_batch (forecast-error score, via ForecastResidualAdapter) + event_detection + 1. time_embed_batch (per-position LogReg, via _TimeEmbedEventAdapter) + 2. embed_batch (causal-windowed LogReg, via _WindowedEventAdapter) + + Override this method if the model requires custom adapter logic. + + # TODO: move this strategy documentation to a "how to add a model" docs + # page; also add a pointer from AGENTS.md for coding agents. Parameters ---------- @@ -130,8 +364,76 @@ def build_adapter(self, task: TaskType, model: Any) -> Any: adapter : BaseTSFMAdapter A fitted (or zero-shot) adapter ready for prediction """ - - def skip(self, task: TaskType, **_) -> tuple[bool, str | None]: + pred_len = self.meta.get("prediction_length", 1) + window_size = getattr(self, "window_size", max(pred_len * 2, 64)) + + match task: + case "forecasting": + if self.can_forecast: + return _SolverForecastAdapter(self) + if self.can_embed: + adapter = _WindowedForecastAdapter(self, window_size, pred_len) + adapter.fit(self.X_train) + return adapter + raise NotImplementedError( + f"{self.name} must implement forecast_batch or embed_batch " + "for task='forecasting'" + ) + + case "classification": + if self.can_embed: + encoder = _SolverEmbedEncoder(self) + elif self.can_time_embed: + encoder = _SolverTimeEmbedPooledEncoder(self) + else: + raise NotImplementedError( + f"{self.name} must implement embed_batch or time_embed_batch " + "for task='classification'" + ) + adapter = LinearProbeAdapter( + encoder, + task="classification", + n_classes=self.meta.get("n_classes"), + classifier=getattr(self, "classifier", "log_reg"), + penalty=getattr(self, "penalty", "l2"), + C=getattr(self, "C", 1.0), + alpha=getattr(self, "alpha", 1.0), + n_estimators=getattr(self, "n_estimators", 100), + ) + adapter.fit(self.X_train, self.y_train) + return adapter + + case "anomaly_detection": + if self.can_embed: + encoder = _SolverEmbedEncoder(self) + adapter = LinearProbeAdapter(encoder, task="anomaly_detection") + adapter.fit(self.X_train, self.y_train) + return adapter + if self.can_forecast: + return ForecastResidualAdapter(_SolverForecastAdapter(self)) + raise NotImplementedError( + f"{self.name} must implement embed_batch or forecast_batch " + "for task='anomaly_detection'" + ) + + case "event_detection": + if self.can_time_embed: + adapter = _TimeEmbedEventAdapter(self) + adapter.fit(self.X_train, self.y_train) + return adapter + if self.can_embed: + adapter = _WindowedEventAdapter(self, window_size) + adapter.fit(self.X_train, self.y_train) + return adapter + raise NotImplementedError( + f"{self.name} must implement time_embed_batch or embed_batch " + "for task='event_detection'" + ) + + case _: + raise NotImplementedError(f"Unknown task: {task!r}") + + def skip(self, task: TaskType, **_: Any) -> tuple[bool, str | None]: """Skip unsupported tasks.""" if task not in self.supported_tasks: return True, f"{self.name} solver does not support task={task!r}" @@ -166,41 +468,62 @@ def set_objective( self.y_train = y_train self.meta = meta - # Select device and dtype # bfloat16 is well-supported on CUDA but poorly on CPU/MPS; # fall back to float32 elsewhere. self.device = "cuda" if torch.cuda.is_available() else "cpu" self.dtype = torch.bfloat16 if self.device == "cuda" else torch.float32 - # Load model (caching is subclass responsibility via load_model) - self.model = self.load_model(self.device, self.dtype) + # Load model, caching by model_id across set_objective calls. + current_id = self.model_id + if self._loaded_model_id != current_id: + self.model = self.load_model(self.device, self.dtype) + self._loaded_model_id = current_id def run(self, _: Any) -> None: - """Build and fit task-specific adapter. - - Calls build_adapter() to create the adapter for the current task. - Subclasses can override to add custom logic, but should typically - just call super().run(_) to set up self._adapter. - """ + """Build and fit task-specific adapter.""" self._adapter = self.build_adapter(self.task, self.model) def get_result(self) -> dict[str, Any]: """Return the fitted adapter.""" return {"model": self._adapter} + @property + def can_forecast(self) -> bool: + # A model is supposed to have forecast capabilities iff + # it overrides the `forecast_batch` method + return type(self).forecast_batch is not BaseTSFMSolver.forecast_batch + + @property + def quantile_levels(self) -> tuple[float, ...]: + """The quantile levels to use for forecasting. + + Must be overridden by subclasses that support forecasting. Implement + this to expose the model's native quantile levels + (e.g. ``tuple(pipeline.quantiles)``). If looking this up is + expensive (e.g. it requires the loaded model), override as a + ``functools.cached_property`` instead. + """ + raise NotImplementedError( + f"{self.name} must implement quantile_levels " + "to support forecasting" + ) + def forecast_batch( - self, inputs: list[torch.Tensor], covariates: Sequence[Covariates] + self, + inputs: list[torch.Tensor], + covariates: Sequence[Covariates], + prediction_length: int, ) -> list[torch.Tensor]: """Forecast on a batch of prepared inputs. - Subclasses must implement this to call their model's inference. - Parameters ---------- inputs Prepared input tensors of shape: (lookback, channel) covariates Corresponding covariates for each input + prediction_length + Number of future steps to forecast Returns ------- @@ -260,7 +583,7 @@ def forecast( # TODO We still do this in batches in case data is very large # Get a list of model outputs aligned with inputs - raw = self.forecast_batch(inputs, covariates) + raw = self.forecast_batch(inputs, covariates, prediction_length) per_series_preds = [ [None] * n_cutoffs for _, n_cutoffs in per_series_shape @@ -271,3 +594,98 @@ def forecast( per_series = [np.stack(preds) for preds in per_series_preds] return ForecastOutput(quantiles=per_series, quantile_levels=quantile_levels) + + @property + def can_embed(self) -> bool: + # A model is supposed to have static embed capabilities iff + # it overrides the `embed_batch` method + return type(self).embed_batch is not BaseTSFMSolver.embed_batch + + def embed_batch(self, inputs: list[torch.Tensor]) -> list[torch.Tensor]: + """Compute static embeddings on a batch of series. + + Parameters + ---------- + inputs : list of torch.Tensor, shape (T, C) + One tensor per series (full series, not windowed). + + Returns + ------- + list of torch.Tensor, shape (D,) + One flat embedding vector per input series. + """ + raise NotImplementedError( + "Subclasses must implement embed_batch to call their model" + ) + + def embed(self, x: list[np.ndarray]) -> np.ndarray: + """Compute static embeddings for a list of series. + + Calls embed_batch() and stacks results into a 2-D array. + + Parameters + ---------- + x : list of np.ndarray, shape (T_i, C) + Input series (variable length). + + Returns + ------- + np.ndarray, shape (N, D) + One flat embedding per input series. + """ + inputs = [] + for series in x: + arr = np.asarray(series, dtype=np.float32) + if arr.ndim == 1: + arr = arr[:, None] # (T,) → (T, 1) + inputs.append(torch.from_numpy(arr)) + results = self.embed_batch(inputs) + return np.stack([r.float().cpu().numpy() for r in results], axis=0) + + @property + def can_time_embed(self) -> bool: + # A model is supposed to have temporal embed capabilities iff + # it overrides the `time_embed_batch` method + return type(self).time_embed_batch is not BaseTSFMSolver.time_embed_batch + + def time_embed_batch(self, inputs: list[torch.Tensor]) -> list[torch.Tensor]: + """Compute temporal embeddings on a batch of series. + + Parameters + ---------- + inputs : list of torch.Tensor, shape (T, C) + One tensor per series (full series, not windowed). + + Returns + ------- + list of torch.Tensor, shape (T', D) + One temporal embedding per input series; T' is model-determined + (depends on stride and windowing). + """ + raise NotImplementedError( + "Subclasses must implement time_embed_batch to call their model" + ) + + def time_embed(self, x: list[np.ndarray]) -> list[np.ndarray]: + """Compute temporal embeddings for a list of series. + + Calls time_embed_batch() and converts results to numpy. + + Parameters + ---------- + x : list of np.ndarray, shape (T_i, C) + Input series (variable length). + + Returns + ------- + list of np.ndarray, shape (T'_i, D) + Temporal embeddings; T'_i is model-determined per series. + """ + inputs = [] + for series in x: + arr = np.asarray(series, dtype=np.float32) + if arr.ndim == 1: + arr = arr[:, None] # (T,) → (T, 1) + inputs.append(torch.from_numpy(arr)) + results = self.time_embed_batch(inputs) + return [r.float().cpu().numpy() for r in results] diff --git a/solvers/chronos.py b/solvers/chronos.py index e0da7ef..49bc03f 100644 --- a/solvers/chronos.py +++ b/solvers/chronos.py @@ -1,15 +1,15 @@ """Chronos solver for the TSFM benchmark (local inference). -Supports: - - forecasting : zero-shot via ChronosPipeline - - classification : linear probe on pooled encoder embeddings - - anomaly_detection : forecast-residual on top of the same forecaster +Directly supports: + - forecasting : zero-shot via ``ChronosPipeline.predict`` (Monte-Carlo sampling) + +Additionally provides: + - ``embed_batch``: pooled encoder embeddings for classification and + anomaly detection via the default adaptation strategies. Model loading is done in ``set_objective`` (untimed). Inference batches -every (series, cutoff) pair into a single ``ChronosPipeline.predict`` -call — the pipeline accepts a list of variable-length tensors and -applies left-padding internally, so all the per-cutoff work happens in -one forward pass. +every (series, cutoff, channel) tuple into chunked ``ChronosPipeline.predict`` +calls — Chronos v1 is univariate so each channel is processed independently. References ---------- @@ -18,255 +18,35 @@ import numpy as np import torch -from benchopt import BaseSolver from chronos import ChronosPipeline -from benchmark_utils.adapters import ( - POOLERS, - Encoder, - LinearProbeAdapter, - UnpooledEncoder, -) -from benchmark_utils.adapters.base import BaseTSFMAdapter -from benchmark_utils.adapters.forecast_residual import ForecastResidualAdapter -from benchmark_utils.inputs import ForecastInput -from benchmark_utils.outputs import ForecastOutput - -SUPPORTED_TASKS = {"forecasting", "classification", "anomaly_detection"} - - -class _ChronosForecaster(BaseTSFMAdapter): - """Batched Chronos v1 adapter; quantiles are derived from sample draws.""" - - DEFAULT_QUANTILE_LEVELS = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9) - - # Cap how many (series, cutoff, channel) histories go through one - # ``pipeline.predict`` call. Chronos v1 left-pads the whole list into a - # single T5 forward whose attention is O(L^2) in the (padded) history - # length, so batching every series at once blows up memory on datasets - # with many long series (e.g. m4_weekly). Chunking bounds peak memory. - PREDICT_BATCH_SIZE = 16 - - def __init__(self, pipeline, prediction_length, quantile_levels=None): - self.pipeline = pipeline - self.prediction_length = prediction_length - self.quantile_levels = quantile_levels or self.DEFAULT_QUANTILE_LEVELS - - # ------------------------------------------------------------------ - # Template method — subclasses override _build_inputs / _assemble - # ------------------------------------------------------------------ - - def predict(self, x: ForecastInput, prediction_length=None) -> ForecastOutput: - horizon = prediction_length or self.prediction_length - inputs, layout, per_series_shape = self._build_inputs(x) - if not inputs: - return ForecastOutput(quantiles=[], quantile_levels=self.quantile_levels) - - # Chronos v1 forecasts by Monte-Carlo sampling, so it is - # non-deterministic by default. Seed before sampling so the same - # history yields the same draws — required for the behavioural - # leakage probe (benchmark_utils.leakage), which compares two - # predict() calls on identical history and would otherwise read - # sampling noise as a leak. Seeding once (before the chunk loop) - # keeps the draw sequence deterministic across calls because the - # inputs — and thus the chunking — are identical between calls. - torch.manual_seed(0) - chunks = [] - with torch.no_grad(): - for start in range(0, len(inputs), self.PREDICT_BATCH_SIZE): - batch = inputs[start:start + self.PREDICT_BATCH_SIZE] - chunks.append( - self.pipeline.predict(batch, prediction_length=horizon) - ) - output = torch.cat(chunks, dim=0) # (n_inputs, num_samples, H) - return self._assemble_output(output, layout, per_series_shape, horizon) - - def _build_inputs(self, x): - """Build list of 1-D tensors (one per channel) and track layout.""" - inputs = [] - layout = [] # (series_idx, cutoff_idx, channel_idx) - per_series_shape = [] # (C, n_cutoffs) - for series_idx, (series, cutoffs) in enumerate(zip(x.x, x.cutoff_indexes)): - series = np.asarray(series, dtype=np.float32) - if series.ndim == 1: - series = series[:, None] - _, C = series.shape - per_series_shape.append((C, len(cutoffs))) - for cutoff_idx, cutoff in enumerate(cutoffs): - hist = series[:cutoff] - for c in range(C): - inputs.append(torch.from_numpy(hist[:, c])) - layout.append((series_idx, cutoff_idx, c)) - return inputs, layout, per_series_shape - - def _assemble_output(self, samples, layout, per_series_shape, prediction_length): - """Derive quantile fan from Monte-Carlo sample draws.""" - # samples: (n_inputs, num_samples, H) - q_arr = np.quantile( - samples.float().cpu().numpy(), - q=list(self.quantile_levels), - axis=1, - ).transpose(1, 0, 2) # (n_inputs, Q, H) - - Q = len(self.quantile_levels) - per_series = [ - np.empty((n_cutoffs, prediction_length, C, Q), dtype=np.float32) - for C, n_cutoffs in per_series_shape - ] - for i, (series_idx, cutoff_idx, c) in enumerate(layout): - # q_arr[i] is (Q, H); store as (H, Q) at channel c. - per_series[series_idx][cutoff_idx, :, c, :] = q_arr[i].T +from benchmark_utils.adapters import POOLERS +from benchmark_utils.base_solver import BaseTSFMSolver - return ForecastOutput( - quantiles=per_series, quantile_levels=self.quantile_levels - ) - - -class _ChronosEmbedEncoder(UnpooledEncoder): - """Default path — uses ``ChronosPipeline.embed``. - - Returns hidden states *after* ``encoder.final_layer_norm``. - """ - - def __init__(self, pipeline: ChronosPipeline): - self.pipeline = pipeline - - def encode(self, X) -> np.ndarray: - # X: (B, T, V) or (T, V). - X = np.asarray(X, dtype=np.float32) - batched = X.ndim == 3 - if not batched: - X = X[None] # (1, T, V) - B, T, V = X.shape - - # Chronos is univariate — flatten B & V into the batch axis. - flat = X.reshape(B * V, T) # (B*V, T) - with torch.no_grad(): - emb, _ = self.pipeline.embed(torch.from_numpy(flat)) # (B*V, T_tok, D) - - # (B*V, T_tok, D) -> (B, T_tok, V, D) - return emb.float().cpu().numpy().reshape(B, -1, V, emb.shape[-1]) - - -class _ChronosHookEncoder(UnpooledEncoder): - """Layer-specific path — forward hook on ``encoder.block[layer]``. - - Returns the *pre-norm* hidden state at the chosen block. Negative - indices are allowed (``-1`` = last block). - """ - - def __init__(self, pipeline: ChronosPipeline, layer: int): - self.pipeline = pipeline - n_blocks = len(pipeline.model.model.encoder.block) - if not -n_blocks <= layer < n_blocks: - raise IndexError( - f"layer {layer} out of range for {n_blocks} encoder blocks" - ) - self._block_idx = layer % n_blocks - def encode(self, X) -> np.ndarray: - X = np.asarray(X, dtype=np.float32) - batched = X.ndim == 3 - if not batched: - X = X[None] # (1, T, V) - B, T, V = X.shape - - flat = X.reshape(B * V, T) # (B*V, T) - context = torch.from_numpy(flat) - token_ids, attn_mask, _ = self.pipeline.tokenizer.context_input_transform( - context - ) - device = self.pipeline.model.device - token_ids = token_ids.to(device) - attn_mask = attn_mask.to(device) - - encoder = self.pipeline.model.model.encoder - captured = {} - - def _hook(_module, _inputs, output): - hidden = output[0] if isinstance(output, tuple) else output - captured["h"] = hidden.detach() - - handle = encoder.block[self._block_idx].register_forward_hook(_hook) - try: - with torch.no_grad(): - encoder(input_ids=token_ids, attention_mask=attn_mask) - finally: - handle.remove() - - # (B*V, T_tok, D) -> (B, T_tok, V, D) - return ( - captured["h"] - .float() - .cpu() - .numpy() - .reshape(B, -1, V, captured["h"].shape[-1]) - ) - - -def ChronosEncoder( - pipeline: ChronosPipeline, layer: int | None = None -) -> UnpooledEncoder: - """Build a Chronos feature extractor. - - Parameters - ---------- - pipeline : ChronosPipeline - A loaded Chronos pipeline. - layer : int, optional - Encoder block index to read hidden states from. ``None`` (default) - uses :meth:`ChronosPipeline.embed`, which returns post-final-norm - states from the full encoder. An integer ``layer`` registers a - forward hook on ``encoder.block[layer]`` and returns the pre-norm - hidden state there. Negative indexing supported. - - Returns - ------- - UnpooledEncoder - Object exposing ``encode(x: np.ndarray (T, C)) -> np.ndarray - (T_tok, C, D)``. Embeddings are *not* pooled. - - Notes - ----- - ``ChronosEncoder(pipeline)`` and ``ChronosEncoder(pipeline, layer=-1)`` - differ only by ``encoder.final_layer_norm`` — they will be close but - not identical. - """ - if layer is None: - return _ChronosEmbedEncoder(pipeline) - return _ChronosHookEncoder(pipeline, layer) - - -# --------------------------------------------------------------------------- -# Solver -# --------------------------------------------------------------------------- - - -class Solver(BaseSolver): - """Chronos zero-shot solver. +class Solver(BaseTSFMSolver): + """Chronos v1 zero-shot solver. Parameters ---------- model_size : str Chronos model variant: "tiny", "mini", "small", "base", "large". - layer : int or None - Encoder block index for classification embeddings. ``None`` uses - ``ChronosPipeline.embed`` (post-final-norm). pooler : {"mean", "max", "last"} - Pooling strategy over the time-token axis for classification. - task_adaptation : str - Per-task usage of the forecaster: - ``"zeroshot"`` — direct forecasting (forecasting only) - ``"forecast_residual"`` — anomaly score = forecast error (AD only) + Pooling strategy over the time-token axis for embed_batch. """ name = "Chronos" requirements = ["pip::chronos-forecasting>=2.2", "pip::torch"] + # Cap how many univariate histories go through one pipeline.predict call. + # Chronos v1 left-pads the whole list into a single T5 forward whose + # attention is O(L^2) in the padded length, so batching everything at once + # blows up memory on datasets with many long series (e.g. m4_weekly). + PREDICT_BATCH_SIZE = 16 + parameters = { "model_size": ["small"], - "layer": [None], "pooler": ["mean"], "classifier": ["log_reg"], "penalty": ["l2"], @@ -275,52 +55,85 @@ class Solver(BaseSolver): "n_estimators": [100], } - def skip(self, task, **kwargs): - if task not in SUPPORTED_TASKS: - return True, f"Chronos solver does not support task={task!r}" - return False, None - - def set_objective(self, X_train, y_train, task, **meta): - self.task = task - self.X_train = X_train - self.y_train = y_train - self.meta = meta - - # bfloat16 is fine on CUDA but poorly supported on CPU / MPS; - # fall back to float32 there so inference doesn't crash or stall. - device = "cuda" if torch.cuda.is_available() else "cpu" - dtype = torch.bfloat16 if device == "cuda" else torch.float32 - model_id = f"amazon/chronos-t5-{self.model_size}" - if not hasattr(self, "_pipeline") or self._loaded_model != model_id: - self._pipeline = ChronosPipeline.from_pretrained( - model_id, - device_map=device, - dtype=dtype, - ) - self._loaded_model = model_id + @property + def supported_tasks(self) -> set: + return {"forecasting", "classification", "anomaly_detection"} - def run(self, _): - pred_len = self.meta.get("prediction_length", 1) - if self.task == "forecasting": - self._adapter = _ChronosForecaster(self._pipeline, pred_len) + @property + def model_id(self) -> str: + return f"amazon/chronos-t5-{self.model_size}" - elif self.task == "classification": - base_encoder = ChronosEncoder(self._pipeline, layer=self.layer) - encoder = Encoder(base_encoder, POOLERS[self.pooler]()) - adapter = LinearProbeAdapter( - encoder, - task="classification", - n_classes=self.meta.get("n_classes"), - ) - adapter.fit(self.X_train, self.y_train) - self._adapter = adapter + def load_model(self, device, dtype): + return ChronosPipeline.from_pretrained( + self.model_id, + device_map=device, + dtype=dtype, + ) - elif self.task == "anomaly_detection": - # AD scores forecast residuals over an adaptive horizon. - self._adapter = ForecastResidualAdapter( - # prediction_length is ignored by the forecaster in AD mode - _ChronosForecaster(self._pipeline, prediction_length=1), + @property + def quantile_levels(self) -> tuple[float, ...]: + return (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9) + + def forecast_batch(self, inputs, covariates, prediction_length): + # Chronos v1 is univariate — split each (T, C) input into C 1-D tensors. + univariate: list[torch.Tensor] = [] + layout: list[tuple[int, int]] = [] # (inp_idx, channel_idx) + + for inp_idx, inp in enumerate(inputs): + x = inp.float().cpu() # (T, C) + if x.ndim == 1: + x = x.unsqueeze(-1) + for c in range(x.shape[1]): + univariate.append(x[:, c]) # (T,) + layout.append((inp_idx, c)) + + # Seed before sampling: Chronos v1 is non-deterministic by default. + # A fixed seed ensures identical histories produce identical draws — + # required by the leakage probe which compares two predict() calls. + torch.manual_seed(0) + all_samples: list[torch.Tensor] = [] + with torch.no_grad(): + for start in range(0, len(univariate), self.PREDICT_BATCH_SIZE): + batch = univariate[start : start + self.PREDICT_BATCH_SIZE] + all_samples.append( + self.model.predict(batch, prediction_length=prediction_length) + ) + samples = torch.cat(all_samples, dim=0).float().cpu().numpy() + # samples: (total_univariate, num_samples, H) + + q_levels = self.quantile_levels + q_arr = np.quantile(samples, q=list(q_levels), axis=1).transpose(1, 0, 2) + # q_arr: (total_univariate, Q, H) + + # Reassemble per-input (H, C, Q) tensors. + by_input: dict[int, dict[int, np.ndarray]] = {} + for k, (inp_idx, c) in enumerate(layout): + by_input.setdefault(inp_idx, {})[c] = q_arr[k] # (Q, H) + + results = [] + for inp_idx in range(len(inputs)): + channels = by_input[inp_idx] + # q_arr[k] is (Q, H); .T gives (H, Q); stack over C → (H, C, Q) + result = np.stack([channels[c].T for c in range(len(channels))], axis=1) + results.append(torch.from_numpy(result)) + return results + + def embed_batch(self, inputs): + pooler = POOLERS[self.pooler]() + results = [] + for inp in inputs: + x = inp.float().cpu().numpy() # (T, C) + if x.ndim == 1: + x = x[:, None] + # Flatten channels into batch axis (pipeline is univariate) + flat = torch.from_numpy(x.T.copy()) # (C, T) + with torch.no_grad(): + emb, _ = self.model.embed(flat) # (C, T_tok, D) + emb_np = emb.float().cpu().numpy() + # (C, T_tok, D) → (1, T_tok, C, D) to match pooler convention + emb_4d = emb_np.transpose(1, 0, 2)[None] + pooled = pooler.pool(emb_4d) # (1, C, D) + results.append( + torch.from_numpy(pooled[0].reshape(-1).astype(np.float32)) ) - - def get_result(self): - return {"model": self._adapter} + return results diff --git a/solvers/chronos2.py b/solvers/chronos2.py index b12c080..ddcc272 100644 --- a/solvers/chronos2.py +++ b/solvers/chronos2.py @@ -1,9 +1,11 @@ """Chronos-2 solver for the TSFM benchmark (local inference). -Supports: - - forecasting : zero-shot via Chronos2Pipeline - - classification : linear probe on pooled encoder embeddings - - anomaly_detection : forecast-residual on top of the same forecaster +Directly supports: + - forecasting : zero-shot via ``Chronos2Pipeline.predict`` + +Additionally provides: + - ``embed_batch``: pooled encoder embeddings for classification and + anomaly detection via the default adaptation strategies. Model loading is done in ``set_objective`` (untimed). Inference batches every (series, cutoff) pair into a single ``Chronos2Pipeline.predict`` @@ -16,168 +18,25 @@ https://github.com/amazon-science/chronos-forecasting """ +import functools + import numpy as np import torch -from benchopt import BaseSolver from chronos.chronos2 import Chronos2Pipeline -from benchmark_utils.adapters import ( - POOLERS, - BaseTSFMAdapter, - Encoder, - LinearProbeAdapter, - UnpooledEncoder, -) -from benchmark_utils.adapters.forecast_residual import ForecastResidualAdapter -from benchmark_utils.inputs import ForecastInput -from benchmark_utils.outputs import ForecastOutput - -SUPPORTED_TASKS = {"forecasting", "classification", "anomaly_detection"} - -# --------------------------------------------------------------------------- -# Chronos-2 encoders — embed() has a different signature than Chronos v1: -# Chronos2Pipeline.embed takes (B, V, T) and returns List[(V, T_tok, D)]. -# --------------------------------------------------------------------------- - - -def _to_context(x): - """Reshape ``(T, V)`` or ``(B, T, V)`` to Chronos-2 input ``(B, V, T)``.""" - x = np.asarray(x, dtype=np.float32) - if x.ndim == 2: - x = x[None] - return x.transpose(0, 2, 1) - - -class _Chronos2Forecaster(BaseTSFMAdapter): - """Chronos-2 forecaster — uses native quantile output from the pipeline.""" - - def __init__(self, pipeline, prediction_length): - self.pipeline = pipeline - self.prediction_length = prediction_length - self.quantile_levels = tuple(float(q) for q in pipeline.quantiles) - - def predict(self, x: ForecastInput, prediction_length=None) -> ForecastOutput: - horizon = prediction_length or self.prediction_length - inputs, layout, per_series_shape = self._build_inputs(x) - if not inputs: - return ForecastOutput(quantiles=[], quantile_levels=self.quantile_levels) - - with torch.no_grad(): - output = self.pipeline.predict( - inputs, - prediction_length=horizon, - ) - return self._assemble_output(output, layout, per_series_shape, horizon) - - def _build_inputs(self, x): - """Build (C, T) tensors (all channels together); layout omits channel idx.""" - inputs = [] - layout = [] # (series_idx, cutoff_idx) - per_series_shape = [] # (C, n_cutoffs) - for series_idx, (series, cutoffs) in enumerate(zip(x.x, x.cutoff_indexes)): - series = np.asarray(series, dtype=np.float32) - if series.ndim == 1: - series = series[:, None] - _, C = series.shape - per_series_shape.append((C, len(cutoffs))) - for cutoff_idx, cutoff in enumerate(cutoffs): - hist = series[:cutoff] - inputs.append(torch.from_numpy(hist.T)) # (C, T_cutoff) - layout.append((series_idx, cutoff_idx)) - return inputs, layout, per_series_shape - - def _assemble_output(self, forecast, layout, per_series_shape, prediction_length): - """Use quantile tensors directly from Chronos-2 pipeline.""" - # forecast: list[(n_variates, Q, prediction_length)] - Q = len(self.quantile_levels) - per_series = [ - np.empty((n_cutoffs, prediction_length, C, Q), dtype=np.float32) - for C, n_cutoffs in per_series_shape - ] - for (series_idx, cutoff_idx), pred in zip(layout, forecast): - arr = pred.float().cpu().numpy() # (C, Q, H) - per_series[series_idx][cutoff_idx] = arr.transpose(2, 0, 1) # (H, C, Q) - return ForecastOutput( - quantiles=per_series, quantile_levels=self.quantile_levels - ) - - -class _Chronos2EmbedEncoder(UnpooledEncoder): - """Uses ``Chronos2Pipeline.embed`` which returns a list of tensors.""" - - def __init__(self, pipeline): - self.pipeline = pipeline - - def encode(self, X) -> np.ndarray: - context = _to_context(X) # (B, V, T) - with torch.no_grad(): - embeddings, _ = self.pipeline.embed(context) # list[(V, T_tok, D)] - stacked = torch.stack(list(embeddings)) # (B, V, T_tok, D) - return stacked.transpose(1, 2).float().cpu().numpy() # (B, T_tok, V, D) - - -class _Chronos2HookEncoder(UnpooledEncoder): - """Forward hook on ``encoder.block[layer]``.""" - - def __init__(self, pipeline, layer: int): - self.pipeline = pipeline - n_blocks = len(pipeline.model.model.encoder.block) - if not -n_blocks <= layer < n_blocks: - raise IndexError( - f"layer {layer} out of range for {n_blocks} encoder blocks" - ) - self._block_idx = layer % n_blocks - - def encode(self, X) -> np.ndarray: - context = _to_context(X) # (B, V, T) - token_ids, attn_mask, _ = self.pipeline.tokenizer.context_input_transform( - torch.from_numpy(context) - ) - device = self.pipeline.model.device - token_ids = token_ids.to(device) - attn_mask = attn_mask.to(device) - - encoder = self.pipeline.model.model.encoder - captured = {} - - def _hook(_module, _inputs, output): - hidden = output[0] if isinstance(output, tuple) else output - captured["h"] = hidden.detach() - - handle = encoder.block[self._block_idx].register_forward_hook(_hook) - try: - with torch.no_grad(): - encoder(input_ids=token_ids, attention_mask=attn_mask) - finally: - handle.remove() - - return captured["h"].float().cpu().numpy() - - -def _Chronos2Encoder(pipeline, layer=None): - """Build a Chronos-2 feature extractor.""" - if layer is None: - return _Chronos2EmbedEncoder(pipeline) - return _Chronos2HookEncoder(pipeline, layer) - +from benchmark_utils.adapters import POOLERS +from benchmark_utils.base_solver import BaseTSFMSolver -# --------------------------------------------------------------------------- -# Solver -# --------------------------------------------------------------------------- - -class Solver(BaseSolver): +class Solver(BaseTSFMSolver): """Chronos-2 zero-shot solver. Parameters ---------- model_size : str Chronos-2 model variant: "tiny", "mini", "small", "base", "large". - layer : int or None - Encoder block index for classification embeddings. ``None`` uses - ``Chronos2Pipeline.embed`` (post-final-norm). pooler : {"mean", "max", "last"} - Pooling strategy over the time-token axis for classification. + Pooling strategy over the time-token axis for embed_batch. """ name = "Chronos2" @@ -186,7 +45,6 @@ class Solver(BaseSolver): parameters = { "model_size": ["small"], - "layer": [None], "pooler": ["mean"], "classifier": ["log_reg"], "penalty": ["l2"], @@ -195,48 +53,49 @@ class Solver(BaseSolver): "n_estimators": [100], } - def skip(self, task, **kwargs): - if task not in SUPPORTED_TASKS: - return True, f"Chronos2 solver does not support task={task!r}" - return False, None - - def set_objective(self, X_train, y_train, task, **meta): - self.task = task - self.X_train = X_train - self.y_train = y_train - self.meta = meta - - device = "cuda" if torch.cuda.is_available() else "cpu" - dtype = torch.bfloat16 if device == "cuda" else torch.float32 - model_id = f"autogluon/chronos-2-{self.model_size}" - if not hasattr(self, "_pipeline") or self._loaded_model != model_id: - self._pipeline = Chronos2Pipeline.from_pretrained( - model_id, - device_map=device, - dtype=dtype, - ) - self._loaded_model = model_id - - def run(self, _): - pred_len = self.meta.get("prediction_length", 1) - if self.task == "forecasting": - self._adapter = _Chronos2Forecaster(self._pipeline, pred_len) - - elif self.task == "classification": - base_encoder = _Chronos2Encoder(self._pipeline, layer=self.layer) - encoder = Encoder(base_encoder, POOLERS[self.pooler]()) - adapter = LinearProbeAdapter( - encoder, - task="classification", - n_classes=self.meta.get("n_classes"), - ) - adapter.fit(self.X_train, self.y_train) - self._adapter = adapter + @property + def supported_tasks(self): + return {"forecasting", "classification", "anomaly_detection"} - elif self.task == "anomaly_detection": - self._adapter = ForecastResidualAdapter( - _Chronos2Forecaster(self._pipeline, prediction_length=1), - ) + @property + def model_id(self): + return f"autogluon/chronos-2-{self.model_size}" + + def load_model(self, device, dtype): + return Chronos2Pipeline.from_pretrained( + self.model_id, + device_map=device, + dtype=dtype, + ) - def get_result(self): - return {"model": self._adapter} + @functools.cached_property + def quantile_levels(self): + return tuple(float(q) for q in self.model.quantiles) + + def forecast_batch(self, inputs, covariates, prediction_length): + # inputs: list of (T, C) tensors — Chronos-2 expects (C, T) + chrono_inputs = [inp.T for inp in inputs] + with torch.no_grad(): + # pipeline.predict returns list of (C, Q, H) tensors + output = self.model.predict( + chrono_inputs, + prediction_length=prediction_length, + ) + # Convert (C, Q, H) → (H, C, Q) for base class assembly + return [pred.permute(2, 0, 1) for pred in output] + + def embed_batch(self, inputs): + pooler = POOLERS[self.pooler]() + results = [] + for inp in inputs: + context = inp.T.unsqueeze(0) # (1, C, T) — Chronos-2 input format + with torch.no_grad(): + embeddings, _ = self.model.embed(context) # list[(C, T_tok, D)] + emb_np = torch.stack(list(embeddings))[0].float().cpu().numpy() + # Reshape to (1, T_tok, C, D) to match pooler convention (..., T, V, D) + emb_4d = emb_np.transpose(1, 0, 2)[None] # (1, T_tok, C, D) + pooled = pooler.pool(emb_4d) # (1, C, D) + results.append( + torch.from_numpy(pooled[0].reshape(-1).astype(np.float32)) # (C*D,) + ) + return results diff --git a/solvers/mantis.py b/solvers/mantis.py index 2bee70b..55be87e 100644 --- a/solvers/mantis.py +++ b/solvers/mantis.py @@ -1,8 +1,8 @@ """Mantis solver for time series classification on UCR datasets. -Uses the official ``mantis-tsfm`` API to load a pretrained Mantis checkpoint, -extract embeddings with ``MantisTrainer.transform``, and train a Random Forest -classifier on top. +Additionally provides: + - ``embed_batch``: embeddings via ``MantisTrainer.transform`` for + classification via the default linear-probe adaptation strategy. References: https://huggingface.co/paris-noah/Mantis-8M @@ -11,26 +11,17 @@ import numpy as np import torch -from benchopt import BaseSolver from mantis.architecture import MantisV1, MantisV2 from mantis.trainer import MantisTrainer -from benchmark_utils.adapters.linear_probe import LinearProbeAdapter +from benchmark_utils.base_solver import BaseTSFMSolver -SUPPORTED_TASKS = {"classification"} - -class Solver(BaseSolver): - """Mantis time series classification solver with Random Forest. - - The model is loaded once in ``set_objective`` (not timed). Training - embeddings are extracted and a classification head is trained. - During ``run`` the predictions are generated for the test set. - """ +class Solver(BaseTSFMSolver): + """Mantis time series classification solver.""" name = "Mantis" - # mantis-tsfm and torch are required to load the model and run inference. requirements = [ "pip::mantis-tsfm>=1.0.0", ] @@ -46,150 +37,50 @@ class Solver(BaseSolver): "n_estimators": [100], } - def _extract_embeddings(self, X): - """Extract embeddings for a batch of time series. - - Parameters - ---------- - X : np.ndarray - Input time series of shape (N, T, C) where N is the number - of series, T is the sequence length, and C the number of channels. - Note that Mantis expects (N, C, T) internally. - batch_size : int - Batch size for processing - - Returns - ------- - np.ndarray - Embeddings of shape (n_samples, embedding_dim) - """ - batch_size = self.batch_size - n_samples = len(X) - all_embeddings = [] - - for batch_idx in range(0, n_samples, batch_size): - batch_end = min(batch_idx + batch_size, n_samples) - X_batch = np.asarray(X[batch_idx:batch_end], dtype=np.float32) - X_batch_processed = self._prepare_inputs(X_batch) - - try: - with torch.no_grad(): - embeddings_np = self._trainer.transform(X_batch_processed) - all_embeddings.append(np.asarray(embeddings_np)) + @property + def supported_tasks(self): + return {"classification"} - except Exception as e: - print(f" Warning: Failed to process batch {batch_idx}: {e}") - if all_embeddings: - embedding_dim = all_embeddings[0].shape[1] - else: - embedding_dim = 128 - all_embeddings.append( - np.zeros((batch_end - batch_idx, embedding_dim), - dtype=np.float32) - ) + @property + def model_id(self): + return self.checkpoint - # Concatenate all embeddings - if all_embeddings: - embeddings_all = np.vstack(all_embeddings) - else: - embeddings_all = np.zeros((n_samples, 768)) - - return embeddings_all + def load_model(self, device, dtype): + MantisBackbone = MantisV2 if "MantisV2" in self.checkpoint else MantisV1 + network = MantisBackbone(device=device) + network = network.from_pretrained(self.checkpoint) + return MantisTrainer(device=device, network=network) def _prepare_inputs(self, X_batch): - """Ensure Mantis-compatible shape and sequence length. - - Mantis expects arrays of shape (N, C, T), and the sequence length - should be divisible by 32. Following official guidance, - we interpolate to ``interpolate_to`` (default 512). - """ - X_in = X_batch.transpose(0, 2, 1) - - current_len = X_in.shape[-1] + """Interpolate to ``interpolate_to`` and transpose to (N, C, T).""" + X_in = X_batch.transpose(0, 2, 1) # (N, T, C) → (N, C, T) target_len = int(self.interpolate_to) - - if current_len != target_len: + if X_in.shape[-1] != target_len: tensor = torch.tensor(X_in, dtype=torch.float32) tensor = torch.nn.functional.interpolate( - tensor, - size=target_len, - mode="linear", - align_corners=False, + tensor, size=target_len, mode="linear", align_corners=False ) X_in = tensor.numpy() - if X_in.shape[-1] % 32 != 0: raise ValueError( "Sequence length must be divisible by 32 for Mantis, " f"got {X_in.shape[-1]}" ) - return X_in - def skip(self, task, **kwargs): - if task not in SUPPORTED_TASKS: - return True, f"Mantis solver does not support task={task!r}" - return False, None - - def set_objective(self, task, X_train, y_train, **meta): - """Prepare the solver for a given dataset configuration. - - Model loading is done here (not inside ``run``) so that the - checkpoint download/loading time is excluded from the benchmark - timing. - """ - self.task = task - self.X_train = X_train - self.y_train = y_train - self.meta = meta - - device = "cuda" if torch.cuda.is_available() else "cpu" - - # Load the model only on the first call for this checkpoint. - should_reload = ( - not hasattr(self, "_network") - or not hasattr(self, "_loaded_checkpoint") - or self._loaded_checkpoint != self.checkpoint - ) - if should_reload: - try: - MantisBackbone = MantisV2 if "MantisV2" in self.checkpoint else MantisV1 - network = MantisBackbone(device=device) - network = network.from_pretrained(self.checkpoint) - - self._network = network - self._trainer = MantisTrainer( - device=device, network=self._network) - self._loaded_checkpoint = self.checkpoint - print( - f"✓ Mantis checkpoint loaded: {self.checkpoint} on device: {device}" - ) - except Exception as e: - raise RuntimeError( - f"Failed to load Mantis checkpoint '{self.checkpoint}' " - f"from Hugging Face: {e}. Make sure you have internet " - "access and the model is available." - ) - - self._device = device - - def run(self, _): - """Fit the linear probe adapter on the training data.""" - self._adapter = LinearProbeAdapter( - encoder=self, - task=self.task, - classifier=self.classifier, - penalty=self.penalty, - C=self.C, - alpha=self.alpha, - n_estimators=self.n_estimators, - ) - self._adapter.fit(self.X_train, self.y_train) - - def encode(self, x): - """Encode a batch of time series into embeddings.""" - return self._extract_embeddings(x) - - def get_result(self): - """Return the fitted adapter.""" - return {"model": self._adapter} + def embed_batch(self, inputs): + # self.model is a MantisTrainer + results = [] + for i in range(0, len(inputs), self.batch_size): + batch = inputs[i : i + self.batch_size] + X_batch = np.stack( + [inp.float().cpu().numpy() for inp in batch] + ) # (B, T, C) + X_prepared = self._prepare_inputs(X_batch) # (B, C, T_interp) + with torch.no_grad(): + emb = np.asarray( + self.model.transform(X_prepared), dtype=np.float32 + ) # (B, D) + for row in emb: + results.append(torch.from_numpy(row)) + return results diff --git a/solvers/moment.py b/solvers/moment.py index 563ff6c..e3c8c62 100644 --- a/solvers/moment.py +++ b/solvers/moment.py @@ -1,13 +1,13 @@ """Moment solver for the TSFM benchmark. -Moment is a time series foundation model from Alibaba. This solver supports: - - forecasting : zero-shot via Moment pipeline - - classification : linear probe on pooled encoder embeddings +Directly supports: + - forecasting : zero-shot point forecast via ``MOMENTPipeline.forecast`` -Model loading is done in ``set_objective`` (untimed). For forecasting, -inference batches every (series, cutoff) pair into a single forward pass. -For classification, training embeddings are extracted and a linear probe -or classifier is trained on top. +Additionally provides: + - ``embed_batch``: pooled patch embeddings for classification via the + default adaptation strategies. + +Model loading is done in ``set_objective`` (untimed). References: https://huggingface.co/AutonLab/MOMENT-1-large @@ -15,154 +15,20 @@ import numpy as np import torch -from benchopt import BaseSolver from momentfm import MOMENTPipeline -from benchmark_utils.adapters import ( - Encoder, - LastPooler, - LinearProbeAdapter, - MaxPooler, - MeanPooler, - UnpooledEncoder, -) -from benchmark_utils.adapters.base import BaseTSFMAdapter -from benchmark_utils.inputs import ForecastInput -from benchmark_utils.outputs import ForecastOutput - -SUPPORTED_TASKS = {"forecasting", "classification"} - -POOLERS = { - "mean": MeanPooler, - "max": MaxPooler, - "last": LastPooler, -} - - -class _MomentForecaster(BaseTSFMAdapter): - """Moment forecasting adapter.""" - - def __init__(self, pipeline, prediction_length): - self.pipeline = pipeline - self.prediction_length = prediction_length - - def predict(self, x: ForecastInput) -> ForecastOutput: - quantiles = [] - for series, cutoffs in zip(x.x, x.cutoff_indexes): - series = np.asarray(series, dtype=np.float32) - if series.ndim == 1: - series = series[:, None] - T, C = series.shape - - preds_per_series = [] - for cutoff in cutoffs: - hist = series[:cutoff] # (T_cutoff, C) - - if hist.ndim == 1: - hist = hist[None, :] - - # Moment expects (B, channels, seq_len) - hist_tensor = ( - torch.from_numpy(hist.transpose(1, 0)).unsqueeze(0).float() - ) - device = next(self.pipeline.parameters()).device - hist_tensor = hist_tensor.to(device) - input_mask = torch.ones( - (hist_tensor.shape[0], hist_tensor.shape[2]), - dtype=torch.float32, - device=device, - ) - - with torch.no_grad(): - outputs = self.pipeline.forecast( - x_enc=hist_tensor, - input_mask=input_mask, - prediction_length=self.prediction_length, - ) - - forecast = outputs.forecast if hasattr(outputs, "forecast") else outputs - if isinstance(forecast, tuple): - forecast = forecast[0] - - arr = forecast.squeeze(0).cpu().numpy() - - if arr.ndim == 1: - arr = arr[:, None] - - # Moment returns (channels, horizon) by default. - if arr.ndim == 2 and arr.shape[0] != self.prediction_length: - arr = arr.T - - if arr.shape[0] > self.prediction_length: - arr = arr[: self.prediction_length] - - if arr.shape[0] != self.prediction_length: - raise ValueError( - f"Unexpected forecast shape after transpose/slice: {arr.shape}" - ) - - preds_per_series.append(arr) - - # Stack predictions: (n_cutoffs, prediction_length, C) - stacked = np.stack(preds_per_series, axis=0) - # Add quantile dimension: (n_cutoffs, prediction_length, C, 1) - quantiles.append(stacked[:, :, :, None]) - - return ForecastOutput(quantiles=quantiles, quantile_levels=(0.5,)) +from benchmark_utils.adapters import POOLERS +from benchmark_utils.base_solver import BaseTSFMSolver -class _MomentEncoder(UnpooledEncoder): - """Moment encoder for extracting embeddings.""" - - def __init__(self, pipeline): - self.pipeline = pipeline - - def encode(self, X) -> np.ndarray: - """Extract embeddings from time series data. - - Args: - X: np.ndarray of shape (T, C) or (B, T, C) - - Returns: - np.ndarray of shape (B, T, V, D) - """ - X = np.asarray(X, dtype=np.float32) - if X.ndim == 2: - X = X[None] - elif X.ndim != 3: - raise ValueError(f"Unexpected input shape for Moment encoder: {X.shape}") - - # Moment expects (B, channels, seq_len) - X = X.transpose(0, 2, 1) - - with torch.no_grad(): - X_tensor = torch.from_numpy(X).float() - device = next(self.pipeline.parameters()).device - X_tensor = X_tensor.to(device) - outputs = self.pipeline.embed(x_enc=X_tensor, reduction="none") - emb = outputs.embeddings - - if isinstance(emb, torch.Tensor): - emb = emb.cpu().numpy() - - if emb.ndim != 4: - raise ValueError(f"Unexpected Moment embedding shape: {emb.shape}") - - # Moment returns (B, channels, n_patches, D); transform to - # (B, n_patches, channels, D) for the benchmark encoder API. - return emb.transpose(0, 2, 1, 3) - - -class Solver(BaseSolver): +class Solver(BaseTSFMSolver): """Moment foundation model solver. - Supports forecasting (zero-shot) and classification (with linear probe). - The model is loaded once in ``set_objective`` (not timed). + Supports forecasting (zero-shot) and classification (linear probe). """ name = "Moment" - # moment-fm package required for the model requirements = [ "pip::momentfm @ git+https://github.com/moment-timeseries-foundation-model/moment.git", ] @@ -171,8 +37,7 @@ class Solver(BaseSolver): parameters = { "checkpoint": ["AutonLab/MOMENT-1-large"], - "pooler": ["mean"], # pooler for classification embeddings - "batch_size": [32], + "pooler": ["mean"], "classifier": ["log_reg"], "penalty": ["l2"], "C": [1.0], @@ -184,76 +49,75 @@ class Solver(BaseSolver): "checkpoint": "AutonLab/MOMENT-1-small", } - def skip(self, task, **kwargs): - if task not in SUPPORTED_TASKS: - return True, f"Moment solver does not support task={task!r}" - return False, None + @property + def supported_tasks(self): + return {"forecasting", "classification"} - def set_objective(self, X_train, y_train, task, **meta): - """Prepare the solver for a given dataset configuration. + @property + def model_id(self): + return self.checkpoint - Model loading is done here (not inside ``run``) so that the - checkpoint download/loading time is excluded from the benchmark - timing. - """ - self.task = task - self.X_train = X_train - self.y_train = y_train - self.meta = meta - self.prediction_length = meta.get("prediction_length") - - device = "cuda" if torch.cuda.is_available() else "cpu" - - # Load the model only on the first call for this checkpoint - should_reload = ( - not hasattr(self, "_pipeline") - or not hasattr(self, "_loaded_checkpoint") - or self._loaded_checkpoint != self.checkpoint + def load_model(self, device, dtype): + pipeline = MOMENTPipeline.from_pretrained( + self.checkpoint, + torch_dtype=torch.float32, ) - if should_reload: - try: - self._pipeline = MOMENTPipeline.from_pretrained( - self.checkpoint, - torch_dtype=torch.float32, - ) - self._pipeline = self._pipeline.to(device) - self._loaded_checkpoint = self.checkpoint - print( - f"✓ Moment checkpoint loaded: {self.checkpoint} on device: {device}" - ) - except Exception as e: - raise RuntimeError( - f"Failed to load Moment checkpoint '{self.checkpoint}' " - f"from Hugging Face: {e}. Make sure you have internet " - "access and the model is available." - ) - - self._device = device - - def run(self, _): - """Fit the model or adapter on the training data.""" - if self.task == "forecasting": - self._adapter = _MomentForecaster( - pipeline=self._pipeline, - prediction_length=self.prediction_length, + return pipeline.to(device) + + @property + def quantile_levels(self): + return (0.5,) # Moment outputs point forecasts only + + def forecast_batch(self, inputs, covariates, prediction_length): + device = next(self.model.parameters()).device + results = [] + for inp in inputs: + # inp: (T, C) — Moment expects (1, C, T) + x_enc = inp.float().T.unsqueeze(0).to(device) # (1, C, T) + input_mask = torch.ones( + 1, inp.shape[0], dtype=torch.float32, device=device ) - elif self.task == "classification": - base_encoder = _MomentEncoder(pipeline=self._pipeline) - encoder = Encoder(base_encoder, POOLERS[self.pooler]()) - self._adapter = LinearProbeAdapter( - encoder=encoder, - task="classification", - n_classes=self.meta.get("n_classes"), - classifier=self.classifier, - n_estimators=self.n_estimators, - C=self.C, - alpha=self.alpha, - penalty=self.penalty, - ) - self._adapter.fit(self.X_train, self.y_train) - else: - raise ValueError(f"Unsupported task: {self.task}") + with torch.no_grad(): + outputs = self.model.forecast( + x_enc=x_enc, + input_mask=input_mask, + prediction_length=prediction_length, + ) - def get_result(self): - return {"model": self._adapter} + forecast = outputs.forecast if hasattr(outputs, "forecast") else outputs + if isinstance(forecast, tuple): + forecast = forecast[0] + + arr = forecast.squeeze(0).float().cpu() # (C, H) or (H,) + if arr.ndim == 1: + arr = arr.unsqueeze(-1) # (H, 1) + elif arr.shape[0] != prediction_length: + arr = arr.T # (H, C) + arr = arr[:prediction_length] + + results.append(arr.unsqueeze(-1)) # (H, C, 1) + return results + + def embed_batch(self, inputs): + pooler = POOLERS[self.pooler]() + device = next(self.model.parameters()).device + results = [] + for inp in inputs: + # inp: (T, C) — Moment expects (1, C, T) + x_enc = inp.float().T.unsqueeze(0).to(device) # (1, C, T) + + with torch.no_grad(): + outputs = self.model.embed(x_enc=x_enc, reduction="none") + emb = outputs.embeddings # (1, C, n_patches, D) + + if isinstance(emb, torch.Tensor): + emb = emb.cpu().numpy() + + # (1, C, n_patches, D) → (1, n_patches, C, D) to match pooler convention + emb_4d = emb.transpose(0, 2, 1, 3) + pooled = pooler.pool(emb_4d) # (1, C, D) + results.append( + torch.from_numpy(pooled[0].reshape(-1).astype(np.float32)) + ) + return results diff --git a/solvers/toto2.py b/solvers/toto2.py index 4de7334..68a1ded 100644 --- a/solvers/toto2.py +++ b/solvers/toto2.py @@ -1,9 +1,11 @@ """Toto-2.0 solver for the TSFM benchmark (local inference). -Supports: - - forecasting : zero-shot via Toto2Model - - classification : linear probe on pooled transformer patch embeddings - - anomaly_detection : forecast-residual on top of the same forecaster +Directly supports: + - forecasting : zero-shot via ``Toto2Model.forecast`` + +Additionally provides: + - ``embed_batch``: pooled transformer patch embeddings for classification + and anomaly detection via the default adaptation strategies. References: https://github.com/datadog/toto @@ -11,252 +13,13 @@ import numpy as np import torch -from benchopt import BaseSolver from toto2 import Toto2Model -from benchmark_utils.adapters import ( - Encoder, - LastPooler, - LinearProbeAdapter, - MaxPooler, - MeanPooler, - UnpooledEncoder, -) -from benchmark_utils.adapters.base import BaseTSFMAdapter -from benchmark_utils.adapters.forecast_residual import ForecastResidualAdapter -from benchmark_utils.inputs import ForecastInput -from benchmark_utils.outputs import ForecastOutput - -SUPPORTED_TASKS = {"forecasting", "classification", "anomaly_detection"} - -POOLERS = { - "mean": MeanPooler, - "max": MaxPooler, - "last": LastPooler, -} - - -# --------------------------------------------------------------------------- -# Thin wrapper exposing the predict() interface expected by the objective -# --------------------------------------------------------------------------- - - -class _Toto2Forecaster(BaseTSFMAdapter): - """Toto2Model adapter for the forecasting contract.""" - - quantile_levels = tuple(float(q) / 10 for q in range(1, 10)) - - def __init__( - self, - model, - device, - prediction_length, - context_length=None, - decode_block_size=None, - patch_size=32, - ): - self.model = model - self.device = device - self.prediction_length = prediction_length - self.context_length = context_length - self.decode_block_size = decode_block_size - self.patch_size = patch_size - - def _forecast_context(self, context: np.ndarray) -> np.ndarray: - import torch - - x = np.asarray(context, dtype=np.float32) - if self.context_length is not None: - x = x[-self.context_length :] - - # Toto expects (batch, n_variates, time_steps). - target_np = np.swapaxes(x, 0, 1)[None, :, :] - finite_mask_np = np.isfinite(target_np) - target_np = np.nan_to_num(target_np, nan=0.0, posinf=0.0, neginf=0.0) - - pad_len = (-target_np.shape[-1]) % self.patch_size - if pad_len: - target_np = np.pad( - target_np, - ((0, 0), (0, 0), (pad_len, 0)), - mode="constant", - constant_values=0.0, - ) - finite_mask_np = np.pad( - finite_mask_np, - ((0, 0), (0, 0), (pad_len, 0)), - mode="constant", - constant_values=False, - ) - - has_missing_values = not bool(finite_mask_np.all()) +from benchmark_utils.adapters import POOLERS +from benchmark_utils.base_solver import BaseTSFMSolver - target = torch.from_numpy(target_np).to(self.device) - target_mask = torch.from_numpy(finite_mask_np).to(self.device, dtype=torch.bool) - series_ids = torch.zeros( - target.shape[0], - target.shape[1], - dtype=torch.long, - device=self.device, - ) - - with torch.inference_mode(): - quantiles = self.model.forecast( - { - "target": target, - "target_mask": target_mask, - "series_ids": series_ids, - }, - horizon=self.prediction_length, - decode_block_size=self.decode_block_size, - has_missing_values=has_missing_values, - ) - - # quantiles[:, 0]: (Q, C, H) -> (H, C, Q) - return quantiles[:, 0].detach().float().cpu().numpy().transpose(2, 1, 0) - - def predict(self, x: ForecastInput) -> ForecastOutput: - per_series = [] - - for series, cutoffs in zip(x.x, x.cutoff_indexes): - series = np.asarray(series, dtype=np.float32) - if series.ndim == 1: - series = series[:, None] - - C = series.shape[1] - forecasts = np.empty( - ( - len(cutoffs), - self.prediction_length, - C, - len(self.quantile_levels), - ), - dtype=np.float32, - ) - for cutoff_idx, cutoff in enumerate(cutoffs): - forecasts[cutoff_idx] = self._forecast_context(series[:cutoff]) - - per_series.append(forecasts) - - return ForecastOutput( - quantiles=per_series, - quantile_levels=self.quantile_levels, - ) - - -class _Toto2EmbedEncoder(UnpooledEncoder): - """Use Toto-2 transformer patch states as sequence embeddings. - - ``layer=None`` captures the final output of ``model.transformer`` - after its output norm. An integer ``layer`` registers a forward hook - on ``model.transformer.layers[layer]`` and returns the output of that - selected Toto transformer layer. Negative indexing is supported - (``-1`` = last layer). - """ - - def __init__(self, model, device, context_length=None, layer=None): - self.model = model - self.device = device - self.context_length = context_length - self.layer = layer - self.patch_size = model.config.patch_size - - def _prepare_batch(self, X): - x = np.asarray(X, dtype=np.float32) - if x.ndim == 1: - x = x[None, :, None] - elif x.ndim == 2: - x = x[None] - elif x.ndim != 3: - raise ValueError( - "Toto-2 classification expects input with shape " - f"(time, variates) or (batch, time, variates); got {x.shape}." - ) - if self.context_length is not None: - x = x[:, -self.context_length :, :] - - # Toto expects (batch, variates, time_steps). - batch = x.transpose(0, 2, 1) - mask = np.isfinite(batch) - batch = np.nan_to_num(batch, nan=0.0, posinf=0.0, neginf=0.0) - - pad_len = (-batch.shape[-1]) % self.patch_size - if pad_len: - batch = np.pad( - batch, - ((0, 0), (0, 0), (pad_len, 0)), - mode="constant", - constant_values=0.0, - ) - mask = np.pad( - mask, - ((0, 0), (0, 0), (pad_len, 0)), - mode="constant", - constant_values=False, - ) - - return batch, mask - - def encode(self, X) -> np.ndarray: - import torch - - target_np, target_mask_np = self._prepare_batch(X) - - target = torch.from_numpy(target_np).to(self.device) - target_mask = torch.from_numpy(target_mask_np).to(self.device, dtype=torch.bool) - cpm_mask = torch.ones_like(target_mask) - series_ids = torch.zeros( - target.shape[0], - target.shape[1], - dtype=torch.long, - device=self.device, - ) - captured_output = {} - - if self.layer is None: - hook_module = self.model.transformer - else: - n_layers = len(self.model.transformer.layers) - if not -n_layers <= self.layer < n_layers: - raise IndexError( - f"layer {self.layer} out of range for {n_layers} " - "Toto-2 transformer layers" - ) - hook_module = self.model.transformer.layers[self.layer % n_layers] - - def _capture_output(_module, _inputs, output): - captured_output["hidden_state"] = output.detach() - - handle = hook_module.register_forward_hook(_capture_output) - with torch.inference_mode(): - try: - self.model( - target=target, - target_mask=target_mask, - cpm_mask=cpm_mask, - series_ids=series_ids, - ) - finally: - handle.remove() - - embeddings = captured_output["hidden_state"] - if embeddings.ndim != 4: - raise ValueError( - "Toto-2 transformer output should have shape " - f"(batch, variates, patches, dim); got {tuple(embeddings.shape)}." - ) - - # Toto: (B, C, T_patch, D) -> benchmark: (B, T_patch, C, D). - return embeddings.transpose(1, 2).float().cpu().numpy() - - -# --------------------------------------------------------------------------- -# Solver -# --------------------------------------------------------------------------- - - -class Solver(BaseSolver): +class Solver(BaseTSFMSolver): """Datadog Toto-2.0 zero-shot solver.""" name = "Toto-2.0" @@ -286,60 +49,125 @@ class Solver(BaseSolver): "checkpoint": "Datadog/Toto-2.0-4m", } - def skip(self, task, **kwargs): - if task not in SUPPORTED_TASKS: - return True, f"Toto-2.0 solver does not support task={task!r}" - return False, None - - def set_objective(self, X_train, y_train, task, **meta): - self.task = task - self.X_train = X_train - self.y_train = y_train - self.meta = meta + @property + def supported_tasks(self): + return {"forecasting", "classification", "anomaly_detection"} + + @property + def model_id(self): + return self.checkpoint + + def load_model(self, device, dtype): + model = Toto2Model.from_pretrained(self.checkpoint) + return model.to(device).eval() + + @property + def quantile_levels(self) -> tuple[float, ...]: + return tuple(float(q) / 10 for q in range(1, 10)) + + def forecast_batch(self, inputs, covariates, prediction_length): + device = next(self.model.parameters()).device + results = [] + for inp in inputs: + x = inp.float().cpu().numpy() # (T, C) + if self.context_length is not None: + x = x[-self.context_length:] + + # Toto expects (B, C, T) + target_np = x.T[None] # (1, C, T) + finite_mask_np = np.isfinite(target_np) + target_np = np.nan_to_num( + target_np, nan=0.0, posinf=0.0, neginf=0.0 + ) - self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - should_reload = ( - not hasattr(self, "_model") or self._loaded_checkpoint != self.checkpoint - ) - if should_reload: - self._model = Toto2Model.from_pretrained(self.checkpoint) - self._model = self._model.to(self._device).eval() - self._loaded_checkpoint = self.checkpoint + pad_len = (-target_np.shape[-1]) % self.patch_size + if pad_len: + target_np = np.pad( + target_np, ((0, 0), (0, 0), (pad_len, 0)), constant_values=0.0 + ) + finite_mask_np = np.pad( + finite_mask_np, + ((0, 0), (0, 0), (pad_len, 0)), + constant_values=False, + ) - def run(self, _): - pred_len = self.meta.get("prediction_length", 1) - forecaster = _Toto2Forecaster( - self._model, - self._device, - prediction_length=pred_len, - context_length=self.context_length, - decode_block_size=self.decode_block_size, - patch_size=self.patch_size, - ) + target = torch.from_numpy(target_np).to(device) + target_mask = torch.from_numpy(finite_mask_np).to(device, dtype=torch.bool) + series_ids = torch.zeros( + 1, target.shape[1], dtype=torch.long, device=device + ) - if self.task == "forecasting": - self._adapter = forecaster + with torch.inference_mode(): + quantiles = self.model.forecast( + { + "target": target, + "target_mask": target_mask, + "series_ids": series_ids, + }, + horizon=prediction_length, + decode_block_size=self.decode_block_size, + has_missing_values=not bool(finite_mask_np.all()), + ) + # quantiles: (Q, B=1, C, H) → (H, C, Q) + results.append(quantiles[:, 0].float().cpu().permute(2, 1, 0)) + return results + + def embed_batch(self, inputs): + pooler = POOLERS[self.pooler]() + device = next(self.model.parameters()).device + results = [] + for inp in inputs: + x = inp.float().cpu().numpy() # (T, C) + if self.context_length is not None: + x = x[-self.context_length:] + + # Toto expects (B, C, T) + batch = x.T[None] # (1, C, T) + mask = np.isfinite(batch) + batch = np.nan_to_num(batch, nan=0.0, posinf=0.0, neginf=0.0) + + pad_len = (-batch.shape[-1]) % self.patch_size + if pad_len: + batch = np.pad( + batch, ((0, 0), (0, 0), (pad_len, 0)), constant_values=0.0 + ) + mask = np.pad( + mask, ((0, 0), (0, 0), (pad_len, 0)), constant_values=False + ) - elif self.task == "classification": - base_encoder = _Toto2EmbedEncoder( - self._model, - self._device, - context_length=self.context_length, - layer=self.layer, - ) - encoder = Encoder(base_encoder, POOLERS[self.pooler]()) - adapter = LinearProbeAdapter( - encoder, - task="classification", - n_classes=self.meta.get("n_classes"), + target = torch.from_numpy(batch).to(device) + target_mask = torch.from_numpy(mask).to(device, dtype=torch.bool) + cpm_mask = torch.ones_like(target_mask) + series_ids = torch.zeros( + 1, target.shape[1], dtype=torch.long, device=device ) - adapter.fit(self.X_train, self.y_train) - self._adapter = adapter - elif self.task == "anomaly_detection": - self._adapter = ForecastResidualAdapter( - forecaster + captured = {} + if self.layer is None: + hook_module = self.model.transformer + else: + n = len(self.model.transformer.layers) + hook_module = self.model.transformer.layers[self.layer % n] + + def _hook(_, __, out): + captured["h"] = out.detach() + + handle = hook_module.register_forward_hook(_hook) + with torch.inference_mode(): + try: + self.model( + target=target, + target_mask=target_mask, + cpm_mask=cpm_mask, + series_ids=series_ids, + ) + finally: + handle.remove() + + # (1, C, T_patch, D) → (1, T_patch, C, D) to match pooler convention + emb_np = captured["h"].transpose(1, 2).float().cpu().numpy() + pooled = pooler.pool(emb_np) # (1, C, D) + results.append( + torch.from_numpy(pooled[0].reshape(-1).astype(np.float32)) ) - - def get_result(self): - return {"model": self._adapter} + return results diff --git a/tests/solvers/test_chronos_encoder.py b/tests/solvers/test_chronos_encoder.py deleted file mode 100644 index acee3d4..0000000 --- a/tests/solvers/test_chronos_encoder.py +++ /dev/null @@ -1,27 +0,0 @@ -"""Shape tests for the Chronos encoder variants.""" - -import chronos -import numpy as np -import pytest -import torch - -from solvers.chronos import ChronosEncoder # noqa: E402 - - -@pytest.fixture(scope="module") -def pipeline(): - return chronos.ChronosPipeline.from_pretrained( - "amazon/chronos-t5-tiny", - device_map="cpu", - torch_dtype=torch.float32, - ) - - -@pytest.mark.parametrize("layer", [None, 0, -1]) -@pytest.mark.parametrize("n_channels", [1, 3]) -def test_encode_shape(pipeline, layer, n_channels): - T = 64 - x = np.random.randn(T, n_channels).astype(np.float32) - emb = ChronosEncoder(pipeline, layer=layer).encode(x) - d_model = pipeline.model.model.config.d_model - assert emb.shape == (1, T + 1, n_channels, d_model)