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
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."""
from .configs import ForecastConfig
from .timesfm_2p5 import timesfm_2p5_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)
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
src/timesfm/timesfm_2p5/timesfm_2p5_base.py
+136 -130
View File
@@ -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]
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):
"""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
src/timesfm/torch/dense.py
+62 -62
View File
@@ -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
src/timesfm/torch/normalization.py
+16 -16
View File
@@ -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
src/timesfm/torch/transformer.py
+257 -251
View File
@@ -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
src/timesfm/torch/util.py
+41 -39
View File
@@ -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)