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

Full pytorch support

Browse Source
This commit is contained in:
Rajat Sen
2024-09-12 23:30:46 +00:00
parent 61fa1b2ef2
commit 1b95563eea
19 changed files with 2402 additions and 1935 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/__init__.py
View File
src/timesfm/__init__.py
+7 -2
View File
@@ -11,7 +11,12 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""TimesFM init file."""
from .timesfm import TimesFm, freq_map
from timesfm.timesfm_base import freq_map, TimesFmCheckpoint, TimesFmHparams, TimesFmBase
try:
from timesfm.timesfm_jax import TimesFmJax as TimesFm
from timesfm import data_loader
except Exception as _:
print("No pax dependencies installed.")
from timesfm.timesfm_torch import TimesFmTorch as TimesFm
src/timesfm/patched_decoder.py
+45 -61
View File
@@ -11,7 +11,6 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Pax ML model for patched time-series decoder.
The file implements Residual MLPs, Patched Decoder layers and PAX ML models.
@@ -36,7 +35,6 @@ from praxis.layers import normalizations
from praxis.layers import stochastics
from praxis.layers import transformers
# PAX shortcuts
NestedMap = py_utils.NestedMap
JTensor = pytypes.JTensor
@@ -44,7 +42,6 @@ JTensor = pytypes.JTensor
LayerTpl = pax_fiddle.Config[base_layer.BaseLayer]
template_field = base_layer.template_field
PAD_VAL = 1123581321.0
DEFAULT_QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
@@ -57,7 +54,6 @@ _FREQ = "freq"
_OUTPUT_TOKENS = "output_tokens"
_STATS = "stats"
# Small numerical value.
_TOLERANCE = 1e-7
@@ -158,9 +154,8 @@ class ResidualBlock(base_layer.BaseLayer):
return output + residual
def _masked_mean_std(
inputs: JTensor, padding: JTensor
) -> Tuple[JTensor, JTensor]:
def _masked_mean_std(inputs: JTensor,
padding: JTensor) -> Tuple[JTensor, JTensor]:
"""Calculates mean and standard deviation of arr across axis 1.
It should exclude values where pad is 1.
@@ -197,7 +192,7 @@ def _masked_mean_std(
# Calculate the masked sum and squared sum of M
masked_sum = jnp.sum(arr * mask, axis=1)
masked_squared_sum = jnp.sum((arr * mask) ** 2, axis=1)
masked_squared_sum = jnp.sum((arr * mask)**2, axis=1)
# Calculate the masked mean and standard deviation
masked_mean = masked_sum / num_valid_elements
@@ -240,8 +235,7 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
quantiles: list[float] = dataclasses.field(default_factory=_create_quantiles)
residual_block_tpl: LayerTpl = template_field(ResidualBlock)
stacked_transformer_params_tpl: LayerTpl = template_field(
transformers.StackedTransformer
)
transformers.StackedTransformer)
use_freq: bool = True
def setup(self) -> None:
@@ -276,9 +270,8 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
self.create_child(
"position_emb",
pax_fiddle.Config(
layers.PositionalEmbedding, embedding_dims=self.model_dims
),
pax_fiddle.Config(layers.PositionalEmbedding,
embedding_dims=self.model_dims),
)
if self.use_freq:
@@ -292,27 +285,24 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
)
def transform_decode_state(
self, transform_fn: base_layer.DecodeStateTransformFn
) -> None:
self, transform_fn: base_layer.DecodeStateTransformFn) -> None:
"""Transforms all decode state variables based on transform_fn."""
self.stacked_transformer_layer.transform_decode_state(transform_fn)
def _forward_transform(
self, inputs: JTensor, patched_pads: JTensor
) -> Tuple[JTensor, Tuple[JTensor, JTensor]]:
self, inputs: JTensor,
patched_pads: JTensor) -> Tuple[JTensor, Tuple[JTensor, JTensor]]:
"""Input is of shape [B, N, P]."""
mu, sigma = _masked_mean_std(inputs, patched_pads)
sigma = jnp.where(sigma < _TOLERANCE, 1.0, sigma)
# Normalize each patch.
outputs = (inputs - mu[:, None, None]) / sigma[:, None, None]
outputs = jnp.where(
jnp.abs(inputs - PAD_VAL) < _TOLERANCE, PAD_VAL, outputs
)
jnp.abs(inputs - PAD_VAL) < _TOLERANCE, PAD_VAL, outputs)
return outputs, (mu, sigma)
def _reverse_transform(
self, outputs: JTensor, stats: Tuple[JTensor, JTensor]
) -> JTensor:
def _reverse_transform(self, outputs: JTensor,
stats: Tuple[JTensor, JTensor]) -> JTensor:
"""Output is of shape [B, N, P, Q]."""
mu, sigma = stats
return outputs * sigma[:, None, None, None] + mu[:, None, None, None]
@@ -326,18 +316,15 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
"""Preprocess input for stacked transformer."""
# Reshape into patches.
patched_inputs = es.jax_einshape("b(np)->bnp", input_ts, p=self.patch_len)
patched_pads = es.jax_einshape(
"b(np)->bnp", input_padding, p=self.patch_len
)
patched_pads = es.jax_einshape("b(np)->bnp",
input_padding,
p=self.patch_len)
patched_inputs = jnp.where(
jnp.abs(patched_pads - 1.0) < _TOLERANCE, 0.0, patched_inputs
)
jnp.abs(patched_pads - 1.0) < _TOLERANCE, 0.0, patched_inputs)
patched_pads = jnp.where(
jnp.abs(patched_inputs - PAD_VAL) < _TOLERANCE, 1, patched_pads
)
patched_inputs, stats = self._forward_transform(
patched_inputs, patched_pads
)
jnp.abs(patched_inputs - PAD_VAL) < _TOLERANCE, 1, patched_pads)
patched_inputs, stats = self._forward_transform(patched_inputs,
patched_pads)
# B x N x D
patched_inputs = patched_inputs * (1.0 - patched_pads)
@@ -367,9 +354,10 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
"""Postprocess output of stacked transformer."""
# B x N x (H.Q)
output_ts = self.horizon_ff_layer(model_output)
output_ts = es.jax_einshape(
"bn(hq)->bnhq", output_ts, q=num_outputs, h=self.horizon_len
)
output_ts = es.jax_einshape("bn(hq)->bnhq",
output_ts,
q=num_outputs,
h=self.horizon_len)
return self._reverse_transform(output_ts, stats)
def __call__(self, inputs: NestedMap) -> NestedMap:
@@ -400,9 +388,11 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
model_output = self.stacked_transformer_layer(model_input, patched_padding)
output_ts = self._postprocess_output(model_output, num_outputs, stats)
return NestedMap(
{_OUTPUT_TOKENS: model_output, _OUTPUT_TS: output_ts, _STATS: stats}
)
return NestedMap({
_OUTPUT_TOKENS: model_output,
_OUTPUT_TS: output_ts,
_STATS: stats
})
def decode(
self,
@@ -443,15 +433,13 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
if paddings.shape[1] != final_out.shape[1] + horizon_len:
raise ValueError(
"Length of paddings must match length of input + horizon_len:"
f" {paddings.shape[1]} != {final_out.shape[1]} + {horizon_len}"
)
f" {paddings.shape[1]} != {final_out.shape[1]} + {horizon_len}")
if output_patch_len is None:
output_patch_len = self.horizon_len
num_decode_patches = (
horizon_len + output_patch_len - 1
) // output_patch_len
num_decode_patches = (horizon_len + output_patch_len -
1) // output_patch_len
for step_index in range(num_decode_patches):
current_padding = paddings[:, 0 : final_out.shape[1]]
current_padding = paddings[:, 0:final_out.shape[1]]
input_ts = final_out[:, -max_len:]
input_padding = current_padding[:, -max_len:]
model_input = NestedMap(
@@ -463,7 +451,7 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
if return_forecast_on_context and step_index == 0:
# For the first decodings step, collect the model forecast on the
# context except the unavailable first input batch forecast.
new_full_ts = fprop_outputs[:, :-1, : self.patch_len, :]
new_full_ts = fprop_outputs[:, :-1, :self.patch_len, :]
new_full_ts = es.jax_einshape("bnph->b(np)h", new_full_ts)
full_outputs.append(new_full_ts)
@@ -477,9 +465,9 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
if return_forecast_on_context:
# `full_outputs` indexing starts at after the first input patch.
full_outputs = jnp.concatenate(full_outputs, axis=1)[
:, : (context_len - self.patch_len + horizon_len), :
]
full_outputs = jnp.concatenate(full_outputs,
axis=1)[:, :(context_len - self.patch_len +
horizon_len), :]
else:
# `full_outputs` indexing starts at the forecast horizon.
full_outputs = jnp.concatenate(full_outputs, axis=1)[:, 0:horizon_len, :]
@@ -506,14 +494,12 @@ class PatchedDecoderFinetuneModel(base_model.BaseModel):
input_padding = jnp.zeros_like(input_ts)
context_len = input_ts.shape[1]
input_patch_len = self.core_layer_tpl.patch_len
context_pad = (
(context_len + input_patch_len - 1) // input_patch_len
) * input_patch_len - context_len
context_pad = ((context_len + input_patch_len - 1) //
input_patch_len) * input_patch_len - context_len
input_ts = jnp.pad(input_ts, [(0, 0), (context_pad, 0)])
input_padding = jnp.pad(
input_padding, [(0, 0), (context_pad, 0)], constant_values=1
)
input_padding = jnp.pad(input_padding, [(0, 0), (context_pad, 0)],
constant_values=1)
freq = jnp.ones([input_ts.shape[0], 1], dtype=jnp.int32) * self.freq
new_input_batch = NestedMap(
input_ts=input_ts,
@@ -522,9 +508,8 @@ class PatchedDecoderFinetuneModel(base_model.BaseModel):
)
return self.core_layer(new_input_batch)
def _quantile_loss(
self, pred: JTensor, actual: JTensor, quantile: float
) -> JTensor:
def _quantile_loss(self, pred: JTensor, actual: JTensor,
quantile: float) -> JTensor:
"""Calculates quantile loss.
Args:
@@ -540,12 +525,11 @@ class PatchedDecoderFinetuneModel(base_model.BaseModel):
loss_second = -dev * (1.0 - quantile)
return 2 * jnp.where(loss_first >= 0, loss_first, loss_second)
def compute_loss(
self, prediction_output: NestedMap, input_batch: NestedMap
) -> Tuple[NestedMap, NestedMap]:
def compute_loss(self, prediction_output: NestedMap,
input_batch: NestedMap) -> Tuple[NestedMap, NestedMap]:
output_ts = prediction_output[_OUTPUT_TS]
actual_ts = input_batch[_TARGET_FUTURE]
pred_ts = output_ts[:, -1, 0 : actual_ts.shape[1], :]
pred_ts = output_ts[:, -1, 0:actual_ts.shape[1], :]
loss = jnp.square(pred_ts[:, :, 0] - actual_ts)
for i, quantile in enumerate(self.core_layer.quantiles):
loss += self._quantile_loss(pred_ts[:, :, i + 1], actual_ts, quantile)
src/timesfm_torch/pytorch_patched_decoder.py → src/timesfm/pytorch_patched_decoder.py
+3 -2
View File
@@ -22,7 +22,7 @@ import torch.nn.functional as F
def _create_quantiles() -> list[float]:
return [0.1, 0.25, 0.5, 0.75, 0.9]
return [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
@dataclasses.dataclass
@@ -192,7 +192,8 @@ def convert_paddings_to_mask(
Returns:
A torch.Tensor of shape [B, 1, 1, T] ready to add to attention logits.
"""
attention_mask = paddings[:, None, None, :] # Equivalent to jnp.newaxis
attention_mask = paddings.detach().clone()
attention_mask = attention_mask[:, None, None, :] # Equivalent to jnp.newaxis
attention_mask *= get_large_negative_number(dtype)
return attention_mask
src/timesfm/timesfm.py → src/timesfm/timesfm_base.py
+103 -360
View File
@@ -11,43 +11,26 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""TimesFM forecast API for inference."""
"""Base class for TimesFM inference. This will be common to PAX and Pytorch."""
import collections
import dataclasses
import logging
import multiprocessing
from os import path
import time
from typing import Any, Literal, Optional, Sequence
from typing import Any, Literal, Sequence
import einshape as es
from huggingface_hub import snapshot_download
import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd
from paxml import checkpoints
from paxml import tasks_lib
from praxis import base_hyperparams
from praxis import base_layer
from praxis import pax_fiddle
from praxis import py_utils
from praxis import pytypes
from praxis.layers import normalizations
from praxis.layers import transformers
from utilsforecast.processing import make_future_dataframe
from . import patched_decoder
from . import xreg_lib
instantiate = base_hyperparams.instantiate
NestedMap = py_utils.NestedMap
JTensor = pytypes.JTensor
Category = xreg_lib.Category
XRegMode = xreg_lib.XRegMode
_TOL = 1e-6
DEFAULT_QUANTILES = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)
def process_group(key, group, value_name, forecast_context_len):
@@ -79,7 +62,7 @@ def freq_map(freq: str):
# Per time series normalization: forward.
def _normalize(batch):
def normalize(batch):
stats = [
(np.mean(x), np.where((w := np.std(x)) > _TOL, w, 1.0)) for x in batch
]
@@ -88,32 +71,19 @@ def _normalize(batch):
# Per time series normalization: inverse.
def _renormalize(batch, stats):
def renormalize(batch, stats):
return [x * stat[1] + stat[0] for x, stat in zip(batch, stats)]
class TimesFm:
"""TimesFM forecast API for inference.
@dataclasses.dataclass(kw_only=True)
class TimesFmHparams:
"""Hparams used to initialize a TimesFM model for inference.
This class is the scaffolding for calling TimesFM forecast. To properly use:
1. Create an instance with the correct hyperparameters of a TimesFM model.
2. Call `load_from_checkpoint` to load a compatible checkpoint.
3. Call `forecast` for inference.
Given the model size, this API does not shard the model weights for SPMD. All
parallelism happens on the data dimension.
Compilation happens during the first time `forecast` is called and uses the
`per_core_batch_size` to set and freeze the input signature. Subsequent calls
to `forecast` reflect the actual inference latency.
These are the sufficient subset of hparams to configure TimesFM inference
agnostic to the checkpoint version, and are not necessarily the same as the
hparams used to train the checkpoint.
Attributes:
per_core_batch_size: Batch size on each core for data parallelism.
backend: One of "cpu", "gpu" or "tpu".
num_devices: Number of cores provided the backend.
global_batch_size: per_core_batch_size * num_devices. Each batch of
inference task will be padded with respect to global_batch_size to
minimize latency.
context_len: Largest context length the model allows for each decode call.
This technically can be any large, but practically should set to the
context length the checkpoint was trained with.
@@ -122,237 +92,96 @@ class TimesFm:
output_patch_len: Output patch len. How many timepoints is taken from a
single step of autoregressive decoding. Can be set as the training horizon
of the checkpoint.
mesh_shape: Shape of the data parallelism mesh.
mesh_name: Names of the data parallelism mesh.
model_p: Configuration of the TimesFM model deduced from the hparams.
num_layers: Number of transformer layers in the model.
model_dims: Model dimension.
per_core_batch_size: Batch size on each core for data parallelism.
backend: One of "cpu", "gpu" or "tpu".
quantiles: Which quantiles are output by the model.
"""
context_len: int = 512
horizon_len: int = 128
input_patch_len: int = 32
output_patch_len: int = 128
num_layers: int = 20
num_heads: int = 16
model_dims: int = 1280
per_core_batch_size: int = 32
backend: Literal["cpu", "gpu", "tpu"] = "cpu"
quantiles: Sequence[float] | None = DEFAULT_QUANTILES
@dataclasses.dataclass(kw_only=True)
class TimesFmCheckpoint:
"""Checkpoint used to initialize a TimesFM model for inference.
Attributes:
version: Version of the checkpoint, e.g. "jax", "torch", "tensorflow", etc.
The factory will create the corresponding TimesFm inference class based on
this version.
path: Path to the checkpoint.
type: If provided, type of the checkpoint used by the specific checkpoint
loader per version.
step: If provided, step of the checkpoint.
"""
version: str = "jax"
path: str | None = None
huggingface_repo_id: str | None = None
type: Any = None
step: int | None = None
class TimesFmBase:
"""Base TimesFM forecast API for inference.
This class is the scaffolding for calling TimesFM forecast. To properly use:
1. Create an instance with the correct hyperparameters of a TimesFM model.
2. Call `load_from_checkpoint` to load a compatible checkpoint.
3. Call `forecast` for inference.
"""
def _logging(self, s):
if self._verbose:
print(s)
print(s)
def __init__(
self,
context_len: int,
horizon_len: int,
input_patch_len: int,
output_patch_len: int,
num_layers: int,
model_dims: int,
per_core_batch_size: int = 32,
backend: Literal["cpu", "gpu", "tpu"] = "cpu",
quantiles: Sequence[float] | None = None,
verbose: bool = True,
) -> None:
def __post_init__(self) -> None:
"""Additional initialization for subclasses before checkpoint loading."""
pass
def __init__(self, hparams: TimesFmHparams,
checkpoint: TimesFmCheckpoint) -> None:
"""Initializes the TimesFM forecast API.
Args:
context_len: Largest context length the model allows for each decode call.
This technically can be any large, but practically should set to the
context length the checkpoint was trained with.
horizon_len: Forecast horizon.
input_patch_len: Input patch len.
output_patch_len: Output patch len. How many timepoints is taken from a
single step of autoregressive decoding. Can be set as the training
horizon of the checkpoint.
num_layers: Number of transformer layers.
model_dims: Model dimension.
per_core_batch_size: Batch size on each core for data parallelism.
backend: One of "cpu", "gpu" or "tpu".
quantiles: list of output quantiles supported by the model.
verbose: Whether to print logging messages.
hparams: Hyperparameters of the model.
checkpoint: Checkpoint to load. Notice `checkpoint.version` will decide
which TimesFM version to use.
"""
self.per_core_batch_size = per_core_batch_size
self.backend = backend
self.num_devices = jax.local_device_count(self.backend)
self.global_batch_size = self.per_core_batch_size * self.num_devices
self.hparams = hparams
# Expand hparams for conciseness within the model code.
self.context_len = hparams.context_len
self.horizon_len = hparams.horizon_len
self.input_patch_len = hparams.input_patch_len
self.output_patch_len = hparams.output_patch_len
self.num_layers = hparams.num_layers
self.model_dims = hparams.model_dims
self.backend = hparams.backend
self.quantiles = hparams.quantiles
self.num_heads = hparams.num_heads
# Rewrite these values in __post_init__ for SPMD.
self.num_cores = 1
self.per_core_batch_size = hparams.per_core_batch_size
self.global_batch_size = hparams.per_core_batch_size
self.context_len = context_len
self.horizon_len = horizon_len
self.input_patch_len = input_patch_len
self.output_patch_len = output_patch_len
self._horizon_start = self.context_len - self.input_patch_len
self.__post_init__()
self.load_from_checkpoint(checkpoint)
self.mesh_shape = [1, self.num_devices, 1]
self.mesh_name = ["replica", "data", "mdl"]
if quantiles is None:
quantiles = patched_decoder.DEFAULT_QUANTILES
self.model_p = pax_fiddle.Config(
patched_decoder.PatchedTimeSeriesDecoder,
name="patched_decoder",
horizon_len=self.output_patch_len,
patch_len=input_patch_len,
model_dims=model_dims,
hidden_dims=model_dims,
residual_block_tpl=pax_fiddle.Config(patched_decoder.ResidualBlock),
quantiles=quantiles,
use_freq=True,
stacked_transformer_params_tpl=pax_fiddle.Config(
transformers.StackedTransformer,
num_heads=16,
num_layers=num_layers,
transformer_layer_params_tpl=pax_fiddle.Config(
transformers.Transformer,
ln_tpl=pax_fiddle.Config(normalizations.RmsNorm,),
),
),
)
self._key1, self._key2 = jax.random.split(jax.random.PRNGKey(42))
self._model = None
self._train_state = None
self._pmapped_decode = None
self._verbose = verbose
self._eval_context = base_layer.JaxContext.HParams(do_eval=True)
try:
multiprocessing.set_start_method("spawn")
except RuntimeError:
print("Multiprocessing context has already been set.")
def _get_sample_inputs(self):
return {
"input_ts":
jnp.zeros(
(
self.per_core_batch_size,
self.context_len + self.output_patch_len,
),
dtype=jnp.float32,
),
"input_padding":
jnp.zeros(
(
self.per_core_batch_size,
self.context_len + self.output_patch_len,
),
dtype=jnp.float32,
),
"freq":
jnp.zeros(
(
self.per_core_batch_size,
1,
),
dtype=jnp.int32,
),
}
def load_from_checkpoint(
self,
checkpoint_path: Optional[str] = None,
repo_id: str = "google/timesfm-1.0-200m",
checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.
FLAX,
step: int | None = None,
) -> None:
"""Loads a checkpoint and compiles the decoder.
Args:
checkpoint_path: Optional path to the checkpoint directory.
repo_id: Hugging Face Hub repo id.
checkpoint_type: type of PAX checkpoint
step: step of the checkpoint to load. If `None`, load latest checkpoint.
"""
# Download the checkpoint from Hugging Face Hub if not given
if checkpoint_path is None:
checkpoint_path = path.join(snapshot_download(repo_id), "checkpoints")
# Initialize the model weights.
self._logging("Constructing model weights.")
start_time = time.time()
self._model = instantiate(self.model_p)
var_weight_hparams = self._model.abstract_init_with_metadata(
self._get_sample_inputs(), do_eval=True)
train_state_partition_specs = tasks_lib.create_state_partition_specs(
var_weight_hparams,
mesh_shape=self.mesh_shape,
mesh_axis_names=self.mesh_name,
discard_opt_states=True,
learners=None,
)
train_state_local_shapes = tasks_lib.create_state_unpadded_shapes(
var_weight_hparams,
discard_opt_states=True,
learners=None,
)
self._logging(
f"Constructed model weights in {time.time() - start_time:.2f} seconds.")
# Load the model weights.
self._logging(f"Restoring checkpoint from {checkpoint_path}.")
start_time = time.time()
self._train_state = checkpoints.restore_checkpoint(
train_state_local_shapes,
checkpoint_dir=checkpoint_path,
checkpoint_type=checkpoint_type,
state_specs=train_state_partition_specs,
step=step,
)
self._logging(
f"Restored checkpoint in {time.time() - start_time:.2f} seconds.")
self.jit_decode()
def jit_decode(self):
"""Jitting decoding function."""
# Initialize and jit the decode fn.
def _decode(inputs):
assert self._model is not None
assert self._train_state is not None
return self._model.apply(
self._train_state.mdl_vars,
inputs,
horizon_len=self.horizon_len,
output_patch_len=self.output_patch_len,
max_len=self.context_len,
return_forecast_on_context=True,
rngs={
base_layer.PARAMS: self._key1,
base_layer.RANDOM: self._key2,
},
method=self._model.decode,
)
self._logging("Jitting decoding.")
start_time = time.time()
self._pmapped_decode = jax.pmap(
_decode,
axis_name="batch",
devices=jax.devices(self.backend),
backend=self.backend,
axis_size=self.num_devices,
)
with base_layer.JaxContext.new_context(hparams=self._eval_context):
_ = self._pmapped_decode(
NestedMap({
"input_ts":
jnp.zeros(
(
self.num_devices,
self.per_core_batch_size,
self.context_len,
),
dtype=jnp.float32,
),
"input_padding":
jnp.zeros(
(
self.num_devices,
self.per_core_batch_size,
self.context_len + self.horizon_len,
),
dtype=jnp.float32,
),
"date_features":
None,
"freq":
jnp.zeros(
(self.num_devices, self.per_core_batch_size, 1),
dtype=jnp.int32,
),
}))
self._logging(f"Jitted decoding in {time.time() - start_time:.2f} seconds.")
def load_from_checkpoint(self, checkpoint: TimesFmCheckpoint) -> None:
"""Loads a checkpoint and compiles the decoder."""
raise NotImplementedError("`load_from_checkpoint` is not implemented.")
def _preprocess(self, inputs: Sequence[np.array],
freq: Sequence[int]) -> tuple[np.array, np.array, int]:
@@ -417,7 +246,7 @@ class TimesFm:
forecast_context_len: int | None = None,
return_forecast_on_context: bool = False,
truncate_negative: bool = False,
) -> tuple[JTensor, JTensor]:
) -> tuple[np.array, np.array]:
"""Forecasts on a list of time series.
Args:
@@ -443,92 +272,7 @@ class TimesFm:
Raises:
ValueError: If the checkpoint is not properly loaded.
"""
if not self._train_state or not self._model:
raise ValueError(
"Checkpoint not loaded. Call `load_from_checkpoint` before"
" `forecast`.")
if forecast_context_len is None:
forecast_context_len = self.context_len
inputs = [np.array(ts)[-forecast_context_len:] for ts in inputs]
inp_min = np.min([np.min(ts) for ts in inputs])
if window_size is not None:
new_inputs = []
for ts in inputs:
new_inputs.extend(moving_average(ts, window_size))
inputs = new_inputs
if freq is None:
logging.info("No frequency provided via `freq`. Default to high (0).")
freq = [0] * len(inputs)
input_ts, input_padding, inp_freq, pmap_pad = self._preprocess(inputs, freq)
with base_layer.JaxContext.new_context(hparams=self._eval_context):
mean_outputs = []
full_outputs = []
assert input_ts.shape[0] % self.global_batch_size == 0
for i in range(input_ts.shape[0] // self.global_batch_size):
input_ts_in = jnp.array(input_ts[i * self.global_batch_size:(i + 1) *
self.global_batch_size])
input_padding_in = jnp.array(
input_padding[i * self.global_batch_size:(i + 1) *
self.global_batch_size],)
inp_freq_in = jnp.array(
inp_freq[i * self.global_batch_size:(i + 1) *
self.global_batch_size, :],
dtype=jnp.int32,
)
pmapped_inputs = NestedMap({
"input_ts":
es.jax_einshape(
"(db)...->db...",
input_ts_in,
d=self.num_devices,
),
"input_padding":
es.jax_einshape(
"(db)...->db...",
input_padding_in,
d=self.num_devices,
),
"date_features":
None,
"freq":
es.jax_einshape(
"(db)...->db...",
inp_freq_in,
d=self.num_devices,
),
})
mean_output, full_output = self._pmapped_decode(pmapped_inputs)
if not return_forecast_on_context:
mean_output = mean_output[:, :, self._horizon_start:, ...]
full_output = full_output[:, :, self._horizon_start:, ...]
mean_output = es.jax_einshape("db...->(db)...",
mean_output,
d=self.num_devices)
full_output = es.jax_einshape("db...->(db)...",
full_output,
d=self.num_devices)
mean_output = np.array(mean_output)
full_output = np.array(full_output)
mean_outputs.append(mean_output)
full_outputs.append(full_output)
mean_outputs = np.concatenate(mean_outputs, axis=0)
full_outputs = np.concatenate(full_outputs, axis=0)
if pmap_pad > 0:
mean_outputs = mean_outputs[:-pmap_pad, ...]
full_outputs = full_outputs[:-pmap_pad, ...]
if window_size is not None:
mean_outputs = mean_outputs[0::2, ...] + mean_outputs[1::2, ...]
full_outputs = full_outputs[0::2, ...] + full_outputs[1::2, ...]
if inp_min >= 0 and truncate_negative:
mean_outputs = np.maximum(mean_outputs, 0.0)
full_outputs = np.maximum(full_outputs, 0.0)
return mean_outputs, full_outputs
raise NotImplementedError("`forecast` is not implemented.")
def forecast_with_covariates(
self,
@@ -663,7 +407,7 @@ class TimesFm:
]
per_instance_stats = None
if normalize_xreg_target_per_input:
targets, per_instance_stats = _normalize(targets)
targets, per_instance_stats = normalize(targets)
xregs = xreg_lib.BatchedInContextXRegLinear(
targets=targets,
train_lens=train_lens,
@@ -686,7 +430,7 @@ class TimesFm:
assert_covariate_shapes=True,
)
if normalize_xreg_target_per_input:
xregs = _renormalize(xregs, per_instance_stats)
xregs = renormalize(xregs, per_instance_stats)
outputs = [
(mean_output[self._horizon_start:(self._horizon_start + test_len)] +
xreg)
@@ -701,7 +445,7 @@ class TimesFm:
]
per_instance_stats = None
if normalize_xreg_target_per_input:
targets, per_instance_stats = _normalize(targets)
targets, per_instance_stats = normalize(targets)
xregs, xregs_on_context, _, _, _ = xreg_lib.BatchedInContextXRegLinear(
targets=targets,
train_lens=train_lens,
@@ -739,7 +483,7 @@ class TimesFm:
for mean_output, test_len, xreg in zip(mean_outputs, test_lens, xregs)
]
if normalize_xreg_target_per_input:
outputs = _renormalize(outputs, per_instance_stats)
outputs = renormalize(outputs, per_instance_stats)
return outputs, xregs
@@ -825,12 +569,11 @@ class TimesFm:
)
fcst_df[model_name] = full_forecast[:, 0:self.horizon_len, 0].reshape(-1, 1)
if self._model.quantiles is not None:
for i, q in enumerate(self._model.quantiles):
q_col = f"{model_name}-q-{q}"
fcst_df[q_col] = full_forecast[:, 0:self.horizon_len,
1 + i].reshape(-1, 1)
if q == 0.5:
fcst_df[model_name] = fcst_df[q_col]
for i, q in enumerate(self.quantiles):
q_col = f"{model_name}-q-{q}"
fcst_df[q_col] = full_forecast[:, 0:self.horizon_len,
1 + i].reshape(-1, 1)
if q == 0.5:
fcst_df[model_name] = fcst_df[q_col]
logging.info("Finished creating output dataframe.")
return fcst_df
src/timesfm/timesfm_jax.py
+358
View File
@@ -0,0 +1,358 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""TimesFM JAX forecast API for inference."""
import logging
import multiprocessing
import time
from os import path
from typing import Any, Sequence
import einshape as es
import jax
import jax.numpy as jnp
import numpy as np
from huggingface_hub import snapshot_download
from paxml import checkpoints, tasks_lib
from praxis import base_hyperparams, base_layer, pax_fiddle, py_utils, pytypes
from praxis.layers import normalizations, transformers
from timesfm import timesfm_base
from timesfm import patched_decoder
instantiate = base_hyperparams.instantiate
NestedMap = py_utils.NestedMap
JTensor = pytypes.JTensor
_TOL = 1e-6
class TimesFmJax(timesfm_base.TimesFmBase):
"""TimesFM forecast API for inference.
This class is the scaffolding for calling TimesFM forecast. To properly use:
1. Create an instance with the correct hyperparameters of a TimesFM model.
2. Call `load_from_checkpoint` to load a compatible checkpoint.
3. Call `forecast` for inference.
Given the model size, this API does not shard the model weights for SPMD. All
parallelism happens on the data dimension.
Compilation happens during the first time `forecast` is called and uses the
`per_core_batch_size` to set and freeze the input signature. Subsequent calls
to `forecast` reflect the actual inference latency.
"""
def _get_sample_inputs(self):
return {
"input_ts":
jnp.zeros(
(
self.per_core_batch_size,
self.context_len + self.output_patch_len,
),
dtype=jnp.float32,
),
"input_padding":
jnp.zeros(
(
self.per_core_batch_size,
self.context_len + self.output_patch_len,
),
dtype=jnp.float32,
),
"freq":
jnp.zeros(
(
self.per_core_batch_size,
1,
),
dtype=jnp.int32,
),
}
def __post_init__(self):
self.num_cores = jax.local_device_count(self.backend)
self.global_batch_size = self.per_core_batch_size * self.num_cores
self._eval_context = base_layer.JaxContext.HParams(do_eval=True)
self._pmapped_decode = None
self._model = None
self._train_state = None
def load_from_checkpoint(
self,
checkpoint: timesfm_base.TimesFmCheckpoint,
) -> None:
"""Loads a checkpoint and compiles the decoder."""
checkpoint_type = (checkpoints.CheckpointType.FLAX
if checkpoint.type is None else checkpoint.type)
checkpoint_path = checkpoint.path
step = checkpoint.step
repo_id = checkpoint.huggingface_repo_id
if checkpoint_path is None:
checkpoint_path = path.join(snapshot_download(repo_id), "checkpoints")
# Rewrite the devices for Jax.
self.mesh_shape = [1, self.num_cores, 1]
self.mesh_name = ["replica", "data", "mdl"]
self.model_p = pax_fiddle.Config(
patched_decoder.PatchedTimeSeriesDecoder,
name="patched_decoder",
horizon_len=self.output_patch_len,
patch_len=self.input_patch_len,
model_dims=self.model_dims,
hidden_dims=self.model_dims,
residual_block_tpl=pax_fiddle.Config(patched_decoder.ResidualBlock),
quantiles=self.quantiles,
use_freq=True,
stacked_transformer_params_tpl=pax_fiddle.Config(
transformers.StackedTransformer,
num_heads=self.num_heads,
num_layers=self.num_layers,
transformer_layer_params_tpl=pax_fiddle.Config(
transformers.Transformer,
ln_tpl=pax_fiddle.Config(normalizations.RmsNorm,),
),
),
)
self._key1, self._key2 = jax.random.split(jax.random.PRNGKey(42))
self._model = None
self._train_state = None
self._pmapped_decode = None
self._eval_context = base_layer.JaxContext.HParams(do_eval=True)
try:
multiprocessing.set_start_method("spawn")
except RuntimeError:
print("Multiprocessing context has already been set.")
# Download the checkpoint from Hugging Face Hub if not given
# Initialize the model weights.
self._logging("Constructing model weights.")
start_time = time.time()
self._model = instantiate(self.model_p)
var_weight_hparams = self._model.abstract_init_with_metadata(
self._get_sample_inputs(), do_eval=True)
train_state_partition_specs = tasks_lib.create_state_partition_specs(
var_weight_hparams,
mesh_shape=self.mesh_shape,
mesh_axis_names=self.mesh_name,
discard_opt_states=True,
learners=None,
)
train_state_local_shapes = tasks_lib.create_state_unpadded_shapes(
var_weight_hparams,
discard_opt_states=True,
learners=None,
)
self._logging(
f"Constructed model weights in {time.time() - start_time:.2f} seconds.")
# Load the model weights.
self._logging(f"Restoring checkpoint from {checkpoint_path}.")
start_time = time.time()
self._train_state = checkpoints.restore_checkpoint(
train_state_local_shapes,
checkpoint_dir=checkpoint_path,
checkpoint_type=checkpoint_type,
state_specs=train_state_partition_specs,
step=step,
)
self._logging(
f"Restored checkpoint in {time.time() - start_time:.2f} seconds.")
self.jit_decode()
def jit_decode(self):
"""Jitting decoding function."""
# Initialize and jit the decode fn.
def _decode(inputs):
assert self._model is not None
assert self._train_state is not None
return self._model.apply(
self._train_state.mdl_vars,
inputs,
horizon_len=self.horizon_len,
output_patch_len=self.output_patch_len,
max_len=self.context_len,
return_forecast_on_context=True,
rngs={
base_layer.PARAMS: self._key1,
base_layer.RANDOM: self._key2,
},
method=self._model.decode,
)
self._logging("Jitting decoding.")
start_time = time.time()
self._pmapped_decode = jax.pmap(
_decode,
axis_name="batch",
devices=jax.devices(self.backend),
backend=self.backend,
axis_size=self.num_cores,
)
with base_layer.JaxContext.new_context(hparams=self._eval_context):
_ = self._pmapped_decode(
NestedMap({
"input_ts":
jnp.zeros(
(
self.num_cores,
self.per_core_batch_size,
self.context_len,
),
dtype=jnp.float32,
),
"input_padding":
jnp.zeros(
(
self.num_cores,
self.per_core_batch_size,
self.context_len + self.horizon_len,
),
dtype=jnp.float32,
),
"date_features":
None,
"freq":
jnp.zeros(
(self.num_cores, self.per_core_batch_size, 1),
dtype=jnp.int32,
),
}))
self._logging(f"Jitted decoding in {time.time() - start_time:.2f} seconds.")
def forecast(
self,
inputs: Sequence[Any],
freq: Sequence[int] | None = None,
window_size: int | None = None,
forecast_context_len: int | None = None,
return_forecast_on_context: bool = False,
truncate_negative: bool = False,
) -> tuple[np.ndarray, np.ndarray]:
"""Forecasts on a list of time series.
Args:
inputs: list of time series forecast contexts. Each context time series
should be in a format convertible to JTensor by `jnp.array`.
freq: frequency of each context time series. 0 for high frequency
(default), 1 for medium, and 2 for low. Notice this is different from
the `freq` required by `forecast_on_df`.
window_size: window size of trend + residual decomposition. If None then
we do not do decomposition.
forecast_context_len: optional max context length.
return_forecast_on_context: True to return the forecast on the context
when available, i.e. after the first input patch.
truncate_negative: truncate to only non-negative values if all the contexts
have non-negative values.
Returns:
A tuple for JTensors:
- the mean forecast of size (# inputs, # forecast horizon),
- the full forecast (mean + quantiles) of size
(# inputs, # forecast horizon, 1 + # quantiles).
Raises:
ValueError: If the checkpoint is not properly loaded.
"""
if not self._train_state or not self._model:
raise ValueError(
"Checkpoint not loaded. Call `load_from_checkpoint` before"
" `forecast`.")
if forecast_context_len is None:
fcontext_len = self.context_len
else:
fcontext_len = forecast_context_len
inputs = [np.array(ts)[-fcontext_len:] for ts in inputs]
inp_min = np.min([np.min(ts) for ts in inputs])
if window_size is not None:
new_inputs = []
for ts in inputs:
new_inputs.extend(timesfm_base.moving_average(ts, window_size))
inputs = new_inputs
if freq is None:
logging.info("No frequency provided via `freq`. Default to high (0).")
freq = [0] * len(inputs)
input_ts, input_padding, inp_freq, pmap_pad = self._preprocess(inputs, freq)
with base_layer.JaxContext.new_context(hparams=self._eval_context):
mean_outputs = []
full_outputs = []
assert input_ts.shape[0] % self.global_batch_size == 0
for i in range(input_ts.shape[0] // self.global_batch_size):
input_ts_in = jnp.array(input_ts[i * self.global_batch_size:(i + 1) *
self.global_batch_size])
input_padding_in = jnp.array(
input_padding[i * self.global_batch_size:(i + 1) *
self.global_batch_size],)
inp_freq_in = jnp.array(
inp_freq[i * self.global_batch_size:(i + 1) *
self.global_batch_size, :],
dtype=jnp.int32,
)
pmapped_inputs = NestedMap({
"input_ts":
es.jax_einshape(
"(db)...->db...",
input_ts_in,
d=self.num_cores,
),
"input_padding":
es.jax_einshape(
"(db)...->db...",
input_padding_in,
d=self.num_cores,
),
"date_features":
None,
"freq":
es.jax_einshape(
"(db)...->db...",
inp_freq_in,
d=self.num_cores,
),
})
mean_output, full_output = self._pmapped_decode(pmapped_inputs)
if not return_forecast_on_context:
mean_output = mean_output[:, :, self._horizon_start:, ...]
full_output = full_output[:, :, self._horizon_start:, ...]
mean_output = es.jax_einshape("db...->(db)...",
mean_output,
d=self.num_cores)
full_output = es.jax_einshape("db...->(db)...",
full_output,
d=self.num_cores)
mean_output = np.array(mean_output)
full_output = np.array(full_output)
mean_outputs.append(mean_output)
full_outputs.append(full_output)
mean_outputs = np.concatenate(mean_outputs, axis=0)
full_outputs = np.concatenate(full_outputs, axis=0)
if pmap_pad > 0:
mean_outputs = mean_outputs[:-pmap_pad, ...]
full_outputs = full_outputs[:-pmap_pad, ...]
if window_size is not None:
mean_outputs = mean_outputs[0::2, ...] + mean_outputs[1::2, ...]
full_outputs = full_outputs[0::2, ...] + full_outputs[1::2, ...]
if inp_min >= 0 and truncate_negative:
mean_outputs = np.maximum(mean_outputs, 0.0)
full_outputs = np.maximum(full_outputs, 0.0)
return mean_outputs, full_outputs
src/timesfm/timesfm_torch.py
+171
View File
@@ -0,0 +1,171 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""TimesFM pytorch forecast API for inference."""
import logging
from os import path
from typing import Any, Sequence
import numpy as np
import torch
from huggingface_hub import snapshot_download
from timesfm import timesfm_base
from . import pytorch_patched_decoder as ppd
_TOL = 1e-6
class TimesFmTorch(timesfm_base.TimesFmBase):
"""TimesFM forecast API for inference."""
def __post_init__(self):
self._model_config = ppd.TimesFMConfig(
num_layers=self.num_layers,
num_heads=self.num_heads,
hidden_size=self.model_dims,
intermediate_size=self.model_dims,
patch_len=self.input_patch_len,
horizon_len=self.output_patch_len,
head_dim=self.model_dims // self.num_heads,
quantiles=self.quantiles,
)
self._model = None
self.num_cores = 1
self.global_batch_size = self.per_core_batch_size
self._device = torch.device("cuda:0" if (
torch.cuda.is_available() and self.backend == "gpu") else "cpu")
def load_from_checkpoint(
self,
checkpoint: timesfm_base.TimesFmCheckpoint,
) -> None:
"""Loads a checkpoint and compiles the decoder."""
checkpoint_path = checkpoint.path
repo_id = checkpoint.huggingface_repo_id
if checkpoint_path is None:
checkpoint_path = path.join(snapshot_download(repo_id),
"torch_model.ckpt")
self._model = ppd.PatchedTimeSeriesDecoder(self._model_config)
loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
logging.info("Loading checkpoint from %s", checkpoint_path)
self._model.load_state_dict(loaded_checkpoint)
logging.info("Sending checkpoint to device %s", f"{self._device}")
self._model.to(self._device)
self._model.eval()
# TODO: add compilation.
def forecast(
self,
inputs: Sequence[Any],
freq: Sequence[int] | None = None,
window_size: int | None = None,
forecast_context_len: int | None = None,
return_forecast_on_context: bool = False,
truncate_negative: bool = False,
) -> tuple[np.ndarray, np.ndarray]:
"""Forecasts on a list of time series.
Args:
inputs: list of time series forecast contexts. Each context time series
should be in a format convertible to JTensor by `jnp.array`.
freq: frequency of each context time series. 0 for high frequency
(default), 1 for medium, and 2 for low. Notice this is different from
the `freq` required by `forecast_on_df`.
window_size: window size of trend + residual decomposition. If None then
we do not do decomposition.
forecast_context_len: optional max context length.
return_forecast_on_context: True to return the forecast on the context
when available, i.e. after the first input patch.
truncate_negative: truncate to only non-negative values if all the contexts
have non-negative values.
Returns:
A tuple for JTensors:
- the mean forecast of size (# inputs, # forecast horizon),
- the full forecast (mean + quantiles) of size
(# inputs, # forecast horizon, 1 + # quantiles).
Raises:
ValueError: If the checkpoint is not properly loaded.
"""
if not self._model:
raise ValueError(
"Checkpoint not loaded. Call `load_from_checkpoint` before"
" `forecast`.")
if forecast_context_len is None:
fcontext_len = self.context_len
else:
fcontext_len = forecast_context_len
inputs = [np.array(ts)[-fcontext_len:] for ts in inputs]
inp_min = np.min([np.min(ts) for ts in inputs])
if window_size is not None:
new_inputs = []
for ts in inputs:
new_inputs.extend(timesfm_base.moving_average(ts, window_size))
inputs = new_inputs
if freq is None:
logging.info("No frequency provided via `freq`. Default to high (0).")
freq = [0] * len(inputs)
input_ts, input_padding, inp_freq, pmap_pad = self._preprocess(inputs, freq)
with torch.no_grad():
mean_outputs = []
full_outputs = []
assert input_ts.shape[0] % self.global_batch_size == 0
for i in range(input_ts.shape[0] // self.global_batch_size):
input_ts_in = torch.from_numpy(
np.array(input_ts[i * self.global_batch_size:(i + 1) *
self.global_batch_size],
dtype=np.float32)).to(self._device)
input_padding_in = torch.from_numpy(
np.array(input_padding[i * self.global_batch_size:(i + 1) *
self.global_batch_size],
dtype=np.float32)).to(self._device)
inp_freq_in = torch.from_numpy(
np.array(inp_freq[
i * self.global_batch_size:(i + 1) * self.global_batch_size,
:,
],
dtype=np.int32)).long().to(self._device)
mean_output, full_output = self._model.decode(
input_ts=input_ts_in,
paddings=input_padding_in,
freq=inp_freq_in,
horizon_len=self.horizon_len,
return_forecast_on_context=return_forecast_on_context,
)
mean_output = mean_output.detach().cpu().numpy()
full_output = full_output.detach().cpu().numpy()
mean_output = np.array(mean_output)
full_output = np.array(full_output)
mean_outputs.append(mean_output)
full_outputs.append(full_output)
mean_outputs = np.concatenate(mean_outputs, axis=0)
full_outputs = np.concatenate(full_outputs, axis=0)
if pmap_pad > 0:
mean_outputs = mean_outputs[:-pmap_pad, ...]
full_outputs = full_outputs[:-pmap_pad, ...]
if window_size is not None:
mean_outputs = mean_outputs[0::2, ...] + mean_outputs[1::2, ...]
full_outputs = full_outputs[0::2, ...] + full_outputs[1::2, ...]
if inp_min >= 0 and truncate_negative:
mean_outputs = np.maximum(mean_outputs, 0.0)
full_outputs = np.maximum(full_outputs, 0.0)
return mean_outputs, full_outputs
src/timesfm_torch/README.md
-9
View File
@@ -1,9 +0,0 @@
# Official Pytorch implementation of TimesFM
TimesFM (Time Series Foundation Model) is a pretrained time-series foundation model developed by Google
Research for time-series forecasting.
* Paper: [A decoder-only foundation model for time-series forecasting](https://arxiv.org/abs/2310.10688), to appear in ICML 2024.
* [Google Research blog](https://research.google/blog/a-decoder-only-foundation-model-for-time-series-forecasting/)
## Stay tuned for all of the functionalities as that of the pax version.