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

To comply with Google OSS style guide.

Browse Source
This commit is contained in:
Yichen Zhou
2025-09-13 05:30:07 +00:00
parent 59d2985681
commit b9e31fd02a
7 changed files with 901 additions and 883 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/timesfm/configs.py
+35 -35
View File
@@ -20,59 +20,59 @@ 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 max_context: int = 0
max_horizon: int = 0 max_horizon: int = 0
normalize_inputs: bool = False normalize_inputs: bool = False
window_size: int = 0 window_size: int = 0
per_core_batch_size: int = 1 per_core_batch_size: int = 1
use_continuous_quantile_head: bool = False use_continuous_quantile_head: bool = False
force_flip_invariance: bool = True force_flip_invariance: bool = True
infer_is_positive: bool = True infer_is_positive: bool = True
fix_quantile_crossing: bool = False fix_quantile_crossing: bool = False
return_backcast: 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
+130 -130
View File
@@ -26,161 +26,161 @@ 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 forecast_config: ForecastConfig | None = None
compiled_decode: Callable[..., Any] | None = None compiled_decode: Callable[..., Any] | None = None
global_batch_size: int = 0 global_batch_size: int = 0
def load_checkpoint(self, path: str): def load_checkpoint(self, path: str):
"""Loads a TimesFM model from a checkpoint.""" """Loads a TimesFM model from a checkpoint."""
raise NotImplementedError() raise NotImplementedError()
def compile(self, forecast_config: ForecastConfig | None = None): def compile(self, forecast_config: ForecastConfig | None = None):
"""Compiles the TimesFM model for fast decoding.""" """Compiles the TimesFM model for fast decoding."""
raise NotImplementedError() raise NotImplementedError()
def forecast( def forecast(
self, horizon: int, inputs: list[np.ndarray] self, horizon: int, inputs: list[np.ndarray]
) -> tuple[np.ndarray, np.ndarray]: ) -> tuple[np.ndarray, np.ndarray]:
"""Forecasts the time series.""" """Forecasts the time series."""
if self.compiled_decode is None: if self.compiled_decode is None:
raise RuntimeError("Model is not compiled. Please call compile() first.") raise RuntimeError("Model is not compiled. Please call compile() first.")
assert self.global_batch_size > 0 assert self.global_batch_size > 0
assert self.forecast_config is not None assert self.forecast_config is not None
context = self.forecast_config.max_context context = self.forecast_config.max_context
num_inputs = len(inputs) num_inputs = len(inputs)
if (w := num_inputs % self.global_batch_size) != 0: if (w := num_inputs % self.global_batch_size) != 0:
inputs += [np.array([0.0] * 3)] * (self.global_batch_size - w) inputs += [np.array([0.0] * 3)] * (self.global_batch_size - w)
output_points = [] output_points = []
output_quantiles = [] 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
+360 -350
View File
@@ -35,384 +35,394 @@ 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(self, horizon: int, inputs: Sequence[np.ndarray]):
decode_caches = [None] * self.x """Forecasts the time series."""
outputs = []
output_embeddings = input_embeddings for each_input in inputs:
new_decode_caches = [] input_t = torch.tensor(each_input, dtype=torch.float32)
for i, layer in enumerate(self.stacked_xf): mask = torch.zeros_like(input_t, dtype=torch.bool)
output_embeddings, new_cache = layer( len_front_mask = self.p - (len(each_input) % self.p)
output_embeddings, masks[..., -1], decode_caches[i] if len_front_mask < self.p:
) input_t = torch.cat(
new_decode_caches.append(new_cache) [torch.zeros(len_front_mask, dtype=torch.float32), input_t], dim=0
output_ts = self.output_projection_point(output_embeddings) )
output_quantile_spread = self.output_projection_quantiles(output_embeddings) mask = torch.cat(
[torch.ones(len_front_mask, dtype=torch.bool), mask], dim=0
return ( )
input_embeddings, input_t = input_t[None, ...]
output_embeddings, mask = mask[None, ...]
output_ts, t_pf, _, t_ar = self.decode(horizon, input_t, mask)
output_quantile_spread, to_concat = [t_pf[:, -1, ...]]
), new_decode_caches if t_ar is not None:
to_concat.append(t_ar.reshape(1, -1, self.q))
def decode(self, horizon: int, inputs, masks): torch_forecast = torch.cat(to_concat, dim=1)[..., :horizon]
"""Decodes the time series.""" torch_forecast = torch_forecast.squeeze(0)
outputs.append(torch_forecast.detach().cpu().numpy())
inputs = inputs.to(self.device) return outputs
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
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",
): ):
"""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: Use another Hugging Face repo ID.
forecast_config.per_core_batch_size * self.model.device_count """
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):
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] )
full_forecast = torch.cat(
[
pf_outputs[:, -1, ...],
ar_outputs.reshape(batch_size, -1, self.model.q),
],
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)
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
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)