To comply with Google OSS style guide.
This commit is contained in:
Yichen Zhou
+35
-35
@@ -20,59 +20,59 @@ 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
|
||||
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
|
||||
|
||||
@@ -26,161 +26,161 @@ 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
|
||||
forecast_config: ForecastConfig | None = None
|
||||
compiled_decode: Callable[..., Any] | None = None
|
||||
global_batch_size: int = 0
|
||||
|
||||
def load_checkpoint(self, path: str):
|
||||
"""Loads a TimesFM model from a checkpoint."""
|
||||
raise NotImplementedError()
|
||||
def load_checkpoint(self, path: str):
|
||||
"""Loads a TimesFM model from a checkpoint."""
|
||||
raise NotImplementedError()
|
||||
|
||||
def compile(self, forecast_config: ForecastConfig | None = None):
|
||||
"""Compiles the TimesFM model for fast decoding."""
|
||||
raise NotImplementedError()
|
||||
def compile(self, forecast_config: ForecastConfig | None = None):
|
||||
"""Compiles the TimesFM model for fast decoding."""
|
||||
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 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.")
|
||||
|
||||
assert self.global_batch_size > 0
|
||||
assert self.forecast_config is not None
|
||||
assert self.global_batch_size > 0
|
||||
assert self.forecast_config is not None
|
||||
|
||||
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)
|
||||
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 = []
|
||||
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]
|
||||
|
||||
@@ -35,384 +35,394 @@ 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
|
||||
|
||||
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]
|
||||
)
|
||||
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
|
||||
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):
|
||||
"""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",
|
||||
):
|
||||
"""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: Use another Hugging Face repo ID.
|
||||
"""
|
||||
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.
|
||||
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
|
||||
)
|
||||
full_forecast = torch.cat(
|
||||
[
|
||||
pf_outputs[:, -1, ...],
|
||||
ar_outputs.reshape(batch_size, -1, self.model.q),
|
||||
],
|
||||
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)
|
||||
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(
|
||||
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
|
||||
|
||||
+62
-62
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+257
-251
@@ -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
|
||||
|
||||
+41
-39
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user