diff --git a/README.md b/README.md index a75587c..f949d8b 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,84 @@ # TimesFM -PLACEHOLDER +TimesFM (Time Series Foundation Model) is a pretrained time-series foundation +model developed by Google Research for time-series forecasting. + +* Paper: + [A decoder-only foundation model for time-series forecasting](https://arxiv.org/abs/2310.10688), + ICML 2024. +* All checkpoints: + [TimesFM Hugging Face Collection](https://huggingface.co/collections/google/timesfm-release-66e4be5fdb56e960c1e482a6). +* [Google Research blog](https://research.google/blog/a-decoder-only-foundation-model-for-time-series-forecasting/). +* [TimesFM in BigQuery](https://cloud.google.com/bigquery/docs/timesfm-model): + an official Google product. + +This open version is not an officially supported Google product. + +**Latest Model Version:** TimesFM 2.5 + +**Archived Model Versions:** + +- 1.0 and 2.0: relevant code archived in the sub directory `v1`. You can `pip + install timesfm==1.3.0` to install an older version of this package to load + them. + +## Update - Sept. 15, 2025 + +TimesFM 2.5 is out! + +Comparing to TimesFM 2.0, this new 2.5 model: + +- uses 200M parameters, down from 500M. +- supports up to 16k context length, up from 2048. +- supports continuous quantile forecast up to 1k horizon via an optional 30M + quantile head. +- gets rid of the `frequency` indicator. +- has a couple of new forecasting flags. + +Along with the model upgrade we have also upgraded the inference API. This repo +will be under construction over the next few weeks to + +1. add support for an upcoming Flax version of the model (faster inference). +2. add back covariate support. +3. populate more docstrings, docs and notebook. + +### Install + +TODO(siriuz42): Package timesfm==2.0.0 and upload to PyPI . + +Run + +```shell +git clone https://github.com/google-research/timesfm.git +cd timesfm +pip install -e . +``` + +### Code Example + +```python +import numpy as np +import timesfm +model = timesfm.TimesFM_2p5_200M_torch() +model.load_checkpoint() +model.compile( + timesfm.ForecastConfig( + max_context=1024, + max_horizon=256, + normalize_inputs=True, + use_continuous_quantile_head=True, + force_flip_invariance=True, + infer_is_positive=True, + fix_quantile_crossing=True, + ) +) +point_forecast, quantile_forecast = model.forecast( + horizon=12, + inputs=[ + np.linspace(0, 1, 100), + np.sin(np.linspace(0, 20, 67)), + ], # Two dummy inputs +) +point_forecast.shape # (2, 12) +quantile_forecast.shape # (2, 12, 10): mean, then 10th to 90th quantiles. +``` diff --git a/src/timesfm/__init__.py b/src/timesfm/__init__.py index cf83c93..85177a2 100644 --- a/src/timesfm/__init__.py +++ b/src/timesfm/__init__.py @@ -14,6 +14,7 @@ """TimesFM API.""" +from .configs import ForecastConfig from .timesfm_2p5 import timesfm_2p5_torch TimesFM_2p5_200M_torch = timesfm_2p5_torch.TimesFM_2p5_200M_torch diff --git a/src/timesfm/configs.py b/src/timesfm/configs.py index 34d3975..f5ba9c5 100644 --- a/src/timesfm/configs.py +++ b/src/timesfm/configs.py @@ -20,59 +20,85 @@ from typing import Literal @dataclasses.dataclass(frozen=False) class ForecastConfig: - """Options for forecasting.""" + """Options for forecasting. - max_context: int = 0 - max_horizon: int = 0 - normalize_inputs: bool = False - window_size: int = 0 - per_core_batch_size: int = 1 - use_continuous_quantile_head: bool = False - force_flip_invariance: bool = True - infer_is_positive: bool = True - fix_quantile_crossing: bool = False - return_backcast: bool = False + Attributes: + max_context: The maximum context length. This is used by the complied decode + function at inference time during batched inference. Any input time series + with length less than max_context will be padded with zeros, and with + length greater than max_context will be truncated. + max_horizon: The maximum horizon length. This is used by the complied decode + function at inference time during batched inference. The compiled cached + decoding function will by default forecast till max_horizon. + normalize_inputs: Whether to normalize the inputs. This is useful when the + raw inputs are of extremely large or small magnitudes which may result in + numerical issues. + window_size: The window size for decomposed forecasting. + TODO(siriuz42):implement it. + per_core_batch_size: The batch size per core. Used at inference time during + batched inference when multiple GPU / TPU devices are used. + use_continuous_quantile_head: Whether to use a separate continuous quantile + head to avoid quantile collapsing. + force_flip_invariance: Whether to force flip invariance. TimesFM guarantees + that TimesFM(aX + b) = a * TimesFM(x) + b for a >= 0 by default. This flag + extends it to a < 0 as well. + infer_is_positive: Whether to guarantee nonnegativity of the output if the + input is nonnegative. + fix_quantile_crossing: Whether to fix quantile crossing. + return_backcast: Whether to return backcast. + """ + + max_context: int = 0 + max_horizon: int = 0 + normalize_inputs: bool = False + window_size: int = 0 + per_core_batch_size: int = 1 + use_continuous_quantile_head: bool = False + force_flip_invariance: bool = True + infer_is_positive: bool = True + fix_quantile_crossing: bool = False + return_backcast: bool = False @dataclasses.dataclass(frozen=True) class ResidualBlockConfig: - """Framework-agnostic config for a residual block.""" + """Framework-agnostic config for a residual block.""" - input_dims: int - hidden_dims: int - output_dims: int - use_bias: bool - activation: Literal["relu", "swish", "none"] + input_dims: int + hidden_dims: int + output_dims: int + use_bias: bool + activation: Literal["relu", "swish", "none"] @dataclasses.dataclass(frozen=True) class RandomFourierFeaturesConfig: - """Framework-agnostic config for random fourier features.""" + """Framework-agnostic config for random fourier features.""" - input_dims: int - output_dims: int - projection_stddev: float - use_bias: bool + input_dims: int + output_dims: int + projection_stddev: float + use_bias: bool @dataclasses.dataclass(frozen=True) class TransformerConfig: - """Framework-agnostic config for a transformer.""" + """Framework-agnostic config for a transformer.""" - model_dims: int - hidden_dims: int - num_heads: int - attention_norm: Literal["rms"] - feedforward_norm: Literal["rms"] - qk_norm: Literal["rms", "none"] - use_bias: bool - use_rotary_position_embeddings: bool - ff_activation: Literal["relu", "swish", "none"] + model_dims: int + hidden_dims: int + num_heads: int + attention_norm: Literal["rms"] + feedforward_norm: Literal["rms"] + qk_norm: Literal["rms", "none"] + use_bias: bool + use_rotary_position_embeddings: bool + ff_activation: Literal["relu", "swish", "none"] @dataclasses.dataclass(frozen=True) class StackedTransformersConfig: - """Framework-agnostic config for a stacked transformers.""" + """Framework-agnostic config for a stacked transformers.""" - num_layers: int - transformer: TransformerConfig + num_layers: int + transformer: TransformerConfig diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_base.py b/src/timesfm/timesfm_2p5/timesfm_2p5_base.py index dfcdb34..ba5902e 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_base.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_base.py @@ -26,161 +26,167 @@ ForecastConfig = configs.ForecastConfig def strip_leading_nans(arr): - """Removes contiguous NaN values from the beginning of a NumPy array. + """Removes contiguous NaN values from the beginning of a NumPy array. - Args: - arr: The input NumPy array. + Args: + arr: The input NumPy array. - Returns: - A new NumPy array with leading NaN values removed. - If the array is all NaNs or empty, returns an empty array. - """ + Returns: + A new NumPy array with leading NaN values removed. + If the array is all NaNs or empty, returns an empty array. + """ - isnan = np.isnan(arr) - first_valid_index = np.argmax(~isnan) - return arr[first_valid_index:] + isnan = np.isnan(arr) + first_valid_index = np.argmax(~isnan) + return arr[first_valid_index:] def linear_interpolation(arr): - """Performs linear interpolation to fill NaN values in a 1D numpy array. + """Performs linear interpolation to fill NaN values in a 1D numpy array. - Args: - arr: The 1D numpy array containing NaN values. + Args: + arr: The 1D numpy array containing NaN values. - Returns: - A new numpy array with NaN values filled using linear interpolation, - or the original array if no NaNs are present. - Returns None if the input is not a 1D array. - Returns the original array if there are no NaN values. - """ + Returns: + A new numpy array with NaN values filled using linear interpolation, + or the original array if no NaNs are present. + Returns None if the input is not a 1D array. + Returns the original array if there are no NaN values. + """ - nans = np.isnan(arr) - if not np.any(nans): # Check if there are any NaNs - return arr - - def x(z): - return z.nonzero()[0] - - nans_indices = x(nans) - non_nans_indices = x(~nans) - non_nans_values = arr[~nans] - - try: - arr[nans] = np.interp(nans_indices, non_nans_indices, non_nans_values) - except ValueError: - if non_nans_values: - mu = np.nanmean(arr) - else: - mu = 0.0 - arr = np.where(np.isfinite(arr), arr, mu) + nans = np.isnan(arr) + if not np.any(nans): # Check if there are any NaNs return arr + def x(z): + return z.nonzero()[0] + + nans_indices = x(nans) + non_nans_indices = x(~nans) + non_nans_values = arr[~nans] + + try: + arr[nans] = np.interp(nans_indices, non_nans_indices, non_nans_values) + except ValueError: + if non_nans_values: + mu = np.nanmean(arr) + else: + mu = 0.0 + arr = np.where(np.isfinite(arr), arr, mu) + return arr + @dataclasses.dataclass(frozen=True) class TimesFM_2p5_200M_Definition: - """Framework-agnostic config of TimesFM 2.5.""" + """Framework-agnostic config of TimesFM 2.5.""" - context_limit = 16384 - input_patch_len: int = 32 - output_patch_len: int = 128 - output_quantile_len: int = 1024 - quantiles: list[float] = dataclasses.field( - default_factory=lambda: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] - ) - decode_index: int = 5 - tokenizer: ResidualBlockConfig = ResidualBlockConfig( - input_dims=64, - hidden_dims=1280, - output_dims=1280, - use_bias=True, - activation="swish", - ) - stacked_transformers: StackedTransformersConfig = StackedTransformersConfig( - num_layers=20, - transformer=TransformerConfig( - model_dims=1280, - hidden_dims=1280, - num_heads=16, - attention_norm="rms", - feedforward_norm="rms", - qk_norm="rms", - use_bias=False, - use_rotary_position_embeddings=True, - ff_activation="swish", - ), - ) - output_projection_point: ResidualBlockConfig = ResidualBlockConfig( - input_dims=1280, - hidden_dims=1280, - output_dims=1280, - use_bias=False, - activation="swish", - ) - output_projection_quantiles: ResidualBlockConfig = ResidualBlockConfig( - input_dims=1280, - hidden_dims=1280, - output_dims=10240, - use_bias=False, - activation="swish", - ) + context_limit = 16384 + input_patch_len: int = 32 + output_patch_len: int = 128 + output_quantile_len: int = 1024 + quantiles: list[float] = dataclasses.field( + default_factory=lambda: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] + ) + decode_index: int = 5 + tokenizer: ResidualBlockConfig = ResidualBlockConfig( + input_dims=64, + hidden_dims=1280, + output_dims=1280, + use_bias=True, + activation="swish", + ) + stacked_transformers: StackedTransformersConfig = StackedTransformersConfig( + num_layers=20, + transformer=TransformerConfig( + model_dims=1280, + hidden_dims=1280, + num_heads=16, + attention_norm="rms", + feedforward_norm="rms", + qk_norm="rms", + use_bias=False, + use_rotary_position_embeddings=True, + ff_activation="swish", + ), + ) + output_projection_point: ResidualBlockConfig = ResidualBlockConfig( + input_dims=1280, + hidden_dims=1280, + output_dims=1280, + use_bias=False, + activation="swish", + ) + output_projection_quantiles: ResidualBlockConfig = ResidualBlockConfig( + input_dims=1280, + hidden_dims=1280, + output_dims=10240, + use_bias=False, + activation="swish", + ) class TimesFM_2p5: - """Abstract base class for TimesFM models.""" + """Abstract base class for TimesFM models. - forecast_config: ForecastConfig | None = None - compiled_decode: Callable[..., Any] | None = None - global_batch_size: int = 0 + Attributes: + forecast_config: Configuration for forecasting flags. + compiled_decode: Compiled decode function. + global_batch_size: Global batch size. + """ - def load_checkpoint(self, path: str): - """Loads a TimesFM model from a checkpoint.""" - raise NotImplementedError() + forecast_config: ForecastConfig | None = None + compiled_decode: Callable[..., Any] | None = None + global_batch_size: int = 0 - def compile(self, forecast_config: ForecastConfig | None = None): - """Compiles the TimesFM model for fast decoding.""" - raise NotImplementedError() + def load_checkpoint(self, path: str): + """Loads a TimesFM model from a checkpoint.""" + raise NotImplementedError() - def forecast( - self, horizon: int, inputs: list[np.ndarray] - ) -> tuple[np.ndarray, np.ndarray]: - """Forecasts the time series.""" - if self.compiled_decode is None: - raise RuntimeError("Model is not compiled. Please call compile() first.") + def compile(self, forecast_config: ForecastConfig | None = None): + """Compiles the TimesFM model for fast decoding.""" + raise NotImplementedError() - assert self.global_batch_size > 0 - assert self.forecast_config is not None + def forecast( + self, horizon: int, inputs: list[np.ndarray] + ) -> tuple[np.ndarray, np.ndarray]: + """Forecasts the time series.""" + if self.compiled_decode is None: + raise RuntimeError("Model is not compiled. Please call compile() first.") - context = self.forecast_config.max_context - num_inputs = len(inputs) - if (w := num_inputs % self.global_batch_size) != 0: - inputs += [np.array([0.0] * 3)] * (self.global_batch_size - w) + assert self.global_batch_size > 0 + assert self.forecast_config is not None - output_points = [] - output_quantiles = [] + context = self.forecast_config.max_context + num_inputs = len(inputs) + if (w := num_inputs % self.global_batch_size) != 0: + inputs += [np.array([0.0] * 3)] * (self.global_batch_size - w) + + output_points = [] + output_quantiles = [] + values = [] + masks = [] + idx = 0 + for each_input in inputs: + value = linear_interpolation(strip_leading_nans(np.array(each_input))) + if (w := len(value)) >= context: + value = value[-context:] + mask = np.zeros_like(value, dtype=bool) + else: + mask = np.array([True] * (context - w) + [False] * w) + value = np.pad(value, (context - w, 0), "constant", constant_values=0.0) + values.append(value) + masks.append(mask) + idx += 1 + if idx == self.global_batch_size: + idx = 0 + point_forecast, quantile_forecast = self.compiled_decode( + horizon, values, masks + ) + output_points.append(point_forecast) + output_quantiles.append(quantile_forecast) values = [] masks = [] - idx = 0 - for each_input in inputs: - value = linear_interpolation(strip_leading_nans(np.array(each_input))) - if (w := len(value)) >= context: - value = value[-context:] - mask = np.zeros_like(value, dtype=bool) - else: - mask = np.array([True] * (context - w) + [False] * w) - value = np.pad(value, (context - w, 0), "constant", constant_values=0.0) - values.append(value) - masks.append(mask) - idx += 1 - if idx == self.global_batch_size: - idx = 0 - point_forecast, quantile_forecast = self.compiled_decode( - horizon, values, masks - ) - output_points.append(point_forecast) - output_quantiles.append(quantile_forecast) - values = [] - masks = [] - output_points = np.concatenate(output_points, axis=0) - output_quantiles = np.concatenate(output_quantiles, axis=0) - return output_points[:num_inputs], output_quantiles[:num_inputs] + output_points = np.concatenate(output_points, axis=0) + output_quantiles = np.concatenate(output_quantiles, axis=0) + return output_points[:num_inputs], output_quantiles[:num_inputs] diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py index 3855d9d..dfb21aa 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py @@ -35,384 +35,414 @@ revin = util.revin class TimesFM_2p5_200M_torch_module(nn.Module): - """TimesFM 2.5 with 200M parameters.""" + """TimesFM 2.5 with 200M parameters.""" - config = timesfm_2p5_base.TimesFM_2p5_200M_Definition() + config = timesfm_2p5_base.TimesFM_2p5_200M_Definition() - def __init__(self): - super().__init__() + def __init__(self): + super().__init__() - # Names constants. - self.p = self.config.input_patch_len # 32 - self.o = self.config.output_patch_len # 128 - self.os = self.config.output_quantile_len # 1024 - self.m = self.o // self.p # 4 - self.x = self.config.stacked_transformers.num_layers # 20 - self.h = self.config.stacked_transformers.transformer.num_heads # 16 - self.md = self.config.stacked_transformers.transformer.model_dims # 1280 - self.hd = self.md // self.h # 80 - self.q = len(self.config.quantiles) + 1 # 10 - self.aridx = self.config.decode_index # 5 + # Names constants. + self.p = self.config.input_patch_len # 32 + self.o = self.config.output_patch_len # 128 + self.os = self.config.output_quantile_len # 1024 + self.m = self.o // self.p # 4 + self.x = self.config.stacked_transformers.num_layers # 20 + self.h = self.config.stacked_transformers.transformer.num_heads # 16 + self.md = self.config.stacked_transformers.transformer.model_dims # 1280 + self.hd = self.md // self.h # 80 + self.q = len(self.config.quantiles) + 1 # 10 + self.aridx = self.config.decode_index # 5 - # Layers. - self.tokenizer = dense.ResidualBlock(self.config.tokenizer) - self.stacked_xf = nn.ModuleList( - [ - transformer.Transformer(self.config.stacked_transformers.transformer) - for _ in range(self.x) - ] + # Layers. + self.tokenizer = dense.ResidualBlock(self.config.tokenizer) + self.stacked_xf = nn.ModuleList([ + transformer.Transformer(self.config.stacked_transformers.transformer) + for _ in range(self.x) + ]) + self.output_projection_point = dense.ResidualBlock( + self.config.output_projection_point + ) + self.output_projection_quantiles = dense.ResidualBlock( + self.config.output_projection_quantiles + ) + + # Device. + if torch.cuda.is_available(): + self.device = torch.device("cuda:0") + self.device_count = torch.cuda.device_count() + else: + self.device = torch.device("cpu") + self.device_count = 1 + + def load_checkpoint(self, path: str): + """Loads a PyTorch TimesFM model from a checkpoint.""" + tensors = load_file(path) + self.load_state_dict(tensors) + self.to(self.device) + + def forward( + self, + inputs: torch.Tensor, + masks: torch.Tensor, + decode_caches: list[util.DecodeCache] | None = None, + ): + tokenizer_inputs = torch.cat([inputs, masks.to(inputs.dtype)], dim=-1) + input_embeddings = self.tokenizer(tokenizer_inputs) + + if decode_caches is None: + decode_caches = [None] * self.x + + output_embeddings = input_embeddings + new_decode_caches = [] + for i, layer in enumerate(self.stacked_xf): + output_embeddings, new_cache = layer( + output_embeddings, masks[..., -1], decode_caches[i] + ) + new_decode_caches.append(new_cache) + output_ts = self.output_projection_point(output_embeddings) + output_quantile_spread = self.output_projection_quantiles(output_embeddings) + + return ( + input_embeddings, + output_embeddings, + output_ts, + output_quantile_spread, + ), new_decode_caches + + def decode(self, horizon: int, inputs, masks): + """Decodes the time series.""" + + inputs = inputs.to(self.device) + masks = masks.to(self.device) + + with torch.no_grad(): + batch_size, context = inputs.shape[0], inputs.shape[1] + num_decode_steps = (horizon - 1) // self.o + num_input_patches = context // self.p + decode_cache_size = num_input_patches + num_decode_steps * self.m + + # Prefill + patched_inputs = torch.reshape(inputs, (batch_size, -1, self.p)) + patched_masks = torch.reshape(masks, (batch_size, -1, self.p)) + + # running stats + n = torch.zeros(batch_size, device=inputs.device) + mu = torch.zeros(batch_size, device=inputs.device) + sigma = torch.zeros(batch_size, device=inputs.device) + patch_mu = [] + patch_sigma = [] + for i in range(num_input_patches): + (n, mu, sigma), _ = util.update_running_stats( + n, mu, sigma, patched_inputs[:, i], patched_masks[:, i] ) - self.output_projection_point = dense.ResidualBlock( - self.config.output_projection_point + patch_mu.append(mu) + patch_sigma.append(sigma) + last_n, last_mu, last_sigma = n, mu, sigma + context_mu = torch.stack(patch_mu, dim=1) + context_sigma = torch.stack(patch_sigma, dim=1) + + decode_caches = [ + util.DecodeCache( + next_index=torch.zeros( + batch_size, dtype=torch.int32, device=inputs.device + ), + num_masked=torch.zeros( + batch_size, dtype=torch.int32, device=inputs.device + ), + key=torch.zeros( + batch_size, + decode_cache_size, + self.h, + self.hd, + device=inputs.device, + ), + value=torch.zeros( + batch_size, + decode_cache_size, + self.h, + self.hd, + device=inputs.device, + ), + ) + for _ in range(self.x) + ] + + normed_inputs = revin( + patched_inputs, context_mu, context_sigma, reverse=False + ) + normed_inputs = torch.where(patched_masks, 0.0, normed_inputs) + (_, _, normed_outputs, normed_quantile_spread), decode_caches = self( + normed_inputs, patched_masks, decode_caches + ) + renormed_outputs = torch.reshape( + revin(normed_outputs, context_mu, context_sigma, reverse=True), + (batch_size, -1, self.o, self.q), + ) + renormed_quantile_spread = torch.reshape( + revin( + normed_quantile_spread, context_mu, context_sigma, reverse=True + ), + (batch_size, -1, self.os, self.q), + )[:, -1, ...] + + # Autogressive decode + ar_outputs = [] + last_renormed_output = renormed_outputs[:, -1, :, self.aridx] + + for _ in range(num_decode_steps): + new_patched_input = torch.reshape( + last_renormed_output, (batch_size, self.m, self.p) ) - self.output_projection_quantiles = dense.ResidualBlock( - self.config.output_projection_quantiles + new_mask = torch.zeros_like(new_patched_input, dtype=torch.bool) + + n, mu, sigma = last_n, last_mu, last_sigma + new_mus, new_sigmas = [], [] + for i in range(self.m): + (n, mu, sigma), _ = util.update_running_stats( + n, mu, sigma, new_patched_input[:, i], new_mask[:, i] + ) + new_mus.append(mu) + new_sigmas.append(sigma) + last_n, last_mu, last_sigma = n, mu, sigma + new_mu = torch.stack(new_mus, dim=1) + new_sigma = torch.stack(new_sigmas, dim=1) + + new_normed_input = revin( + new_patched_input, new_mu, new_sigma, reverse=False + ) + (_, _, new_normed_output, _), decode_caches = self( + new_normed_input, new_mask, decode_caches ) - # Device. - if torch.cuda.is_available(): - self.device = torch.device("cuda:0") - self.device_count = torch.cuda.device_count() - else: - self.device = torch.device("cpu") - self.device_count = 1 + new_renormed_output = torch.reshape( + revin(new_normed_output, new_mu, new_sigma, reverse=True), + (batch_size, self.m, self.o, self.q), + ) + ar_outputs.append(new_renormed_output[:, -1, ...]) + last_renormed_output = new_renormed_output[:, -1, :, self.aridx] - def load_checkpoint(self, path: str): - """Loads a PyTorch TimesFM model from a checkpoint.""" - tensors = load_file(path) - self.load_state_dict(tensors) - self.to(self.device) + if num_decode_steps > 0: + ar_renormed_outputs = torch.stack(ar_outputs, dim=1) + else: + ar_renormed_outputs = None - def forward( - self, - inputs: torch.Tensor, - masks: torch.Tensor, - decode_caches: list[util.DecodeCache] | None = None, - ): - tokenizer_inputs = torch.cat([inputs, masks.to(inputs.dtype)], dim=-1) - input_embeddings = self.tokenizer(tokenizer_inputs) + return renormed_outputs, renormed_quantile_spread, ar_renormed_outputs - if decode_caches is None: - decode_caches = [None] * self.x + def forecast_naive( + self, horizon: int, inputs: Sequence[np.ndarray] + ) -> list[np.ndarray]: + """Forecasts the time series. - output_embeddings = input_embeddings - new_decode_caches = [] - for i, layer in enumerate(self.stacked_xf): - output_embeddings, new_cache = layer( - output_embeddings, masks[..., -1], decode_caches[i] - ) - new_decode_caches.append(new_cache) - output_ts = self.output_projection_point(output_embeddings) - output_quantile_spread = self.output_projection_quantiles(output_embeddings) + This is a naive implementation for debugging purposes. No forecasting + flags are used here. Forecasting quality can be subpar. - return ( - input_embeddings, - output_embeddings, - output_ts, - output_quantile_spread, - ), new_decode_caches + Args: + horizon: The number of time points to forecast. + inputs: A sequence of numpy arrays, each representing a time series to + query forecast for. - def decode(self, horizon: int, inputs, masks): - """Decodes the time series.""" - - inputs = inputs.to(self.device) - masks = masks.to(self.device) - - with torch.no_grad(): - batch_size, context = inputs.shape[0], inputs.shape[1] - num_decode_steps = (horizon - 1) // self.o - num_input_patches = context // self.p - decode_cache_size = num_input_patches + num_decode_steps * self.m - - # Prefill - patched_inputs = torch.reshape(inputs, (batch_size, -1, self.p)) - patched_masks = torch.reshape(masks, (batch_size, -1, self.p)) - - # running stats - n = torch.zeros(batch_size, device=inputs.device) - mu = torch.zeros(batch_size, device=inputs.device) - sigma = torch.zeros(batch_size, device=inputs.device) - patch_mu = [] - patch_sigma = [] - for i in range(num_input_patches): - (n, mu, sigma), _ = util.update_running_stats( - n, mu, sigma, patched_inputs[:, i], patched_masks[:, i] - ) - patch_mu.append(mu) - patch_sigma.append(sigma) - last_n, last_mu, last_sigma = n, mu, sigma - context_mu = torch.stack(patch_mu, dim=1) - context_sigma = torch.stack(patch_sigma, dim=1) - - decode_caches = [ - util.DecodeCache( - next_index=torch.zeros( - batch_size, dtype=torch.int32, device=inputs.device - ), - num_masked=torch.zeros( - batch_size, dtype=torch.int32, device=inputs.device - ), - key=torch.zeros( - batch_size, - decode_cache_size, - self.h, - self.hd, - device=inputs.device, - ), - value=torch.zeros( - batch_size, - decode_cache_size, - self.h, - self.hd, - device=inputs.device, - ), - ) - for _ in range(self.x) - ] - - normed_inputs = revin( - patched_inputs, context_mu, context_sigma, reverse=False - ) - normed_inputs = torch.where(patched_masks, 0.0, normed_inputs) - (_, _, normed_outputs, normed_quantile_spread), decode_caches = self( - normed_inputs, patched_masks, decode_caches - ) - renormed_outputs = torch.reshape( - revin(normed_outputs, context_mu, context_sigma, reverse=True), - (batch_size, -1, self.o, self.q), - ) - renormed_quantile_spread = torch.reshape( - revin(normed_quantile_spread, context_mu, context_sigma, reverse=True), - (batch_size, -1, self.os, self.q), - )[:, -1, ...] - - # Autogressive decode - ar_outputs = [] - last_renormed_output = renormed_outputs[:, -1, :, self.aridx] - - for _ in range(num_decode_steps): - new_patched_input = torch.reshape( - last_renormed_output, (batch_size, self.m, self.p) - ) - new_mask = torch.zeros_like(new_patched_input, dtype=torch.bool) - - n, mu, sigma = last_n, last_mu, last_sigma - new_mus, new_sigmas = [], [] - for i in range(self.m): - (n, mu, sigma), _ = util.update_running_stats( - n, mu, sigma, new_patched_input[:, i], new_mask[:, i] - ) - new_mus.append(mu) - new_sigmas.append(sigma) - last_n, last_mu, last_sigma = n, mu, sigma - new_mu = torch.stack(new_mus, dim=1) - new_sigma = torch.stack(new_sigmas, dim=1) - - new_normed_input = revin( - new_patched_input, new_mu, new_sigma, reverse=False - ) - (_, _, new_normed_output, _), decode_caches = self( - new_normed_input, new_mask, decode_caches - ) - - new_renormed_output = torch.reshape( - revin(new_normed_output, new_mu, new_sigma, reverse=True), - (batch_size, self.m, self.o, self.q), - ) - ar_outputs.append(new_renormed_output[:, -1, ...]) - last_renormed_output = new_renormed_output[:, -1, :, self.aridx] - - if num_decode_steps > 0: - ar_renormed_outputs = torch.stack(ar_outputs, dim=1) - else: - ar_renormed_outputs = None - - return renormed_outputs, renormed_quantile_spread, ar_renormed_outputs - - def forecast_naive(self, horizon: int, inputs: Sequence[np.ndarray]): - """Forecasts the time series.""" - outputs = [] - for each_input in inputs: - input_t = torch.tensor(each_input, dtype=torch.float32) - mask = torch.zeros_like(input_t, dtype=torch.bool) - len_front_mask = self.p - (len(each_input) % self.p) - if len_front_mask < self.p: - input_t = torch.cat( - [torch.zeros(len_front_mask, dtype=torch.float32), input_t], dim=0 - ) - mask = torch.cat( - [torch.ones(len_front_mask, dtype=torch.bool), mask], dim=0 - ) - input_t = input_t[None, ...] - mask = mask[None, ...] - t_pf, _, t_ar = self.decode(horizon, input_t, mask) - to_concat = [t_pf[:, -1, ...]] - if t_ar is not None: - to_concat.append(t_ar.reshape(1, -1, self.q)) - torch_forecast = torch.cat(to_concat, dim=1)[..., :horizon] - torch_forecast = torch_forecast.squeeze(0) - outputs.append(torch_forecast.detach().cpu().numpy()) - return outputs + Returns: + A list of numpy arrays of forecasts. + """ + outputs = [] + for each_input in inputs: + input_t = torch.tensor(each_input, dtype=torch.float32) + mask = torch.zeros_like(input_t, dtype=torch.bool) + len_front_mask = self.p - (len(each_input) % self.p) + if len_front_mask < self.p: + input_t = torch.cat( + [torch.zeros(len_front_mask, dtype=torch.float32), input_t], dim=0 + ) + mask = torch.cat( + [torch.ones(len_front_mask, dtype=torch.bool), mask], dim=0 + ) + input_t = input_t[None, ...] + mask = mask[None, ...] + t_pf, _, t_ar = self.decode(horizon, input_t, mask) + to_concat = [t_pf[:, -1, ...]] + if t_ar is not None: + to_concat.append(t_ar.reshape(1, -1, self.q)) + torch_forecast = torch.cat(to_concat, dim=1)[..., :horizon] + torch_forecast = torch_forecast.squeeze(0) + outputs.append(torch_forecast.detach().cpu().numpy()) + return outputs class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5): - """PyTorch implementation of TimesFM 2.5 with 200M parameters.""" + """PyTorch implementation of TimesFM 2.5 with 200M parameters.""" - model: nn.Module = TimesFM_2p5_200M_torch_module() + model: nn.Module = TimesFM_2p5_200M_torch_module() - def load_checkpoint( - self, - *, - path: str | None = None, - hf_repo_id: str | None = "google/timesfm-2.5-200m-pytorch", - ): - """Loads a PyTorch safetensors TimesFM model.""" - if path: - pass - elif hf_repo_id: - logging.info("Downloading checkpoint from Hugging Face repo %s", hf_repo_id) - path = os.path.join( - huggingface_hub.snapshot_download(hf_repo_id), "model.safetensors" - ) - logging.info("Loading checkpoint from: %s", path) - else: - raise ValueError("Either path or hf_repo_id must be provided.") - self.model.load_checkpoint(path) + def load_checkpoint( + self, + *, + path: str | None = None, + hf_repo_id: str | None = "google/timesfm-2.5-200m-pytorch", + ) -> None: + """Loads a PyTorch safetensors TimesFM model. - def compile(self, forecast_config: configs.ForecastConfig, **kwargs): - if kwargs.get("backend", None) is not None: - self.model.compile(**kwargs) - self.global_batch_size = ( - forecast_config.per_core_batch_size * self.model.device_count + Args: + path: Path to a local checkpoint. If not provided, will try to download + from the default Hugging Face repo. + hf_repo_id: If provided, will download from the specified Hugging Face + repo instead. + """ + if path: + pass + elif hf_repo_id: + logging.info( + "Downloading checkpoint from Hugging Face repo %s", hf_repo_id + ) + path = os.path.join( + huggingface_hub.snapshot_download(hf_repo_id), "model.safetensors" + ) + logging.info("Loading checkpoint from: %s", path) + else: + raise ValueError("Either path or hf_repo_id must be provided.") + self.model.load_checkpoint(path) + + def compile(self, forecast_config: configs.ForecastConfig, **kwargs) -> None: + """Attempts to compile the model for fast decoding. + + See configs.ForecastConfig for more details on the supported flags. + + Args: + forecast_config: Configuration for forecasting flags. + **kwargs: Additional keyword arguments to pass to model.compile(). + """ + + if kwargs.get("backend", None) is not None: + self.model.compile(**kwargs) + self.global_batch_size = ( + forecast_config.per_core_batch_size * self.model.device_count + ) + + # Shortcut. + fc = forecast_config + + if fc.max_context % self.model.p != 0: + logging.info( + "When compiling, max context needs to be multiple of the patch size" + " %d. Using max context = %d instead.", + self.model.p, + new_context := math.ceil(fc.max_context / self.model.p) + * self.model.p, + ) + fc.max_context = new_context + if fc.max_horizon % self.model.o != 0: + logging.info( + "When compiling, max horizon needs to be multiple of the output patch" + " size %d. Using max horizon = %d instead.", + self.model.o, + new_horizon := math.ceil(fc.max_horizon / self.model.o) + * self.model.o, + ) + fc.max_horizon = new_horizon + if fc.max_context + fc.max_horizon > self.model.config.context_limit: + raise ValueError( + "Context + horizon must be less than the context limit." + f" {fc.max_context} + {fc.max_horizon} >" + f" {self.model.config.context_limit}." + ) + if fc.use_continuous_quantile_head and (fc.max_horizon > self.model.os): + raise ValueError( + "Continuous quantile head is not supported for horizons >" + f" {self.model.os}." + ) + self.forecast_config = fc + + def _compiled_decode(horizon, inputs, masks): + if horizon > fc.max_horizon: + raise ValueError( + "Horizon must be less than the max horizon." + f" {horizon} > {fc.max_horizon}." ) - # Shortcut. - fc = forecast_config + inputs = torch.Tensor(np.array(inputs)).to(self.model.device) + masks = torch.Tensor(np.array(masks)).to(self.model.device).to(torch.bool) + batch_size = inputs.shape[0] - if fc.max_context % self.model.p != 0: - logging.info( - "When compiling, max context needs to be multiple of the patch size" - " %d. Using max context = %d instead.", - self.model.p, - new_context := math.ceil(fc.max_context / self.model.p) * self.model.p, - ) - fc.max_context = new_context - if fc.max_horizon % self.model.o != 0: - logging.info( - "When compiling, max horizon needs to be multiple of the output patch" - " size %d. Using max horizon = %d instead.", - self.model.o, - new_horizon := math.ceil(fc.max_horizon / self.model.o) * self.model.o, - ) - fc.max_horizon = new_horizon - if fc.max_context + fc.max_horizon > self.model.config.context_limit: - raise ValueError( - "Context + horizon must be less than the context limit." - f" {fc.max_context} + {fc.max_horizon} >" - f" {self.model.config.context_limit}." - ) - if fc.use_continuous_quantile_head and (fc.max_horizon > self.model.os): - raise ValueError( - "Continuous quantile head is not supported for horizons >" - f" {self.model.os}." - ) - self.forecast_config = fc + if fc.infer_is_positive: + is_positive = torch.all(inputs >= 0, dim=-1, keepdim=True) + else: + is_positive = None - def _compiled_decode(horizon, inputs, masks): - if horizon > fc.max_horizon: - raise ValueError( - "Horizon must be less than the max horizon." - f" {horizon} > {fc.max_horizon}." - ) + if fc.normalize_inputs: + mu = torch.mean(inputs, dim=-1, keepdim=True) + sigma = torch.std(inputs, dim=-1, keepdim=True) + inputs = revin(inputs, mu, sigma, reverse=False) + else: + mu, sigma = None, None - inputs = torch.Tensor(np.array(inputs)).to(self.model.device) - masks = torch.Tensor(np.array(masks)).to(self.model.device).to(torch.bool) - batch_size = inputs.shape[0] + pf_outputs, quantile_spreads, ar_outputs = self.model.decode( + forecast_config.max_horizon, inputs, masks + ) + to_cat = [pf_outputs[:, -1, ...]] + if ar_outputs is not None: + to_cat.append(ar_outputs.reshape(batch_size, -1, self.model.q)) + full_forecast = torch.cat(to_cat, dim=1) - if fc.infer_is_positive: - is_positive = torch.all(inputs >= 0, dim=-1, keepdim=True) - else: - is_positive = None + flip_quantile_fn = lambda x: torch.cat( + [x[..., :1], torch.flip(x[..., 1:], dims=(-1,))], dim=-1 + ) - if fc.normalize_inputs: - mu = torch.mean(inputs, dim=-1, keepdim=True) - sigma = torch.std(inputs, dim=-1, keepdim=True) - inputs = revin(inputs, mu, sigma, reverse=False) - else: - mu, sigma = None, None + if fc.force_flip_invariance: + flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = ( + self.model.decode(forecast_config.max_horizon, -inputs, masks) + ) + flipped_quantile_spreads = flip_quantile_fn(flipped_quantile_spreads) + flipped_pf_outputs = flip_quantile_fn(flipped_pf_outputs) + to_cat = [flipped_pf_outputs[:, -1, ...]] + if flipped_ar_outputs is not None: + to_cat.append( + flipped_ar_outputs.reshape(batch_size, -1, self.model.q) + ) + flipped_full_forecast = torch.cat(to_cat, dim=1) + quantile_spreads = (quantile_spreads - flipped_quantile_spreads) / 2 + pf_outputs = (pf_outputs - flipped_pf_outputs) / 2 + full_forecast = (full_forecast - flipped_full_forecast) / 2 - pf_outputs, quantile_spreads, ar_outputs = self.model.decode( - forecast_config.max_horizon, inputs, masks - ) - full_forecast = torch.cat( - [ - pf_outputs[:, -1, ...], - ar_outputs.reshape(batch_size, -1, self.model.q), - ], - dim=1, - ) + if fc.use_continuous_quantile_head: + for quantile_index in [1, 2, 3, 4, 6, 7, 8, 9]: + full_forecast[:, :, quantile_index] = ( + quantile_spreads[:, : fc.max_horizon, quantile_index] + - quantile_spreads[:, : fc.max_horizon, 5] + + full_forecast[:, : fc.max_horizon, 5] + ) + full_forecast = full_forecast[:, :horizon, :] - flip_quantile_fn = lambda x: torch.cat( - [x[..., :1], torch.flip(x[..., 1:], dims=(-1,))], dim=-1 - ) + if fc.return_backcast: + full_backcast = pf_outputs[:, :-1, : self.model.p, :].reshape( + batch_size, -1, self.model.q + ) + full_forecast = torch.cat([full_backcast, full_forecast], dim=1) - if fc.force_flip_invariance: - flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = ( - self.model.decode(forecast_config.max_horizon, -inputs, masks) - ) - flipped_quantile_spreads = flip_quantile_fn(flipped_quantile_spreads) - flipped_pf_outputs = flip_quantile_fn(flipped_pf_outputs) - flipped_full_forecast = torch.cat( - [ - flipped_pf_outputs[:, -1, ...], - flipped_ar_outputs.reshape(batch_size, -1, self.model.q), - ], - dim=1, - ) - quantile_spreads = (quantile_spreads - flipped_quantile_spreads) / 2 - pf_outputs = (pf_outputs - flipped_pf_outputs) / 2 - full_forecast = (full_forecast - flipped_full_forecast) / 2 + if fc.fix_quantile_crossing: + for i in [4, 3, 2, 1]: + full_forecast[:, :, i] = torch.where( + full_forecast[:, :, i] < full_forecast[:, :, i + 1], + full_forecast[:, :, i], + full_forecast[:, :, i + 1], + ) + for i in [6, 7, 8, 9]: + full_forecast[:, :, i] = torch.where( + full_forecast[:, :, i] > full_forecast[:, :, i - 1], + full_forecast[:, :, i], + full_forecast[:, :, i - 1], + ) - if fc.use_continuous_quantile_head: - for quantile_index in [1, 2, 3, 4, 6, 7, 8, 9]: - full_forecast[:, :, quantile_index] = ( - quantile_spreads[:, : fc.max_horizon, quantile_index] - - quantile_spreads[:, : fc.max_horizon, 5] - + full_forecast[:, : fc.max_horizon, 5] - ) - full_forecast = full_forecast[:, :horizon, :] + if fc.normalize_inputs: + full_forecast = revin(full_forecast, mu, sigma, reverse=True) - if fc.return_backcast: - full_backcast = pf_outputs[:, :-1, : self.model.p, :].reshape( - batch_size, -1, self.model.q - ) - full_forecast = torch.cat([full_backcast, full_forecast], dim=1) + if is_positive is not None: + full_forecast = torch.where( + is_positive[..., None], + torch.maximum(full_forecast, torch.zeros_like(full_forecast)), + full_forecast, + ) - if fc.fix_quantile_crossing: - for i in [4, 3, 2, 1]: - full_forecast[:, :, i] = torch.where( - full_forecast[:, :, i] < full_forecast[:, :, i + 1], - full_forecast[:, :, i], - full_forecast[:, :, i + 1], - ) - for i in [6, 7, 8, 9]: - full_forecast[:, :, i] = torch.where( - full_forecast[:, :, i] > full_forecast[:, :, i - 1], - full_forecast[:, :, i], - full_forecast[:, :, i - 1], - ) + full_forecast = full_forecast.detach().cpu().numpy() + return full_forecast[..., 5], full_forecast - if fc.normalize_inputs: - full_forecast = revin(full_forecast, mu, sigma, reverse=True) - - if is_positive is not None: - full_forecast = torch.where( - is_positive[..., None], - torch.maximum(full_forecast, torch.zeros_like(full_forecast)), - full_forecast, - ) - - full_forecast = full_forecast.detach().cpu().numpy() - return full_forecast[..., 5], full_forecast - - self.compiled_decode = _compiled_decode + self.compiled_decode = _compiled_decode diff --git a/src/timesfm/torch/dense.py b/src/timesfm/torch/dense.py index 9ebcfcd..bcc0e5c 100644 --- a/src/timesfm/torch/dense.py +++ b/src/timesfm/torch/dense.py @@ -21,74 +21,74 @@ from .. import configs class ResidualBlock(nn.Module): - """Residual block with two linear layers and a linear residual connection.""" + """Residual block with two linear layers and a linear residual connection.""" - def __init__(self, config: configs.ResidualBlockConfig): - super().__init__() - self.config = config - self.hidden_layer = nn.Linear( - in_features=config.input_dims, - out_features=config.hidden_dims, - bias=config.use_bias, - ) - self.output_layer = nn.Linear( - in_features=config.hidden_dims, - out_features=config.output_dims, - bias=config.use_bias, - ) - self.residual_layer = nn.Linear( - in_features=config.input_dims, - out_features=config.output_dims, - bias=config.use_bias, - ) - if config.activation == "relu": - self.activation = nn.ReLU() - elif config.activation == "swish": - self.activation = nn.SiLU() - elif config.activation == "none": - self.activation = nn.Identity() - else: - raise ValueError(f"Activation: {config.activation} not supported.") + def __init__(self, config: configs.ResidualBlockConfig): + super().__init__() + self.config = config + self.hidden_layer = nn.Linear( + in_features=config.input_dims, + out_features=config.hidden_dims, + bias=config.use_bias, + ) + self.output_layer = nn.Linear( + in_features=config.hidden_dims, + out_features=config.output_dims, + bias=config.use_bias, + ) + self.residual_layer = nn.Linear( + in_features=config.input_dims, + out_features=config.output_dims, + bias=config.use_bias, + ) + if config.activation == "relu": + self.activation = nn.ReLU() + elif config.activation == "swish": + self.activation = nn.SiLU() + elif config.activation == "none": + self.activation = nn.Identity() + else: + raise ValueError(f"Activation: {config.activation} not supported.") - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.output_layer( - self.activation(self.hidden_layer(x)) - ) + self.residual_layer(x) + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.output_layer( + self.activation(self.hidden_layer(x)) + ) + self.residual_layer(x) class RandomFourierFeatures(nn.Module): - """Random Fourier features layer.""" + """Random Fourier features layer.""" - def __init__(self, config: configs.RandomFourierFeaturesConfig): - super().__init__() - self.config = config + def __init__(self, config: configs.RandomFourierFeaturesConfig): + super().__init__() + self.config = config - if config.output_dims % 4 != 0: - raise ValueError( - f"Output dims must be a multiple of 4: {config.output_dims} % 4 != 0." - ) - num_projected_features = config.output_dims // 4 + if config.output_dims % 4 != 0: + raise ValueError( + f"Output dims must be a multiple of 4: {config.output_dims} % 4 != 0." + ) + num_projected_features = config.output_dims // 4 - self.phase_shifts = nn.Parameter(torch.zeros(2, num_projected_features)) - self.projection_layer = nn.Linear( - in_features=config.input_dims, - out_features=num_projected_features, - bias=config.use_bias, - ) - self.residual_layer = nn.Linear( - in_features=config.input_dims, - out_features=config.output_dims, - bias=config.use_bias, - ) + self.phase_shifts = nn.Parameter(torch.zeros(2, num_projected_features)) + self.projection_layer = nn.Linear( + in_features=config.input_dims, + out_features=num_projected_features, + bias=config.use_bias, + ) + self.residual_layer = nn.Linear( + in_features=config.input_dims, + out_features=config.output_dims, + bias=config.use_bias, + ) - def forward(self, x: torch.Tensor) -> torch.Tensor: - projected = self.projection_layer(x) - cos_features = torch.cos(projected) - sin_features = torch.sin(projected) - sq_wave_1 = torch.sign(torch.sin(projected + self.phase_shifts[0, :])) - sq_wave_2 = torch.sign(torch.sin(projected + self.phase_shifts[1, :])) - fourier_features = torch.cat( - [cos_features, sin_features, sq_wave_1, sq_wave_2], dim=-1 - ) - residual = self.residual_layer(x) - return fourier_features + residual + def forward(self, x: torch.Tensor) -> torch.Tensor: + projected = self.projection_layer(x) + cos_features = torch.cos(projected) + sin_features = torch.sin(projected) + sq_wave_1 = torch.sign(torch.sin(projected + self.phase_shifts[0, :])) + sq_wave_2 = torch.sign(torch.sin(projected + self.phase_shifts[1, :])) + fourier_features = torch.cat( + [cos_features, sin_features, sq_wave_1, sq_wave_2], dim=-1 + ) + residual = self.residual_layer(x) + return fourier_features + residual diff --git a/src/timesfm/torch/normalization.py b/src/timesfm/torch/normalization.py index 6836423..409eee8 100644 --- a/src/timesfm/torch/normalization.py +++ b/src/timesfm/torch/normalization.py @@ -19,21 +19,21 @@ from torch import nn class RMSNorm(nn.Module): - """RMS normalization.""" + """RMS normalization.""" - def __init__( - self, - num_features: int, - *, - epsilon: float = 1e-6, - ): - super().__init__() - self.scale = nn.Parameter(torch.zeros(num_features)) - self.num_features = num_features - self.epsilon = epsilon + def __init__( + self, + num_features: int, + *, + epsilon: float = 1e-6, + ): + super().__init__() + self.scale = nn.Parameter(torch.zeros(num_features)) + self.num_features = num_features + self.epsilon = epsilon - def forward(self, inputs: torch.Tensor) -> torch.Tensor: - var = torch.mean(torch.square(inputs), dim=-1, keepdim=True) - normed_inputs = inputs * torch.rsqrt(var + self.epsilon) - normed_inputs = normed_inputs * self.scale - return normed_inputs + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + var = torch.mean(torch.square(inputs), dim=-1, keepdim=True) + normed_inputs = inputs * torch.rsqrt(var + self.epsilon) + normed_inputs = normed_inputs * self.scale + return normed_inputs diff --git a/src/timesfm/torch/transformer.py b/src/timesfm/torch/transformer.py index 326002a..df7c6d0 100644 --- a/src/timesfm/torch/transformer.py +++ b/src/timesfm/torch/transformer.py @@ -36,80 +36,81 @@ def make_attn_mask( query_index_offset: torch.Tensor | None = None, kv_length: int = 0, ) -> torch.Tensor: - """Makes attention mask.""" - if kv_length == 0: - kv_length = query_length + """Makes attention mask.""" + if kv_length == 0: + kv_length = query_length - q_index = torch.arange(query_length, device=num_all_masked_kv.device)[ - None, None, :, None - ] - if query_index_offset is not None: - q_index = q_index + query_index_offset[:, None, None, None] - kv_index = torch.arange(kv_length, device=num_all_masked_kv.device)[ - None, None, None, : - ] - return torch.logical_and( - q_index >= kv_index, - kv_index >= num_all_masked_kv[:, None, None, None], - ) + q_index = torch.arange(query_length, device=num_all_masked_kv.device)[ + None, None, :, None + ] + if query_index_offset is not None: + q_index = q_index + query_index_offset[:, None, None, None] + kv_index = torch.arange(kv_length, device=num_all_masked_kv.device)[ + None, None, None, : + ] + return torch.logical_and( + q_index >= kv_index, + kv_index >= num_all_masked_kv[:, None, None, None], + ) class RotaryPositionalEmbedding(nn.Module): - """Rotary positional embedding.""" + """Rotary positional embedding.""" - def __init__( - self, - embedding_dims: int, - min_timescale: float = 1.0, - max_timescale: float = 10000.0, - ): - super().__init__() - self.embedding_dims = embedding_dims - self.min_timescale = min_timescale - self.max_timescale = max_timescale + def __init__( + self, + embedding_dims: int, + min_timescale: float = 1.0, + max_timescale: float = 10000.0, + ): + super().__init__() + self.embedding_dims = embedding_dims + self.min_timescale = min_timescale + self.max_timescale = max_timescale - def forward( - self, - inputs: torch.Tensor, - position: torch.Tensor | None = None, - ): - """Generates a JTensor of sinusoids with different frequencies.""" - if self.embedding_dims != inputs.shape[-1]: - raise ValueError( - "The embedding dims of the rotary position embedding" - "must match the hidden dimension of the inputs." - ) - half_embedding_dim = self.embedding_dims // 2 - fraction = ( - 2 - * torch.arange(0, half_embedding_dim, device=inputs.device) - / self.embedding_dims - ) - timescale = ( - self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction - ).to(inputs.device) - if position is None: - seq_length = inputs.shape[1] - position = torch.arange( - seq_length, dtype=torch.float32, device=inputs.device - )[None, :] + def forward( + self, + inputs: torch.Tensor, + position: torch.Tensor | None = None, + ): + """Generates a JTensor of sinusoids with different frequencies.""" + if self.embedding_dims != inputs.shape[-1]: + raise ValueError( + "The embedding dims of the rotary position embedding" + "must match the hidden dimension of the inputs." + ) + half_embedding_dim = self.embedding_dims // 2 + fraction = ( + 2 + * torch.arange(0, half_embedding_dim, device=inputs.device) + / self.embedding_dims + ) + timescale = ( + self.min_timescale + * (self.max_timescale / self.min_timescale) ** fraction + ).to(inputs.device) + if position is None: + seq_length = inputs.shape[1] + position = torch.arange( + seq_length, dtype=torch.float32, device=inputs.device + )[None, :] - if len(inputs.shape) == 4: - position = position[..., None, None] - timescale = timescale[None, None, None, :] - elif len(inputs.shape) == 3: - position = position[..., None] - timescale = timescale[None, None, :] - else: - raise ValueError("Inputs must be of rank 3 or 4.") + if len(inputs.shape) == 4: + position = position[..., None, None] + timescale = timescale[None, None, None, :] + elif len(inputs.shape) == 3: + position = position[..., None] + timescale = timescale[None, None, :] + else: + raise ValueError("Inputs must be of rank 3 or 4.") - sinusoid_inp = position / timescale - sin = torch.sin(sinusoid_inp) - cos = torch.cos(sinusoid_inp) - first_half, second_half = torch.chunk(inputs, 2, dim=-1) - first_part = first_half * cos - second_half * sin - second_part = second_half * cos + first_half * sin - return torch.cat([first_part, second_part], dim=-1) + sinusoid_inp = position / timescale + sin = torch.sin(sinusoid_inp) + cos = torch.cos(sinusoid_inp) + first_half, second_half = torch.chunk(inputs, 2, dim=-1) + first_part = first_half * cos - second_half * sin + second_part = second_half * cos + first_half * sin + return torch.cat([first_part, second_part], dim=-1) def _dot_product_attention( @@ -118,219 +119,224 @@ def _dot_product_attention( value, mask=None, ): - """Computes dot-product attention given query, key, and value.""" - attn_weights = torch.einsum("...qhd,...khd->...hqk", query, key) - if mask is not None: - attn_weights = torch.where( - mask, attn_weights, -torch.finfo(attn_weights.dtype).max / 2 - ) + """Computes dot-product attention given query, key, and value.""" + attn_weights = torch.einsum("...qhd,...khd->...hqk", query, key) + if mask is not None: + attn_weights = torch.where( + mask, attn_weights, -torch.finfo(attn_weights.dtype).max / 2 + ) - attn_weights = F.softmax(attn_weights, dim=-1) + attn_weights = F.softmax(attn_weights, dim=-1) - return torch.einsum("...hqk,...khd->...qhd", attn_weights, value) + return torch.einsum("...hqk,...khd->...qhd", attn_weights, value) class PerDimScale(nn.Module): - """Per-dimension scaling.""" + """Per-dimension scaling.""" - def __init__(self, num_dims: int): - super().__init__() - self.num_dims = num_dims - self.per_dim_scale = nn.Parameter(torch.zeros(num_dims)) + def __init__(self, num_dims: int): + super().__init__() + self.num_dims = num_dims + self.per_dim_scale = nn.Parameter(torch.zeros(num_dims)) - def forward(self, x: torch.Tensor) -> torch.Tensor: - scale_factor = ( - 1.442695041 / math.sqrt(self.num_dims) * F.softplus(self.per_dim_scale) - ) - return x * scale_factor + def forward(self, x: torch.Tensor) -> torch.Tensor: + scale_factor = ( + 1.442695041 / math.sqrt(self.num_dims) * F.softplus(self.per_dim_scale) + ) + return x * scale_factor class MultiHeadAttention(nn.Module): - """Multi-head attention.""" + """Multi-head attention.""" - def __init__( - self, - num_heads: int, - in_features: int, - *, - use_per_dim_scale: bool = True, - use_rotary_position_embeddings: bool = True, - use_bias: bool = False, - attention_fn: Callable[..., torch.Tensor] = _dot_product_attention, - qk_norm: str = "rms", - ): - super().__init__() - self.num_heads = num_heads - self.in_features = in_features - self.head_dim = in_features // num_heads - self.use_bias = use_bias - self.attention_fn = attention_fn - self.qk_norm = qk_norm + def __init__( + self, + num_heads: int, + in_features: int, + *, + use_per_dim_scale: bool = True, + use_rotary_position_embeddings: bool = True, + use_bias: bool = False, + attention_fn: Callable[..., torch.Tensor] = _dot_product_attention, + qk_norm: str = "rms", + ): + super().__init__() + self.num_heads = num_heads + self.in_features = in_features + self.head_dim = in_features // num_heads + self.use_bias = use_bias + self.attention_fn = attention_fn + self.qk_norm = qk_norm - if self.in_features % self.num_heads != 0: - raise ValueError( - f"Memory dimension ({self.in_features}) must be divisible by " - f"'num_heads' heads ({self.num_heads})." - ) + if self.in_features % self.num_heads != 0: + raise ValueError( + f"Memory dimension ({self.in_features}) must be divisible by " + f"'num_heads' heads ({self.num_heads})." + ) - self.query = nn.Linear(self.in_features, self.in_features, bias=use_bias) - self.key = nn.Linear(self.in_features, self.in_features, bias=use_bias) - self.value = nn.Linear(self.in_features, self.in_features, bias=use_bias) - self.out = nn.Linear(self.in_features, self.in_features, bias=use_bias) + self.query = nn.Linear(self.in_features, self.in_features, bias=use_bias) + self.key = nn.Linear(self.in_features, self.in_features, bias=use_bias) + self.value = nn.Linear(self.in_features, self.in_features, bias=use_bias) + self.out = nn.Linear(self.in_features, self.in_features, bias=use_bias) - if self.qk_norm == "rms": - self.query_ln = RMSNorm(self.head_dim) - self.key_ln = RMSNorm(self.head_dim) - else: - self.query_ln = nn.Identity() - self.key_ln = nn.Identity() + if self.qk_norm == "rms": + self.query_ln = RMSNorm(self.head_dim) + self.key_ln = RMSNorm(self.head_dim) + else: + self.query_ln = nn.Identity() + self.key_ln = nn.Identity() - self.use_rotary_position_embeddings = use_rotary_position_embeddings - if self.use_rotary_position_embeddings: - self.rotary_position_embedding = RotaryPositionalEmbedding( - embedding_dims=self.head_dim, - ) + self.use_rotary_position_embeddings = use_rotary_position_embeddings + if self.use_rotary_position_embeddings: + self.rotary_position_embedding = RotaryPositionalEmbedding( + embedding_dims=self.head_dim, + ) - self.use_per_dim_scale = use_per_dim_scale - if use_per_dim_scale: - self.per_dim_scale = PerDimScale(num_dims=self.head_dim) + self.use_per_dim_scale = use_per_dim_scale + if use_per_dim_scale: + self.per_dim_scale = PerDimScale(num_dims=self.head_dim) - def forward( - self, - inputs_q: torch.Tensor, - *, - decode_cache: DecodeCache | None = None, - patch_mask: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, DecodeCache | None]: - b, n_patches, _ = inputs_q.shape - if patch_mask is None: - patch_mask = torch.zeros( - b, n_patches, dtype=torch.bool, device=inputs_q.device - ) + def forward( + self, + inputs_q: torch.Tensor, + *, + decode_cache: DecodeCache | None = None, + patch_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, DecodeCache | None]: + b, n_patches, _ = inputs_q.shape + if patch_mask is None: + patch_mask = torch.zeros( + b, n_patches, dtype=torch.bool, device=inputs_q.device + ) - query = self.query(inputs_q).view(b, n_patches, self.num_heads, self.head_dim) - key = self.key(inputs_q).view(b, n_patches, self.num_heads, self.head_dim) - value = self.value(inputs_q).view(b, n_patches, self.num_heads, self.head_dim) + query = self.query(inputs_q).view( + b, n_patches, self.num_heads, self.head_dim + ) + key = self.key(inputs_q).view(b, n_patches, self.num_heads, self.head_dim) + value = self.value(inputs_q).view( + b, n_patches, self.num_heads, self.head_dim + ) - if decode_cache is None: - num_masked = torch.sum(patch_mask.to(torch.int32), dim=-1) - next_index = torch.zeros_like(num_masked, dtype=torch.int32) - else: - num_masked = ( - torch.sum(patch_mask.to(torch.int32), dim=-1) + decode_cache.num_masked - ) - next_index = decode_cache.next_index.clone() + if decode_cache is None: + num_masked = torch.sum(patch_mask.to(torch.int32), dim=-1) + next_index = torch.zeros_like(num_masked, dtype=torch.int32) + else: + num_masked = ( + torch.sum(patch_mask.to(torch.int32), dim=-1) + + decode_cache.num_masked + ) + next_index = decode_cache.next_index.clone() - if self.use_rotary_position_embeddings: - position = ( - torch.arange(n_patches, device=inputs_q.device)[None, :] - + next_index[:, None] - - num_masked[:, None] - ) - query = self.rotary_position_embedding(query, position) - key = self.rotary_position_embedding(key, position) + if self.use_rotary_position_embeddings: + position = ( + torch.arange(n_patches, device=inputs_q.device)[None, :] + + next_index[:, None] + - num_masked[:, None] + ) + query = self.rotary_position_embedding(query, position) + key = self.rotary_position_embedding(key, position) - query = self.query_ln(query) - key = self.key_ln(key) + query = self.query_ln(query) + key = self.key_ln(key) - if self.use_per_dim_scale: - query = self.per_dim_scale(query) + if self.use_per_dim_scale: + query = self.per_dim_scale(query) - if decode_cache is not None: - _, decode_cache_size, _, _ = decode_cache.value.shape - for i in range(b): - start = decode_cache.next_index[i] - end = start + n_patches - decode_cache.key[i, start:end] = key[i].clone() - decode_cache.value[i, start:end] = value[i].clone() - key = decode_cache.key.clone() - value = decode_cache.value.clone() - decode_cache.next_index += n_patches - decode_cache.num_masked = num_masked - attn_mask = make_attn_mask( - query_length=n_patches, - num_all_masked_kv=num_masked, - query_index_offset=next_index, - kv_length=decode_cache_size, - ) - else: - attn_mask = make_attn_mask( - query_length=n_patches, num_all_masked_kv=num_masked - ) + if decode_cache is not None: + _, decode_cache_size, _, _ = decode_cache.value.shape + for i in range(b): + start = decode_cache.next_index[i] + end = start + n_patches + decode_cache.key[i, start:end] = key[i].clone() + decode_cache.value[i, start:end] = value[i].clone() + key = decode_cache.key.clone() + value = decode_cache.value.clone() + decode_cache.next_index += n_patches + decode_cache.num_masked = num_masked + attn_mask = make_attn_mask( + query_length=n_patches, + num_all_masked_kv=num_masked, + query_index_offset=next_index, + kv_length=decode_cache_size, + ) + else: + attn_mask = make_attn_mask( + query_length=n_patches, num_all_masked_kv=num_masked + ) - x = self.attention_fn( - query, - key, - value, - mask=attn_mask, - ) - x = x.reshape(b, n_patches, self.in_features) - out = self.out(x) - return out, decode_cache + x = self.attention_fn( + query, + key, + value, + mask=attn_mask, + ) + x = x.reshape(b, n_patches, self.in_features) + out = self.out(x) + return out, decode_cache class Transformer(nn.Module): - """Classic Transformer used in TimesFM.""" + """Classic Transformer used in TimesFM.""" - def __init__(self, config: configs.TransformerConfig): - super().__init__() - self.config = config + def __init__(self, config: configs.TransformerConfig): + super().__init__() + self.config = config - if config.attention_norm == "rms": - self.pre_attn_ln = RMSNorm(num_features=config.model_dims) - self.post_attn_ln = RMSNorm(num_features=config.model_dims) - else: - raise ValueError(f"Layer norm: {config.attention_norm} not supported.") + if config.attention_norm == "rms": + self.pre_attn_ln = RMSNorm(num_features=config.model_dims) + self.post_attn_ln = RMSNorm(num_features=config.model_dims) + else: + raise ValueError(f"Layer norm: {config.attention_norm} not supported.") - self.attn = MultiHeadAttention( - num_heads=config.num_heads, - in_features=config.model_dims, - use_per_dim_scale=True, - use_rotary_position_embeddings=config.use_rotary_position_embeddings, - qk_norm=config.qk_norm, + self.attn = MultiHeadAttention( + num_heads=config.num_heads, + in_features=config.model_dims, + use_per_dim_scale=True, + use_rotary_position_embeddings=config.use_rotary_position_embeddings, + qk_norm=config.qk_norm, + ) + + if config.feedforward_norm == "rms": + self.pre_ff_ln = RMSNorm(num_features=config.model_dims) + self.post_ff_ln = RMSNorm(num_features=config.model_dims) + else: + raise ValueError(f"Layer norm: {config.feedforward_norm} not supported.") + + self.ff0 = nn.Linear( + in_features=config.model_dims, + out_features=config.hidden_dims, + bias=config.use_bias, + ) + self.ff1 = nn.Linear( + in_features=config.hidden_dims, + out_features=config.model_dims, + bias=config.use_bias, + ) + if config.ff_activation == "relu": + self.activation = nn.ReLU() + elif config.ff_activation == "swish": + self.activation = nn.SiLU() + elif config.ff_activation == "none": + self.activation = nn.Identity() + else: + raise ValueError(f"Activation: {config.ff_activation} not supported.") + + def forward( + self, + input_embeddings: torch.Tensor, + patch_mask: torch.Tensor, + decode_cache: DecodeCache | None = None, + ) -> tuple[torch.Tensor, DecodeCache | None]: + attn_output, decode_cache = self.attn( + inputs_q=self.pre_attn_ln(input_embeddings), + decode_cache=decode_cache, + patch_mask=patch_mask, + ) + attn_output = self.post_attn_ln(attn_output) + input_embeddings + output_embeddings = ( + self.post_ff_ln( + self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output)))) ) - - if config.feedforward_norm == "rms": - self.pre_ff_ln = RMSNorm(num_features=config.model_dims) - self.post_ff_ln = RMSNorm(num_features=config.model_dims) - else: - raise ValueError(f"Layer norm: {config.feedforward_norm} not supported.") - - self.ff0 = nn.Linear( - in_features=config.model_dims, - out_features=config.hidden_dims, - bias=config.use_bias, - ) - self.ff1 = nn.Linear( - in_features=config.hidden_dims, - out_features=config.model_dims, - bias=config.use_bias, - ) - if config.ff_activation == "relu": - self.activation = nn.ReLU() - elif config.ff_activation == "swish": - self.activation = nn.SiLU() - elif config.ff_activation == "none": - self.activation = nn.Identity() - else: - raise ValueError(f"Activation: {config.ff_activation} not supported.") - - def forward( - self, - input_embeddings: torch.Tensor, - patch_mask: torch.Tensor, - decode_cache: DecodeCache | None = None, - ) -> tuple[torch.Tensor, DecodeCache | None]: - attn_output, decode_cache = self.attn( - inputs_q=self.pre_attn_ln(input_embeddings), - decode_cache=decode_cache, - patch_mask=patch_mask, - ) - attn_output = self.post_attn_ln(attn_output) + input_embeddings - output_embeddings = ( - self.post_ff_ln( - self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output)))) - ) - + attn_output - ) - return output_embeddings, decode_cache + + attn_output + ) + return output_embeddings, decode_cache diff --git a/src/timesfm/torch/util.py b/src/timesfm/torch/util.py index 0505bda..81efe29 100644 --- a/src/timesfm/torch/util.py +++ b/src/timesfm/torch/util.py @@ -22,12 +22,12 @@ _TOLERANCE = 1e-6 @dataclasses.dataclass(frozen=False) class DecodeCache: - """Cache for decoding.""" + """Cache for decoding.""" - next_index: torch.Tensor - num_masked: torch.Tensor - key: torch.Tensor - value: torch.Tensor + next_index: torch.Tensor + num_masked: torch.Tensor + key: torch.Tensor + value: torch.Tensor def update_running_stats( @@ -40,36 +40,38 @@ def update_running_stats( tuple[torch.Tensor, torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor, torch.Tensor], ]: - """Updates the running stats.""" - is_legit = torch.logical_not(mask) - inc_n = torch.sum(is_legit.to(x.dtype), dim=-1) + """Updates the running stats.""" + is_legit = torch.logical_not(mask) + inc_n = torch.sum(is_legit.to(x.dtype), dim=-1) - inc_mu_numerator = torch.sum(x * is_legit, dim=-1) - inc_n_safe = torch.where(inc_n == 0, 1.0, inc_n) - inc_mu = inc_mu_numerator / inc_n_safe - inc_mu = torch.where(inc_n == 0, 0.0, inc_mu) + inc_mu_numerator = torch.sum(x * is_legit, dim=-1) + inc_n_safe = torch.where(inc_n == 0, 1.0, inc_n) + inc_mu = inc_mu_numerator / inc_n_safe + inc_mu = torch.where(inc_n == 0, 0.0, inc_mu) - inc_var_numerator = torch.sum(((x - inc_mu.unsqueeze(-1)) ** 2) * is_legit, dim=-1) - inc_var = inc_var_numerator / inc_n_safe - inc_var = torch.where(inc_n == 0, 0.0, inc_var) - inc_sigma = torch.sqrt(inc_var) + inc_var_numerator = torch.sum( + ((x - inc_mu.unsqueeze(-1)) ** 2) * is_legit, dim=-1 + ) + inc_var = inc_var_numerator / inc_n_safe + inc_var = torch.where(inc_n == 0, 0.0, inc_var) + inc_sigma = torch.sqrt(inc_var) - new_n = n + inc_n - new_n_safe = torch.where(new_n == 0, 1.0, new_n) + new_n = n + inc_n + new_n_safe = torch.where(new_n == 0, 1.0, new_n) - new_mu = (n * mu + inc_mu * inc_n) / new_n_safe - new_mu = torch.where(new_n == 0, 0.0, new_mu) + new_mu = (n * mu + inc_mu * inc_n) / new_n_safe + new_mu = torch.where(new_n == 0, 0.0, new_mu) - term1 = n * sigma.pow(2) - term2 = inc_n * inc_sigma.pow(2) - term3 = n * (mu - new_mu).pow(2) - term4 = inc_n * (inc_mu - new_mu).pow(2) + term1 = n * sigma.pow(2) + term2 = inc_n * inc_sigma.pow(2) + term3 = n * (mu - new_mu).pow(2) + term4 = inc_n * (inc_mu - new_mu).pow(2) - new_var = (term1 + term2 + term3 + term4) / new_n_safe - new_var = torch.where(new_n == 0, 0.0, new_var) - new_sigma = torch.sqrt(torch.clamp(new_var, min=0.0)) + new_var = (term1 + term2 + term3 + term4) / new_n_safe + new_var = torch.where(new_n == 0, 0.0, new_var) + new_sigma = torch.sqrt(torch.clamp(new_var, min=0.0)) - return (w := (new_n, new_mu, new_sigma), w) + return (w := (new_n, new_mu, new_sigma), w) def revin( @@ -78,15 +80,15 @@ def revin( sigma: torch.Tensor, reverse: bool = False, ): - """Reversible instance normalization.""" - if len(mu.shape) == len(x.shape) - 1: - mu = mu[..., None] - sigma = sigma[..., None] - elif len(mu.shape) == len(x.shape) - 2: - mu = mu[..., None, None] - sigma = sigma[..., None, None] + """Reversible instance normalization.""" + if len(mu.shape) == len(x.shape) - 1: + mu = mu[..., None] + sigma = sigma[..., None] + elif len(mu.shape) == len(x.shape) - 2: + mu = mu[..., None, None] + sigma = sigma[..., None, None] - if reverse: - return x * sigma + mu - else: - return (x - mu) / torch.where(sigma < _TOLERANCE, 1.0, sigma) + if reverse: + return x * sigma + mu + else: + return (x - mu) / torch.where(sigma < _TOLERANCE, 1.0, sigma)