freedak/timesfm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #299 from google-research/siriuz42-2.0-pr

Browse Source 
MVP readme
This commit is contained in:
Yichen Zhou
2025-09-15 09:57:09 -07:00
committed by GitHub
parent 59d2985681 099cefe8b9
commit b810bbdf9f
9 changed files with 1033 additions and 881 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+82 -1
View File
@@ -1,3 +1,84 @@
# TimesFM # 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.
```
src/timesfm/__init__.py
+1
View File
@@ -14,6 +14,7 @@
"""TimesFM API.""" """TimesFM API."""
from .configs import ForecastConfig
from .timesfm_2p5 import timesfm_2p5_torch from .timesfm_2p5 import timesfm_2p5_torch
TimesFM_2p5_200M_torch = timesfm_2p5_torch.TimesFM_2p5_200M_torch TimesFM_2p5_200M_torch = timesfm_2p5_torch.TimesFM_2p5_200M_torch
src/timesfm/configs.py
+61 -35
View File
@@ -20,59 +20,85 @@ from typing import Literal
@dataclasses.dataclass(frozen=False) @dataclasses.dataclass(frozen=False)
class ForecastConfig: class ForecastConfig:
"""Options for forecasting.""" """Options for forecasting.
max_context: int = 0 Attributes:
max_horizon: int = 0 max_context: The maximum context length. This is used by the complied decode
normalize_inputs: bool = False function at inference time during batched inference. Any input time series
window_size: int = 0 with length less than max_context will be padded with zeros, and with
per_core_batch_size: int = 1 length greater than max_context will be truncated.
use_continuous_quantile_head: bool = False max_horizon: The maximum horizon length. This is used by the complied decode
force_flip_invariance: bool = True function at inference time during batched inference. The compiled cached
infer_is_positive: bool = True decoding function will by default forecast till max_horizon.
fix_quantile_crossing: bool = False normalize_inputs: Whether to normalize the inputs. This is useful when the
return_backcast: bool = False 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) @dataclasses.dataclass(frozen=True)
class ResidualBlockConfig: class ResidualBlockConfig:
"""Framework-agnostic config for a residual block.""" """Framework-agnostic config for a residual block."""
input_dims: int input_dims: int
hidden_dims: int hidden_dims: int
output_dims: int output_dims: int
use_bias: bool use_bias: bool
activation: Literal["relu", "swish", "none"] activation: Literal["relu", "swish", "none"]
@dataclasses.dataclass(frozen=True) @dataclasses.dataclass(frozen=True)
class RandomFourierFeaturesConfig: class RandomFourierFeaturesConfig:
"""Framework-agnostic config for random fourier features.""" """Framework-agnostic config for random fourier features."""
input_dims: int input_dims: int
output_dims: int output_dims: int
projection_stddev: float projection_stddev: float
use_bias: bool use_bias: bool
@dataclasses.dataclass(frozen=True) @dataclasses.dataclass(frozen=True)
class TransformerConfig: class TransformerConfig:
"""Framework-agnostic config for a transformer.""" """Framework-agnostic config for a transformer."""
model_dims: int model_dims: int
hidden_dims: int hidden_dims: int
num_heads: int num_heads: int
attention_norm: Literal["rms"] attention_norm: Literal["rms"]
feedforward_norm: Literal["rms"] feedforward_norm: Literal["rms"]
qk_norm: Literal["rms", "none"] qk_norm: Literal["rms", "none"]
use_bias: bool use_bias: bool
use_rotary_position_embeddings: bool use_rotary_position_embeddings: bool
ff_activation: Literal["relu", "swish", "none"] ff_activation: Literal["relu", "swish", "none"]
@dataclasses.dataclass(frozen=True) @dataclasses.dataclass(frozen=True)
class StackedTransformersConfig: class StackedTransformersConfig:
"""Framework-agnostic config for a stacked transformers.""" """Framework-agnostic config for a stacked transformers."""
num_layers: int num_layers: int
transformer: TransformerConfig transformer: TransformerConfig
src/timesfm/timesfm_2p5/timesfm_2p5_base.py
+136 -130
View File
@@ -26,161 +26,167 @@ ForecastConfig = configs.ForecastConfig
def strip_leading_nans(arr): 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: Args:
arr: The input NumPy array. arr: The input NumPy array.
Returns: Returns:
A new NumPy array with leading NaN values removed. A new NumPy array with leading NaN values removed.
If the array is all NaNs or empty, returns an empty array. If the array is all NaNs or empty, returns an empty array.
""" """
isnan = np.isnan(arr) isnan = np.isnan(arr)
first_valid_index = np.argmax(~isnan) first_valid_index = np.argmax(~isnan)
return arr[first_valid_index:] return arr[first_valid_index:]
def linear_interpolation(arr): 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: Args:
arr: The 1D numpy array containing NaN values. arr: The 1D numpy array containing NaN values.
Returns: Returns:
A new numpy array with NaN values filled using linear interpolation, A new numpy array with NaN values filled using linear interpolation,
or the original array if no NaNs are present. or the original array if no NaNs are present.
Returns None if the input is not a 1D array. Returns None if the input is not a 1D array.
Returns the original array if there are no NaN values. Returns the original array if there are no NaN values.
""" """
nans = np.isnan(arr) nans = np.isnan(arr)
if not np.any(nans): # Check if there are any NaNs 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 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) @dataclasses.dataclass(frozen=True)
class TimesFM_2p5_200M_Definition: class TimesFM_2p5_200M_Definition:
"""Framework-agnostic config of TimesFM 2.5.""" """Framework-agnostic config of TimesFM 2.5."""
context_limit = 16384 context_limit = 16384
input_patch_len: int = 32 input_patch_len: int = 32
output_patch_len: int = 128 output_patch_len: int = 128
output_quantile_len: int = 1024 output_quantile_len: int = 1024
quantiles: list[float] = dataclasses.field( 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] 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 decode_index: int = 5
tokenizer: ResidualBlockConfig = ResidualBlockConfig( tokenizer: ResidualBlockConfig = ResidualBlockConfig(
input_dims=64, input_dims=64,
hidden_dims=1280, hidden_dims=1280,
output_dims=1280, output_dims=1280,
use_bias=True, use_bias=True,
activation="swish", activation="swish",
) )
stacked_transformers: StackedTransformersConfig = StackedTransformersConfig( stacked_transformers: StackedTransformersConfig = StackedTransformersConfig(
num_layers=20, num_layers=20,
transformer=TransformerConfig( transformer=TransformerConfig(
model_dims=1280, model_dims=1280,
hidden_dims=1280, hidden_dims=1280,
num_heads=16, num_heads=16,
attention_norm="rms", attention_norm="rms",
feedforward_norm="rms", feedforward_norm="rms",
qk_norm="rms", qk_norm="rms",
use_bias=False, use_bias=False,
use_rotary_position_embeddings=True, use_rotary_position_embeddings=True,
ff_activation="swish", ff_activation="swish",
), ),
) )
output_projection_point: ResidualBlockConfig = ResidualBlockConfig( output_projection_point: ResidualBlockConfig = ResidualBlockConfig(
input_dims=1280, input_dims=1280,
hidden_dims=1280, hidden_dims=1280,
output_dims=1280, output_dims=1280,
use_bias=False, use_bias=False,
activation="swish", activation="swish",
) )
output_projection_quantiles: ResidualBlockConfig = ResidualBlockConfig( output_projection_quantiles: ResidualBlockConfig = ResidualBlockConfig(
input_dims=1280, input_dims=1280,
hidden_dims=1280, hidden_dims=1280,
output_dims=10240, output_dims=10240,
use_bias=False, use_bias=False,
activation="swish", activation="swish",
) )
class TimesFM_2p5: class TimesFM_2p5:
"""Abstract base class for TimesFM models.""" """Abstract base class for TimesFM models.
forecast_config: ForecastConfig | None = None Attributes:
compiled_decode: Callable[..., Any] | None = None forecast_config: Configuration for forecasting flags.
global_batch_size: int = 0 compiled_decode: Compiled decode function.
global_batch_size: Global batch size.
"""
def load_checkpoint(self, path: str): forecast_config: ForecastConfig | None = None
"""Loads a TimesFM model from a checkpoint.""" compiled_decode: Callable[..., Any] | None = None
raise NotImplementedError() global_batch_size: int = 0
def compile(self, forecast_config: ForecastConfig | None = None): def load_checkpoint(self, path: str):
"""Compiles the TimesFM model for fast decoding.""" """Loads a TimesFM model from a checkpoint."""
raise NotImplementedError() raise NotImplementedError()
def forecast( def compile(self, forecast_config: ForecastConfig | None = None):
self, horizon: int, inputs: list[np.ndarray] """Compiles the TimesFM model for fast decoding."""
) -> tuple[np.ndarray, np.ndarray]: raise NotImplementedError()
"""Forecasts the time series."""
if self.compiled_decode is None:
raise RuntimeError("Model is not compiled. Please call compile() first.")
assert self.global_batch_size > 0 def forecast(
assert self.forecast_config is not None 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 assert self.global_batch_size > 0
num_inputs = len(inputs) assert self.forecast_config is not None
if (w := num_inputs % self.global_batch_size) != 0:
inputs += [np.array([0.0] * 3)] * (self.global_batch_size - w)
output_points = [] context = self.forecast_config.max_context
output_quantiles = [] 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 = [] values = []
masks = [] 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_points = np.concatenate(output_points, axis=0)
output_quantiles = np.concatenate(output_quantiles, axis=0) output_quantiles = np.concatenate(output_quantiles, axis=0)
return output_points[:num_inputs], output_quantiles[:num_inputs] return output_points[:num_inputs], output_quantiles[:num_inputs]
src/timesfm/timesfm_2p5/timesfm_2p5_torch.py
+377 -347
View File
@@ -35,384 +35,414 @@ revin = util.revin
class TimesFM_2p5_200M_torch_module(nn.Module): 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): def __init__(self):
super().__init__() super().__init__()
# Names constants. # Names constants.
self.p = self.config.input_patch_len # 32 self.p = self.config.input_patch_len # 32
self.o = self.config.output_patch_len # 128 self.o = self.config.output_patch_len # 128
self.os = self.config.output_quantile_len # 1024 self.os = self.config.output_quantile_len # 1024
self.m = self.o // self.p # 4 self.m = self.o // self.p # 4
self.x = self.config.stacked_transformers.num_layers # 20 self.x = self.config.stacked_transformers.num_layers # 20
self.h = self.config.stacked_transformers.transformer.num_heads # 16 self.h = self.config.stacked_transformers.transformer.num_heads # 16
self.md = self.config.stacked_transformers.transformer.model_dims # 1280 self.md = self.config.stacked_transformers.transformer.model_dims # 1280
self.hd = self.md // self.h # 80 self.hd = self.md // self.h # 80
self.q = len(self.config.quantiles) + 1 # 10 self.q = len(self.config.quantiles) + 1 # 10
self.aridx = self.config.decode_index # 5 self.aridx = self.config.decode_index # 5
# Layers. # Layers.
self.tokenizer = dense.ResidualBlock(self.config.tokenizer) self.tokenizer = dense.ResidualBlock(self.config.tokenizer)
self.stacked_xf = nn.ModuleList( self.stacked_xf = nn.ModuleList([
[ transformer.Transformer(self.config.stacked_transformers.transformer)
transformer.Transformer(self.config.stacked_transformers.transformer) for _ in range(self.x)
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( patch_mu.append(mu)
self.config.output_projection_point 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( new_mask = torch.zeros_like(new_patched_input, dtype=torch.bool)
self.config.output_projection_quantiles
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. new_renormed_output = torch.reshape(
if torch.cuda.is_available(): revin(new_normed_output, new_mu, new_sigma, reverse=True),
self.device = torch.device("cuda:0") (batch_size, self.m, self.o, self.q),
self.device_count = torch.cuda.device_count() )
else: ar_outputs.append(new_renormed_output[:, -1, ...])
self.device = torch.device("cpu") last_renormed_output = new_renormed_output[:, -1, :, self.aridx]
self.device_count = 1
def load_checkpoint(self, path: str): if num_decode_steps > 0:
"""Loads a PyTorch TimesFM model from a checkpoint.""" ar_renormed_outputs = torch.stack(ar_outputs, dim=1)
tensors = load_file(path) else:
self.load_state_dict(tensors) ar_renormed_outputs = None
self.to(self.device)
def forward( return renormed_outputs, renormed_quantile_spread, ar_renormed_outputs
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: def forecast_naive(
decode_caches = [None] * self.x self, horizon: int, inputs: Sequence[np.ndarray]
) -> list[np.ndarray]:
"""Forecasts the time series.
output_embeddings = input_embeddings This is a naive implementation for debugging purposes. No forecasting
new_decode_caches = [] flags are used here. Forecasting quality can be subpar.
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 ( Args:
input_embeddings, horizon: The number of time points to forecast.
output_embeddings, inputs: A sequence of numpy arrays, each representing a time series to
output_ts, query forecast for.
output_quantile_spread,
), new_decode_caches
def decode(self, horizon: int, inputs, masks): Returns:
"""Decodes the time series.""" A list of numpy arrays of forecasts.
"""
inputs = inputs.to(self.device) outputs = []
masks = masks.to(self.device) for each_input in inputs:
input_t = torch.tensor(each_input, dtype=torch.float32)
with torch.no_grad(): mask = torch.zeros_like(input_t, dtype=torch.bool)
batch_size, context = inputs.shape[0], inputs.shape[1] len_front_mask = self.p - (len(each_input) % self.p)
num_decode_steps = (horizon - 1) // self.o if len_front_mask < self.p:
num_input_patches = context // self.p input_t = torch.cat(
decode_cache_size = num_input_patches + num_decode_steps * self.m [torch.zeros(len_front_mask, dtype=torch.float32), input_t], dim=0
)
# Prefill mask = torch.cat(
patched_inputs = torch.reshape(inputs, (batch_size, -1, self.p)) [torch.ones(len_front_mask, dtype=torch.bool), mask], dim=0
patched_masks = torch.reshape(masks, (batch_size, -1, self.p)) )
input_t = input_t[None, ...]
# running stats mask = mask[None, ...]
n = torch.zeros(batch_size, device=inputs.device) t_pf, _, t_ar = self.decode(horizon, input_t, mask)
mu = torch.zeros(batch_size, device=inputs.device) to_concat = [t_pf[:, -1, ...]]
sigma = torch.zeros(batch_size, device=inputs.device) if t_ar is not None:
patch_mu = [] to_concat.append(t_ar.reshape(1, -1, self.q))
patch_sigma = [] torch_forecast = torch.cat(to_concat, dim=1)[..., :horizon]
for i in range(num_input_patches): torch_forecast = torch_forecast.squeeze(0)
(n, mu, sigma), _ = util.update_running_stats( outputs.append(torch_forecast.detach().cpu().numpy())
n, mu, sigma, patched_inputs[:, i], patched_masks[:, i] return outputs
)
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
class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5): 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( def load_checkpoint(
self, self,
*, *,
path: str | None = None, path: str | None = None,
hf_repo_id: str | None = "google/timesfm-2.5-200m-pytorch", hf_repo_id: str | None = "google/timesfm-2.5-200m-pytorch",
): ) -> None:
"""Loads a PyTorch safetensors TimesFM model.""" """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 compile(self, forecast_config: configs.ForecastConfig, **kwargs): Args:
if kwargs.get("backend", None) is not None: path: Path to a local checkpoint. If not provided, will try to download
self.model.compile(**kwargs) from the default Hugging Face repo.
self.global_batch_size = ( hf_repo_id: If provided, will download from the specified Hugging Face
forecast_config.per_core_batch_size * self.model.device_count 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. inputs = torch.Tensor(np.array(inputs)).to(self.model.device)
fc = forecast_config 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: if fc.infer_is_positive:
logging.info( is_positive = torch.all(inputs >= 0, dim=-1, keepdim=True)
"When compiling, max context needs to be multiple of the patch size" else:
" %d. Using max context = %d instead.", is_positive = None
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 fc.normalize_inputs:
if horizon > fc.max_horizon: mu = torch.mean(inputs, dim=-1, keepdim=True)
raise ValueError( sigma = torch.std(inputs, dim=-1, keepdim=True)
"Horizon must be less than the max horizon." inputs = revin(inputs, mu, sigma, reverse=False)
f" {horizon} > {fc.max_horizon}." else:
) mu, sigma = None, None
inputs = torch.Tensor(np.array(inputs)).to(self.model.device) pf_outputs, quantile_spreads, ar_outputs = self.model.decode(
masks = torch.Tensor(np.array(masks)).to(self.model.device).to(torch.bool) forecast_config.max_horizon, inputs, masks
batch_size = inputs.shape[0] )
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: flip_quantile_fn = lambda x: torch.cat(
is_positive = torch.all(inputs >= 0, dim=-1, keepdim=True) [x[..., :1], torch.flip(x[..., 1:], dims=(-1,))], dim=-1
else: )
is_positive = None
if fc.normalize_inputs: if fc.force_flip_invariance:
mu = torch.mean(inputs, dim=-1, keepdim=True) flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = (
sigma = torch.std(inputs, dim=-1, keepdim=True) self.model.decode(forecast_config.max_horizon, -inputs, masks)
inputs = revin(inputs, mu, sigma, reverse=False) )
else: flipped_quantile_spreads = flip_quantile_fn(flipped_quantile_spreads)
mu, sigma = None, None 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( if fc.use_continuous_quantile_head:
forecast_config.max_horizon, inputs, masks for quantile_index in [1, 2, 3, 4, 6, 7, 8, 9]:
) full_forecast[:, :, quantile_index] = (
full_forecast = torch.cat( quantile_spreads[:, : fc.max_horizon, quantile_index]
[ - quantile_spreads[:, : fc.max_horizon, 5]
pf_outputs[:, -1, ...], + full_forecast[:, : fc.max_horizon, 5]
ar_outputs.reshape(batch_size, -1, self.model.q), )
], full_forecast = full_forecast[:, :horizon, :]
dim=1,
)
flip_quantile_fn = lambda x: torch.cat( if fc.return_backcast:
[x[..., :1], torch.flip(x[..., 1:], dims=(-1,))], dim=-1 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: if fc.fix_quantile_crossing:
flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = ( for i in [4, 3, 2, 1]:
self.model.decode(forecast_config.max_horizon, -inputs, masks) full_forecast[:, :, i] = torch.where(
) full_forecast[:, :, i] < full_forecast[:, :, i + 1],
flipped_quantile_spreads = flip_quantile_fn(flipped_quantile_spreads) full_forecast[:, :, i],
flipped_pf_outputs = flip_quantile_fn(flipped_pf_outputs) full_forecast[:, :, i + 1],
flipped_full_forecast = torch.cat( )
[ for i in [6, 7, 8, 9]:
flipped_pf_outputs[:, -1, ...], full_forecast[:, :, i] = torch.where(
flipped_ar_outputs.reshape(batch_size, -1, self.model.q), full_forecast[:, :, i] > full_forecast[:, :, i - 1],
], full_forecast[:, :, i],
dim=1, full_forecast[:, :, i - 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.use_continuous_quantile_head: if fc.normalize_inputs:
for quantile_index in [1, 2, 3, 4, 6, 7, 8, 9]: full_forecast = revin(full_forecast, mu, sigma, reverse=True)
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.return_backcast: if is_positive is not None:
full_backcast = pf_outputs[:, :-1, : self.model.p, :].reshape( full_forecast = torch.where(
batch_size, -1, self.model.q is_positive[..., None],
) torch.maximum(full_forecast, torch.zeros_like(full_forecast)),
full_forecast = torch.cat([full_backcast, full_forecast], dim=1) full_forecast,
)
if fc.fix_quantile_crossing: full_forecast = full_forecast.detach().cpu().numpy()
for i in [4, 3, 2, 1]: return full_forecast[..., 5], full_forecast
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.normalize_inputs: self.compiled_decode = _compiled_decode
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
src/timesfm/torch/dense.py
+62 -62
View File
@@ -21,74 +21,74 @@ from .. import configs
class ResidualBlock(nn.Module): 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): def __init__(self, config: configs.ResidualBlockConfig):
super().__init__() super().__init__()
self.config = config self.config = config
self.hidden_layer = nn.Linear( self.hidden_layer = nn.Linear(
in_features=config.input_dims, in_features=config.input_dims,
out_features=config.hidden_dims, out_features=config.hidden_dims,
bias=config.use_bias, bias=config.use_bias,
) )
self.output_layer = nn.Linear( self.output_layer = nn.Linear(
in_features=config.hidden_dims, in_features=config.hidden_dims,
out_features=config.output_dims, out_features=config.output_dims,
bias=config.use_bias, bias=config.use_bias,
) )
self.residual_layer = nn.Linear( self.residual_layer = nn.Linear(
in_features=config.input_dims, in_features=config.input_dims,
out_features=config.output_dims, out_features=config.output_dims,
bias=config.use_bias, bias=config.use_bias,
) )
if config.activation == "relu": if config.activation == "relu":
self.activation = nn.ReLU() self.activation = nn.ReLU()
elif config.activation == "swish": elif config.activation == "swish":
self.activation = nn.SiLU() self.activation = nn.SiLU()
elif config.activation == "none": elif config.activation == "none":
self.activation = nn.Identity() self.activation = nn.Identity()
else: else:
raise ValueError(f"Activation: {config.activation} not supported.") raise ValueError(f"Activation: {config.activation} not supported.")
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.output_layer( return self.output_layer(
self.activation(self.hidden_layer(x)) self.activation(self.hidden_layer(x))
) + self.residual_layer(x) ) + self.residual_layer(x)
class RandomFourierFeatures(nn.Module): class RandomFourierFeatures(nn.Module):
"""Random Fourier features layer.""" """Random Fourier features layer."""
def __init__(self, config: configs.RandomFourierFeaturesConfig): def __init__(self, config: configs.RandomFourierFeaturesConfig):
super().__init__() super().__init__()
self.config = config self.config = config
if config.output_dims % 4 != 0: if config.output_dims % 4 != 0:
raise ValueError( raise ValueError(
f"Output dims must be a multiple of 4: {config.output_dims} % 4 != 0." f"Output dims must be a multiple of 4: {config.output_dims} % 4 != 0."
) )
num_projected_features = config.output_dims // 4 num_projected_features = config.output_dims // 4
self.phase_shifts = nn.Parameter(torch.zeros(2, num_projected_features)) self.phase_shifts = nn.Parameter(torch.zeros(2, num_projected_features))
self.projection_layer = nn.Linear( self.projection_layer = nn.Linear(
in_features=config.input_dims, in_features=config.input_dims,
out_features=num_projected_features, out_features=num_projected_features,
bias=config.use_bias, bias=config.use_bias,
) )
self.residual_layer = nn.Linear( self.residual_layer = nn.Linear(
in_features=config.input_dims, in_features=config.input_dims,
out_features=config.output_dims, out_features=config.output_dims,
bias=config.use_bias, bias=config.use_bias,
) )
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
projected = self.projection_layer(x) projected = self.projection_layer(x)
cos_features = torch.cos(projected) cos_features = torch.cos(projected)
sin_features = torch.sin(projected) sin_features = torch.sin(projected)
sq_wave_1 = torch.sign(torch.sin(projected + self.phase_shifts[0, :])) sq_wave_1 = torch.sign(torch.sin(projected + self.phase_shifts[0, :]))
sq_wave_2 = torch.sign(torch.sin(projected + self.phase_shifts[1, :])) sq_wave_2 = torch.sign(torch.sin(projected + self.phase_shifts[1, :]))
fourier_features = torch.cat( fourier_features = torch.cat(
[cos_features, sin_features, sq_wave_1, sq_wave_2], dim=-1 [cos_features, sin_features, sq_wave_1, sq_wave_2], dim=-1
) )
residual = self.residual_layer(x) residual = self.residual_layer(x)
return fourier_features + residual return fourier_features + residual
src/timesfm/torch/normalization.py
+16 -16
View File
@@ -19,21 +19,21 @@ from torch import nn
class RMSNorm(nn.Module): class RMSNorm(nn.Module):
"""RMS normalization.""" """RMS normalization."""
def __init__( def __init__(
self, self,
num_features: int, num_features: int,
*, *,
epsilon: float = 1e-6, epsilon: float = 1e-6,
): ):
super().__init__() super().__init__()
self.scale = nn.Parameter(torch.zeros(num_features)) self.scale = nn.Parameter(torch.zeros(num_features))
self.num_features = num_features self.num_features = num_features
self.epsilon = epsilon self.epsilon = epsilon
def forward(self, inputs: torch.Tensor) -> torch.Tensor: def forward(self, inputs: torch.Tensor) -> torch.Tensor:
var = torch.mean(torch.square(inputs), dim=-1, keepdim=True) var = torch.mean(torch.square(inputs), dim=-1, keepdim=True)
normed_inputs = inputs * torch.rsqrt(var + self.epsilon) normed_inputs = inputs * torch.rsqrt(var + self.epsilon)
normed_inputs = normed_inputs * self.scale normed_inputs = normed_inputs * self.scale
return normed_inputs return normed_inputs
src/timesfm/torch/transformer.py
+257 -251
View File
@@ -36,80 +36,81 @@ def make_attn_mask(
query_index_offset: torch.Tensor | None = None, query_index_offset: torch.Tensor | None = None,
kv_length: int = 0, kv_length: int = 0,
) -> torch.Tensor: ) -> torch.Tensor:
"""Makes attention mask.""" """Makes attention mask."""
if kv_length == 0: if kv_length == 0:
kv_length = query_length kv_length = query_length
q_index = torch.arange(query_length, device=num_all_masked_kv.device)[ q_index = torch.arange(query_length, device=num_all_masked_kv.device)[
None, None, :, None None, None, :, None
] ]
if query_index_offset is not None: if query_index_offset is not None:
q_index = q_index + query_index_offset[:, None, None, None] q_index = q_index + query_index_offset[:, None, None, None]
kv_index = torch.arange(kv_length, device=num_all_masked_kv.device)[ kv_index = torch.arange(kv_length, device=num_all_masked_kv.device)[
None, None, None, : None, None, None, :
] ]
return torch.logical_and( return torch.logical_and(
q_index >= kv_index, q_index >= kv_index,
kv_index >= num_all_masked_kv[:, None, None, None], kv_index >= num_all_masked_kv[:, None, None, None],
) )
class RotaryPositionalEmbedding(nn.Module): class RotaryPositionalEmbedding(nn.Module):
"""Rotary positional embedding.""" """Rotary positional embedding."""
def __init__( def __init__(
self, self,
embedding_dims: int, embedding_dims: int,
min_timescale: float = 1.0, min_timescale: float = 1.0,
max_timescale: float = 10000.0, max_timescale: float = 10000.0,
): ):
super().__init__() super().__init__()
self.embedding_dims = embedding_dims self.embedding_dims = embedding_dims
self.min_timescale = min_timescale self.min_timescale = min_timescale
self.max_timescale = max_timescale self.max_timescale = max_timescale
def forward( def forward(
self, self,
inputs: torch.Tensor, inputs: torch.Tensor,
position: torch.Tensor | None = None, position: torch.Tensor | None = None,
): ):
"""Generates a JTensor of sinusoids with different frequencies.""" """Generates a JTensor of sinusoids with different frequencies."""
if self.embedding_dims != inputs.shape[-1]: if self.embedding_dims != inputs.shape[-1]:
raise ValueError( raise ValueError(
"The embedding dims of the rotary position embedding" "The embedding dims of the rotary position embedding"
"must match the hidden dimension of the inputs." "must match the hidden dimension of the inputs."
) )
half_embedding_dim = self.embedding_dims // 2 half_embedding_dim = self.embedding_dims // 2
fraction = ( fraction = (
2 2
* torch.arange(0, half_embedding_dim, device=inputs.device) * torch.arange(0, half_embedding_dim, device=inputs.device)
/ self.embedding_dims / self.embedding_dims
) )
timescale = ( timescale = (
self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction self.min_timescale
).to(inputs.device) * (self.max_timescale / self.min_timescale) ** fraction
if position is None: ).to(inputs.device)
seq_length = inputs.shape[1] if position is None:
position = torch.arange( seq_length = inputs.shape[1]
seq_length, dtype=torch.float32, device=inputs.device position = torch.arange(
)[None, :] seq_length, dtype=torch.float32, device=inputs.device
)[None, :]
if len(inputs.shape) == 4: if len(inputs.shape) == 4:
position = position[..., None, None] position = position[..., None, None]
timescale = timescale[None, None, None, :] timescale = timescale[None, None, None, :]
elif len(inputs.shape) == 3: elif len(inputs.shape) == 3:
position = position[..., None] position = position[..., None]
timescale = timescale[None, None, :] timescale = timescale[None, None, :]
else: else:
raise ValueError("Inputs must be of rank 3 or 4.") raise ValueError("Inputs must be of rank 3 or 4.")
sinusoid_inp = position / timescale sinusoid_inp = position / timescale
sin = torch.sin(sinusoid_inp) sin = torch.sin(sinusoid_inp)
cos = torch.cos(sinusoid_inp) cos = torch.cos(sinusoid_inp)
first_half, second_half = torch.chunk(inputs, 2, dim=-1) first_half, second_half = torch.chunk(inputs, 2, dim=-1)
first_part = first_half * cos - second_half * sin first_part = first_half * cos - second_half * sin
second_part = second_half * cos + first_half * sin second_part = second_half * cos + first_half * sin
return torch.cat([first_part, second_part], dim=-1) return torch.cat([first_part, second_part], dim=-1)
def _dot_product_attention( def _dot_product_attention(
@@ -118,219 +119,224 @@ def _dot_product_attention(
value, value,
mask=None, mask=None,
): ):
"""Computes dot-product attention given query, key, and value.""" """Computes dot-product attention given query, key, and value."""
attn_weights = torch.einsum("...qhd,...khd->...hqk", query, key) attn_weights = torch.einsum("...qhd,...khd->...hqk", query, key)
if mask is not None: if mask is not None:
attn_weights = torch.where( attn_weights = torch.where(
mask, attn_weights, -torch.finfo(attn_weights.dtype).max / 2 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): class PerDimScale(nn.Module):
"""Per-dimension scaling.""" """Per-dimension scaling."""
def __init__(self, num_dims: int): def __init__(self, num_dims: int):
super().__init__() super().__init__()
self.num_dims = num_dims self.num_dims = num_dims
self.per_dim_scale = nn.Parameter(torch.zeros(num_dims)) self.per_dim_scale = nn.Parameter(torch.zeros(num_dims))
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
scale_factor = ( scale_factor = (
1.442695041 / math.sqrt(self.num_dims) * F.softplus(self.per_dim_scale) 1.442695041 / math.sqrt(self.num_dims) * F.softplus(self.per_dim_scale)
) )
return x * scale_factor return x * scale_factor
class MultiHeadAttention(nn.Module): class MultiHeadAttention(nn.Module):
"""Multi-head attention.""" """Multi-head attention."""
def __init__( def __init__(
self, self,
num_heads: int, num_heads: int,
in_features: int, in_features: int,
*, *,
use_per_dim_scale: bool = True, use_per_dim_scale: bool = True,
use_rotary_position_embeddings: bool = True, use_rotary_position_embeddings: bool = True,
use_bias: bool = False, use_bias: bool = False,
attention_fn: Callable[..., torch.Tensor] = _dot_product_attention, attention_fn: Callable[..., torch.Tensor] = _dot_product_attention,
qk_norm: str = "rms", qk_norm: str = "rms",
): ):
super().__init__() super().__init__()
self.num_heads = num_heads self.num_heads = num_heads
self.in_features = in_features self.in_features = in_features
self.head_dim = in_features // num_heads self.head_dim = in_features // num_heads
self.use_bias = use_bias self.use_bias = use_bias
self.attention_fn = attention_fn self.attention_fn = attention_fn
self.qk_norm = qk_norm self.qk_norm = qk_norm
if self.in_features % self.num_heads != 0: if self.in_features % self.num_heads != 0:
raise ValueError( raise ValueError(
f"Memory dimension ({self.in_features}) must be divisible by " f"Memory dimension ({self.in_features}) must be divisible by "
f"'num_heads' heads ({self.num_heads})." f"'num_heads' heads ({self.num_heads})."
) )
self.query = 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.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.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.out = nn.Linear(self.in_features, self.in_features, bias=use_bias)
if self.qk_norm == "rms": if self.qk_norm == "rms":
self.query_ln = RMSNorm(self.head_dim) self.query_ln = RMSNorm(self.head_dim)
self.key_ln = RMSNorm(self.head_dim) self.key_ln = RMSNorm(self.head_dim)
else: else:
self.query_ln = nn.Identity() self.query_ln = nn.Identity()
self.key_ln = nn.Identity() self.key_ln = nn.Identity()
self.use_rotary_position_embeddings = use_rotary_position_embeddings self.use_rotary_position_embeddings = use_rotary_position_embeddings
if self.use_rotary_position_embeddings: if self.use_rotary_position_embeddings:
self.rotary_position_embedding = RotaryPositionalEmbedding( self.rotary_position_embedding = RotaryPositionalEmbedding(
embedding_dims=self.head_dim, embedding_dims=self.head_dim,
) )
self.use_per_dim_scale = use_per_dim_scale self.use_per_dim_scale = use_per_dim_scale
if use_per_dim_scale: if use_per_dim_scale:
self.per_dim_scale = PerDimScale(num_dims=self.head_dim) self.per_dim_scale = PerDimScale(num_dims=self.head_dim)
def forward( def forward(
self, self,
inputs_q: torch.Tensor, inputs_q: torch.Tensor,
*, *,
decode_cache: DecodeCache | None = None, decode_cache: DecodeCache | None = None,
patch_mask: torch.Tensor | None = None, patch_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, DecodeCache | None]: ) -> tuple[torch.Tensor, DecodeCache | None]:
b, n_patches, _ = inputs_q.shape b, n_patches, _ = inputs_q.shape
if patch_mask is None: if patch_mask is None:
patch_mask = torch.zeros( patch_mask = torch.zeros(
b, n_patches, dtype=torch.bool, device=inputs_q.device 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) query = self.query(inputs_q).view(
key = self.key(inputs_q).view(b, n_patches, self.num_heads, self.head_dim) b, n_patches, self.num_heads, self.head_dim
value = self.value(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: if decode_cache is None:
num_masked = torch.sum(patch_mask.to(torch.int32), dim=-1) num_masked = torch.sum(patch_mask.to(torch.int32), dim=-1)
next_index = torch.zeros_like(num_masked, dtype=torch.int32) next_index = torch.zeros_like(num_masked, dtype=torch.int32)
else: else:
num_masked = ( num_masked = (
torch.sum(patch_mask.to(torch.int32), dim=-1) + decode_cache.num_masked torch.sum(patch_mask.to(torch.int32), dim=-1)
) + decode_cache.num_masked
next_index = decode_cache.next_index.clone() )
next_index = decode_cache.next_index.clone()
if self.use_rotary_position_embeddings: if self.use_rotary_position_embeddings:
position = ( position = (
torch.arange(n_patches, device=inputs_q.device)[None, :] torch.arange(n_patches, device=inputs_q.device)[None, :]
+ next_index[:, None] + next_index[:, None]
- num_masked[:, None] - num_masked[:, None]
) )
query = self.rotary_position_embedding(query, position) query = self.rotary_position_embedding(query, position)
key = self.rotary_position_embedding(key, position) key = self.rotary_position_embedding(key, position)
query = self.query_ln(query) query = self.query_ln(query)
key = self.key_ln(key) key = self.key_ln(key)
if self.use_per_dim_scale: if self.use_per_dim_scale:
query = self.per_dim_scale(query) query = self.per_dim_scale(query)
if decode_cache is not None: if decode_cache is not None:
_, decode_cache_size, _, _ = decode_cache.value.shape _, decode_cache_size, _, _ = decode_cache.value.shape
for i in range(b): for i in range(b):
start = decode_cache.next_index[i] start = decode_cache.next_index[i]
end = start + n_patches end = start + n_patches
decode_cache.key[i, start:end] = key[i].clone() decode_cache.key[i, start:end] = key[i].clone()
decode_cache.value[i, start:end] = value[i].clone() decode_cache.value[i, start:end] = value[i].clone()
key = decode_cache.key.clone() key = decode_cache.key.clone()
value = decode_cache.value.clone() value = decode_cache.value.clone()
decode_cache.next_index += n_patches decode_cache.next_index += n_patches
decode_cache.num_masked = num_masked decode_cache.num_masked = num_masked
attn_mask = make_attn_mask( attn_mask = make_attn_mask(
query_length=n_patches, query_length=n_patches,
num_all_masked_kv=num_masked, num_all_masked_kv=num_masked,
query_index_offset=next_index, query_index_offset=next_index,
kv_length=decode_cache_size, kv_length=decode_cache_size,
) )
else: else:
attn_mask = make_attn_mask( attn_mask = make_attn_mask(
query_length=n_patches, num_all_masked_kv=num_masked query_length=n_patches, num_all_masked_kv=num_masked
) )
x = self.attention_fn( x = self.attention_fn(
query, query,
key, key,
value, value,
mask=attn_mask, mask=attn_mask,
) )
x = x.reshape(b, n_patches, self.in_features) x = x.reshape(b, n_patches, self.in_features)
out = self.out(x) out = self.out(x)
return out, decode_cache return out, decode_cache
class Transformer(nn.Module): class Transformer(nn.Module):
"""Classic Transformer used in TimesFM.""" """Classic Transformer used in TimesFM."""
def __init__(self, config: configs.TransformerConfig): def __init__(self, config: configs.TransformerConfig):
super().__init__() super().__init__()
self.config = config self.config = config
if config.attention_norm == "rms": if config.attention_norm == "rms":
self.pre_attn_ln = RMSNorm(num_features=config.model_dims) self.pre_attn_ln = RMSNorm(num_features=config.model_dims)
self.post_attn_ln = RMSNorm(num_features=config.model_dims) self.post_attn_ln = RMSNorm(num_features=config.model_dims)
else: else:
raise ValueError(f"Layer norm: {config.attention_norm} not supported.") raise ValueError(f"Layer norm: {config.attention_norm} not supported.")
self.attn = MultiHeadAttention( self.attn = MultiHeadAttention(
num_heads=config.num_heads, num_heads=config.num_heads,
in_features=config.model_dims, in_features=config.model_dims,
use_per_dim_scale=True, use_per_dim_scale=True,
use_rotary_position_embeddings=config.use_rotary_position_embeddings, use_rotary_position_embeddings=config.use_rotary_position_embeddings,
qk_norm=config.qk_norm, 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))))
) )
+ attn_output
if config.feedforward_norm == "rms": )
self.pre_ff_ln = RMSNorm(num_features=config.model_dims) return output_embeddings, decode_cache
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
src/timesfm/torch/util.py
+41 -39
View File
@@ -22,12 +22,12 @@ _TOLERANCE = 1e-6
@dataclasses.dataclass(frozen=False) @dataclasses.dataclass(frozen=False)
class DecodeCache: class DecodeCache:
"""Cache for decoding.""" """Cache for decoding."""
next_index: torch.Tensor next_index: torch.Tensor
num_masked: torch.Tensor num_masked: torch.Tensor
key: torch.Tensor key: torch.Tensor
value: torch.Tensor value: torch.Tensor
def update_running_stats( 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],
tuple[torch.Tensor, torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor, torch.Tensor],
]: ]:
"""Updates the running stats.""" """Updates the running stats."""
is_legit = torch.logical_not(mask) is_legit = torch.logical_not(mask)
inc_n = torch.sum(is_legit.to(x.dtype), dim=-1) inc_n = torch.sum(is_legit.to(x.dtype), dim=-1)
inc_mu_numerator = torch.sum(x * is_legit, 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_n_safe = torch.where(inc_n == 0, 1.0, inc_n)
inc_mu = inc_mu_numerator / inc_n_safe inc_mu = inc_mu_numerator / inc_n_safe
inc_mu = torch.where(inc_n == 0, 0.0, inc_mu) 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_numerator = torch.sum(
inc_var = inc_var_numerator / inc_n_safe ((x - inc_mu.unsqueeze(-1)) ** 2) * is_legit, dim=-1
inc_var = torch.where(inc_n == 0, 0.0, inc_var) )
inc_sigma = torch.sqrt(inc_var) 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 = n + inc_n
new_n_safe = torch.where(new_n == 0, 1.0, new_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 = (n * mu + inc_mu * inc_n) / new_n_safe
new_mu = torch.where(new_n == 0, 0.0, new_mu) new_mu = torch.where(new_n == 0, 0.0, new_mu)
term1 = n * sigma.pow(2) term1 = n * sigma.pow(2)
term2 = inc_n * inc_sigma.pow(2) term2 = inc_n * inc_sigma.pow(2)
term3 = n * (mu - new_mu).pow(2) term3 = n * (mu - new_mu).pow(2)
term4 = inc_n * (inc_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 = (term1 + term2 + term3 + term4) / new_n_safe
new_var = torch.where(new_n == 0, 0.0, new_var) new_var = torch.where(new_n == 0, 0.0, new_var)
new_sigma = torch.sqrt(torch.clamp(new_var, min=0.0)) 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( def revin(
@@ -78,15 +80,15 @@ def revin(
sigma: torch.Tensor, sigma: torch.Tensor,
reverse: bool = False, reverse: bool = False,
): ):
"""Reversible instance normalization.""" """Reversible instance normalization."""
if len(mu.shape) == len(x.shape) - 1: if len(mu.shape) == len(x.shape) - 1:
mu = mu[..., None] mu = mu[..., None]
sigma = sigma[..., None] sigma = sigma[..., None]
elif len(mu.shape) == len(x.shape) - 2: elif len(mu.shape) == len(x.shape) - 2:
mu = mu[..., None, None] mu = mu[..., None, None]
sigma = sigma[..., None, None] sigma = sigma[..., None, None]
if reverse: if reverse:
return x * sigma + mu return x * sigma + mu
else: else:
return (x - mu) / torch.where(sigma < _TOLERANCE, 1.0, sigma) return (x - mu) / torch.where(sigma < _TOLERANCE, 1.0, sigma)