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

Merge branch 'master' of github.com:google-research/timesfm

Browse Source
This commit is contained in:
Justin Guese
2024-07-12 11:26:04 +02:00
parent 16b899c4a5 03c71634e6
commit eb0f28b3f6
20 changed files with 1956 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+24 -2
View File
@@ -24,7 +24,7 @@ timesfm-1.0-200m is the first open model checkpoint:
## Benchmarks
Please refer to our result tables on the [extended benchmarks](./experiments/extended_benchmarks/tfm_results.png) and the [long horizon benchmarks](./experiments/long_horizon_benchmarks/tfm_long_horizon.png).
Please refer to our result tables on the [extended benchmarks](https://github.com/google-research/timesfm/tree/master/experiments/extended_benchmarks) and the [long horizon benchmarks](https://github.com/google-research/timesfm/tree/master/experiments/long_horizon_benchmarks).
Please look into the README files in the respective benchmark directories within `experiments/` for instructions for running TimesFM on the respective benchmarks.
@@ -51,6 +51,8 @@ This will install the environment in the local .venv folder (depends on the conf
### Conda / GPU installation
We recommend at least 16GB RAM to load TimesFM dependencies.
For calling TimesFM, We have two environment files. Inside `timesfm`, for
GPU installation (assuming CUDA 12 has been setup), you can create a conda
environment `tfm_env` from the base folder through:
@@ -184,4 +186,24 @@ forecast_df = tfm.forecast_on_df(
freq="M", # monthly
value_name="y",
num_jobs=-1,
)```
)
```
## Finetuning
We have provided an example of finetuning the model on a new dataset in `notebooks/finetuning.ipynb`.
## Contribution Style guide
If you would like to submit a PR please make sure that you use our formatting style. We use [yapf](https://github.com/google/yapf) for formatting with the following options,
```
[style]
based_on_style = google
# Add your custom style rules here
indent_width = 2
spaces_before_comment = 2
```
Please run `yapf --in-place --recursive <filename>` on all affected files.
environment.yml
+1
View File
@@ -15,3 +15,4 @@ dependencies:
- paxml
- jax[cuda12]==0.4.26
- einshape
- scikit-learn
environment_cpu.yml
+1
View File
@@ -15,3 +15,4 @@ dependencies:
- paxml
- jax[cpu]==0.4.26
- einshape
- scikit-learn
experiments/baselines/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# 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.
experiments/baselines/timegpt_pipeline.py
+259
View File
@@ -0,0 +1,259 @@
# 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.
import os
from time import time
from typing import List, Optional, Tuple
from dotenv import load_dotenv
from gluonts.time_feature.seasonality import get_seasonality as _get_seasonality
from nixtla import NixtlaClient
import pandas as pd
from tqdm import tqdm
from utilsforecast.processing import (
backtest_splits,
drop_index_if_pandas,
join,
maybe_compute_sort_indices,
take_rows,
vertical_concat,
)
def get_seasonality(freq: str) -> int:
return _get_seasonality(freq, seasonalities={"D": 7})
def maybe_convert_col_to_datetime(
df: pd.DataFrame, col_name: str
) -> pd.DataFrame:
if not pd.api.types.is_datetime64_any_dtype(df[col_name]):
df = df.copy()
df[col_name] = pd.to_datetime(df[col_name])
return df
def zero_pad_time_series(df, freq, min_length=36):
"""If time_series length is less than min_length, front pad it with zeros."""
# 1. Calculate required padding for each unique_id
value_counts = df["unique_id"].value_counts()
to_pad = value_counts[value_counts < min_length].index
# 2. Create a new DataFrame to hold padded data
padded_data = []
for unique_id in to_pad:
# 2a. Filter data for the specific unique_id
subset = df[df["unique_id"] == unique_id]
if len(subset) > min_length:
padded_data.append(subset)
else:
# 2b. Determine earliest date and calculate padding dates
start_date = subset["ds"].min()
padding_dates = pd.date_range(
end=start_date,
periods=min_length - len(subset) + 1,
freq=freq, # 'MS' for month start
)[
:-1
] # Exclude the start_date itself
# 2c. Create padding data
padding_df = pd.DataFrame(
{"ds": padding_dates, "unique_id": unique_id, "y": 0} # Zero padding
)
# 2d. Combine original and padding data, and append to the list
padded_data.append(pd.concat([padding_df, subset]).sort_values("ds"))
# 3. Combine all padded data and original data (unchanged)
result_df = pd.concat(padded_data + [df[~df["unique_id"].isin(to_pad)]])
return result_df
class Forecaster:
"""Borrowed from
https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models.
"""
def forecast(
self,
df: pd.DataFrame,
h: int,
freq: str,
) -> pd.DataFrame:
raise NotImplementedError
def cross_validation(
self,
df: pd.DataFrame,
h: int,
freq: str,
n_windows: int = 1,
step_size: int | None = None,
) -> pd.DataFrame:
df = maybe_convert_col_to_datetime(df, "ds")
# mlforecast cv code
results = []
sort_idxs = maybe_compute_sort_indices(df, "unique_id", "ds")
if sort_idxs is not None:
df = take_rows(df, sort_idxs)
splits = backtest_splits(
df,
n_windows=n_windows,
h=h,
id_col="unique_id",
time_col="ds",
freq=pd.tseries.frequencies.to_offset(freq),
step_size=h if step_size is None else step_size,
)
for _, (cutoffs, train, valid) in tqdm(enumerate(splits)):
if len(valid.columns) > 3:
raise NotImplementedError(
"Cross validation with exogenous variables is not yet supported."
)
y_pred = self.forecast(
df=train,
h=h,
freq=freq,
)
y_pred = join(y_pred, cutoffs, on="unique_id", how="left")
result = join(
valid[["unique_id", "ds", "y"]],
y_pred,
on=["unique_id", "ds"],
)
if result.shape[0] < valid.shape[0]:
raise ValueError(
"Cross validation result produced less results than expected."
" Please verify that the frequency parameter (freq) matches your"
" series' and that there aren't any missing periods."
)
results.append(result)
out = vertical_concat(results)
out = drop_index_if_pandas(out)
first_out_cols = ["unique_id", "ds", "cutoff", "y"]
remaining_cols = [c for c in out.columns if c not in first_out_cols]
fcst_cv_df = out[first_out_cols + remaining_cols]
return fcst_cv_df
class TimeGPT(Forecaster):
"""Borrowed from
https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models.
We modify the class to take care of edge cases.
"""
def __init__(
self,
api_key: str | None = None,
base_url: Optional[str] = None,
max_retries: int = 1,
model: str = "timegpt-1",
alias: str = "TimeGPT",
):
self.api_key = api_key
self.base_url = base_url
self.max_retries = max_retries
self.model = model
self.alias = alias
def _get_client(self) -> NixtlaClient:
if self.api_key is None:
api_key = os.environ["NIXTLA_API_KEY"]
else:
api_key = self.api_key
return NixtlaClient(
api_key=api_key,
base_url=self.base_url,
max_retries=self.max_retries,
)
def forecast(
self,
df: pd.DataFrame,
h: int,
freq: str,
level: List = [90.0],
chunk_size: Optional[int] = None,
) -> pd.DataFrame:
client = self._get_client()
fcst_df = None
if chunk_size is None:
fcst_df = client.forecast(
df=df,
h=h,
freq=freq,
level=level,
model=self.model,
)
else:
all_unique_ids = df["unique_id"].unique()
all_fcst_df = []
for i in range(0, len(all_unique_ids), chunk_size):
chunk_ids = all_unique_ids[i : i + chunk_size]
chunk_df = df[df["unique_id"].isin(chunk_ids)]
fct_chunk_df = client.forecast(
df=chunk_df,
h=h,
freq=freq,
level=level,
)
all_fcst_df.append(fct_chunk_df)
fcst_df = pd.concat(all_fcst_df)
fcst_df["ds"] = pd.to_datetime(fcst_df["ds"])
replace_dict = {}
for col in fcst_df.columns:
if col.startswith("TimeGPT"):
replace_dict[col] = col.replace("TimeGPT", self.alias)
fcst_df = fcst_df.rename(columns=replace_dict)
return fcst_df
def run_timegpt(
train_df: pd.DataFrame,
horizon: int,
freq: str,
seasonality: int,
level: List[int],
dataset: str,
model: str = "timegpt-1",
) -> Tuple[pd.DataFrame, float, str]:
os.environ["NIXTLA_ID_AS_COL"] = "true"
model = TimeGPT(model="timegpt-1", alias=model)
padded_train_df = zero_pad_time_series(train_df, freq)
init_time = time()
# For these datasets the API fails if we do not chunk.
if dataset in ["m5", "m4_quarterly"]:
chunk_size = 5000
else:
chunk_size = None
fcsts_df = model.forecast(
df=padded_train_df,
h=horizon,
level=level,
freq=freq,
chunk_size=chunk_size,
)
total_time = time() - init_time
# In case levels are not returned we replace the levels with the mean predictions.
# Note that this does not affect the results table as we only compare on point
# forecastign metrics.
for lvl in level:
if f"{model.alias}-lo-{lvl}" not in fcsts_df.columns:
fcsts_df[f"{model.alias}-lo-{lvl}"] = fcsts_df[model.alias]
if f"{model.alias}-hi-{lvl}" not in fcsts_df.columns:
fcsts_df[f"{model.alias}-hi-{lvl}"] = fcsts_df[model.alias]
return fcsts_df, total_time, model.alias
experiments/environment.yml
+4
View File
@@ -22,3 +22,7 @@ dependencies:
- paxml
- jax[cuda12]==0.4.26
- einshape
- python-dotenv
- nixtla>=0.5.1
- rich
- scikit-learn
experiments/environment_cpu.yml
+4
View File
@@ -22,3 +22,7 @@ dependencies:
- paxml
- jax[cpu]==0.4.26
- einshape
- python-dotenv
- nixtla>=0.5.1
- rich
- scikit-learn
experiments/extended_benchmarks/README.md
+5 -2
View File
@@ -2,7 +2,6 @@
The benchmark setting has been borrowed from Nixtla's original [benchmarking](https://github.com/AzulGarza/nixtla/tree/main/experiments/amazon-chronos) of time-series foundation models against a strong statistical ensemble. Later more datasets were added by the Chronos team in this [pull request](https://github.com/shchur/nixtla/tree/chronos-full-eval/experiments/amazon-chronos). We compare on all the datasets in this extended benchmarks.
All experiments were performed on a [g2-standard-32](https://cloud.google.com/compute/docs/gpus).
## Running TimesFM on the benchmark
@@ -19,7 +18,11 @@ Note: In the current version of TimesFM we focus on point forecasts and therefor
## Benchmark Results
![Benchmark Results Table](./tfm_results.png)
![Benchmark Results Table](./tfm_extended_new.png)
__Update:__ We have added TimeGPT-1 to the benchmark results. We had to remove the Dominick dataset as we were not able to run TimeGPT-1 on this benchmark. Note that the previous results including Dominick remain available at `./tfm_results.png`. In order to reproduce the results for TimeGPT-1, please run `run_timegpt.py`.
_Remark:_ All baselines except the ones involving TimeGPT were run performed on a [g2-standard-32](https://cloud.google.com/compute/docs/gpus). Since TimeGPT-1 can only be accessed by an API, the time column might not reflect the true speed of the model as it also includes the communication cost. Moreover, we are not sure about the exact backend hardware for TimeGPT.
We can see that TimesFM performs the best in terms of both mase and smape. More importantly it is much faster than the other methods, in particular it is more than 600x faster than StatisticalEnsemble and 80x faster than Chronos (Large).
experiments/extended_benchmarks/run_timegpt.py
+108
View File
@@ -0,0 +1,108 @@
# 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.
"""Evaluation script for timegpt."""
import os
import sys
import time
from absl import flags
import numpy as np
import pandas as pd
from ..baselines.timegpt_pipeline import run_timegpt
from .utils import ExperimentHandler
dataset_names = [
"m1_monthly",
"m1_quarterly",
"m1_yearly",
"m3_monthly",
"m3_other",
"m3_quarterly",
"m3_yearly",
"m4_quarterly",
"m4_yearly",
"tourism_monthly",
"tourism_quarterly",
"tourism_yearly",
"nn5_daily_without_missing",
"m5",
"nn5_weekly",
"traffic",
"weather",
"australian_electricity_demand",
"car_parts_without_missing",
"cif_2016",
"covid_deaths",
"ercot",
"ett_small_15min",
"ett_small_1h",
"exchange_rate",
"fred_md",
"hospital",
]
_MODEL_NAME = flags.DEFINE_string(
"model_name",
"timegpt-1-long-horizon",
"Path to model, can also be set to timegpt-1",
)
_SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory")
QUANTILES = list(np.arange(1, 10) / 10.0)
def main():
results_list = []
run_id = np.random.randint(100000)
model_name = _MODEL_NAME.value
for dataset in dataset_names:
print(f"Evaluating model {model_name} on dataset {dataset}", flush=True)
exp = ExperimentHandler(dataset, quantiles=QUANTILES)
train_df = exp.train_df
horizon = exp.horizon
seasonality = exp.seasonality
freq = exp.freq
level = exp.level
fcsts_df, total_time, model_name = run_timegpt(
train_df=train_df,
horizon=exp.horizon,
model=model_name,
seasonality=seasonality,
freq=freq,
dataset=dataset,
level=level,
)
time_df = pd.DataFrame({"time": [total_time], "model": model_name})
fcsts_df = exp.fcst_from_level_to_quantiles(fcsts_df, model_name)
results = exp.evaluate_from_predictions(
models=[model_name], fcsts_df=fcsts_df, times_df=time_df
)
print(results, flush=True)
results_list.append(results)
results_full = pd.concat(results_list)
save_path = os.path.join(_SAVE_DIR.value, str(run_id))
print(f"Saving results to {save_path}", flush=True)
os.makedirs(save_path, exist_ok=True)
results_full.to_csv(f"{save_path}/results.csv")
if __name__ == "__main__":
FLAGS = flags.FLAGS
FLAGS(sys.argv)
main()
experiments/extended_benchmarks/run_timesfm.py
-1
View File
@@ -45,7 +45,6 @@ dataset_names = [
"nn5_weekly",
"traffic",
"weather",
"dominick",
"australian_electricity_demand",
"car_parts_without_missing",
"cif_2016",
experiments/extended_benchmarks/tfm_extended_new.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 301 KiB

experiments/long_horizon_benchmarks/run_eval.py
+1 -1
View File
@@ -24,9 +24,9 @@ import numpy as np
import pandas as pd
from paxml import checkpoints
import timesfm
from timesfm import data_loader
import torch
import tqdm
from . import data_loader
FLAGS = flags.FLAGS
notebooks/finetuning.ipynb
+612
View File
@@ -0,0 +1,612 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Importing relevant packages for finetuning"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'\n",
"os.environ['JAX_PMAP_USE_TENSORSTORE'] = 'false'"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import timesfm\n",
"import gc\n",
"import numpy as np\n",
"import pandas as pd\n",
"from timesfm import patched_decoder\n",
"from timesfm import data_loader"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from tqdm import tqdm\n",
"import dataclasses\n",
"import IPython\n",
"import IPython.display\n",
"import matplotlib as mpl\n",
"import matplotlib.pyplot as plt\n",
"mpl.rcParams['figure.figsize'] = (8, 6)\n",
"mpl.rcParams['axes.grid'] = False"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Loading TimesFM pretrained checkpoint"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"tfm = timesfm.TimesFm(\n",
" context_len=512,\n",
" horizon_len=128,\n",
" input_patch_len=32,\n",
" output_patch_len=128,\n",
" num_layers=20,\n",
" model_dims=1280,\n",
" backend=\"gpu\",\n",
")\n",
"tfm.load_from_checkpoint(repo_id=\"google/timesfm-1.0-200m\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Evaluating pretrained checkpoint on ETT datasets"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"DATA_DICT = {\n",
" \"ettm2\": {\n",
" \"boundaries\": [34560, 46080, 57600],\n",
" \"data_path\": \"../datasets/ETT-small/ETTm2.csv\",\n",
" \"freq\": \"15min\",\n",
" },\n",
" \"ettm1\": {\n",
" \"boundaries\": [34560, 46080, 57600],\n",
" \"data_path\": \"../datasets/ETT-small/ETTm1.csv\",\n",
" \"freq\": \"15min\",\n",
" },\n",
" \"etth2\": {\n",
" \"boundaries\": [8640, 11520, 14400],\n",
" \"data_path\": \"../datasets/ETT-small/ETTh2.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"etth1\": {\n",
" \"boundaries\": [8640, 11520, 14400],\n",
" \"data_path\": \"../datasets/ETT-small/ETTh1.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"elec\": {\n",
" \"boundaries\": [18413, 21044, 26304],\n",
" \"data_path\": \"../datasets/electricity/electricity.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"traffic\": {\n",
" \"boundaries\": [12280, 14036, 17544],\n",
" \"data_path\": \"../datasets/traffic/traffic.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"weather\": {\n",
" \"boundaries\": [36887, 42157, 52696],\n",
" \"data_path\": \"../datasets/weather/weather.csv\",\n",
" \"freq\": \"10min\",\n",
" },\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"dataset = \"ettm1\"\n",
"data_path = DATA_DICT[dataset][\"data_path\"]\n",
"freq = DATA_DICT[dataset][\"freq\"]\n",
"int_freq = timesfm.freq_map(freq)\n",
"boundaries = DATA_DICT[dataset][\"boundaries\"]\n",
"\n",
"data_df = pd.read_csv(open(data_path, \"r\"))\n",
"\n",
"\n",
"ts_cols = [col for col in data_df.columns if col != \"date\"]\n",
"num_cov_cols = None\n",
"cat_cov_cols = None\n",
"\n",
"context_len = 512\n",
"pred_len = 96\n",
"\n",
"num_ts = len(ts_cols)\n",
"batch_size = 16\n",
"\n",
"dtl = data_loader.TimeSeriesdata(\n",
" data_path=data_path,\n",
" datetime_col=\"date\",\n",
" num_cov_cols=num_cov_cols,\n",
" cat_cov_cols=cat_cov_cols,\n",
" ts_cols=np.array(ts_cols),\n",
" train_range=[0, boundaries[0]],\n",
" val_range=[boundaries[0], boundaries[1]],\n",
" test_range=[boundaries[1], boundaries[2]],\n",
" hist_len=context_len,\n",
" pred_len=pred_len,\n",
" batch_size=num_ts,\n",
" freq=freq,\n",
" normalize=True,\n",
" epoch_len=None,\n",
" holiday=False,\n",
" permute=True,\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_batches = dtl.tf_dataset(mode=\"train\", shift=1).batch(batch_size)\n",
"val_batches = dtl.tf_dataset(mode=\"val\", shift=pred_len)\n",
"test_batches = dtl.tf_dataset(mode=\"test\", shift=pred_len)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"for tbatch in tqdm(train_batches.as_numpy_iterator()):\n",
" pass\n",
"print(tbatch[0].shape)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### MAE on the test split for the pretrained TimesFM model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"mae_losses = []\n",
"for batch in tqdm(test_batches.as_numpy_iterator()):\n",
" past = batch[0]\n",
" actuals = batch[3]\n",
" _, forecasts = tfm.forecast(list(past), [0] * past.shape[0])\n",
" forecasts = forecasts[:, 0 : actuals.shape[1], 5]\n",
" mae_losses.append(np.abs(forecasts - actuals).mean())\n",
"\n",
"print(f\"MAE: {np.mean(mae_losses)}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Finetuning the model on the ETT dataset"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import jax\n",
"from jax import numpy as jnp\n",
"from praxis import pax_fiddle\n",
"from praxis import py_utils\n",
"from praxis import pytypes\n",
"from praxis import base_model\n",
"from praxis import optimizers\n",
"from praxis import schedules\n",
"from praxis import base_hyperparams\n",
"from praxis import base_layer\n",
"from paxml import tasks_lib\n",
"from paxml import trainer_lib\n",
"from paxml import checkpoints\n",
"from paxml import learners\n",
"from paxml import partitioning\n",
"from paxml import checkpoint_types"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# PAX shortcuts\n",
"NestedMap = py_utils.NestedMap\n",
"WeightInit = base_layer.WeightInit\n",
"WeightHParams = base_layer.WeightHParams\n",
"InstantiableParams = py_utils.InstantiableParams\n",
"JTensor = pytypes.JTensor\n",
"NpTensor = pytypes.NpTensor\n",
"WeightedScalars = pytypes.WeightedScalars\n",
"instantiate = base_hyperparams.instantiate\n",
"LayerTpl = pax_fiddle.Config[base_layer.BaseLayer]\n",
"AuxLossStruct = base_layer.AuxLossStruct\n",
"\n",
"AUX_LOSS = base_layer.AUX_LOSS\n",
"template_field = base_layer.template_field\n",
"\n",
"# Standard prng key names\n",
"PARAMS = base_layer.PARAMS\n",
"RANDOM = base_layer.RANDOM\n",
"\n",
"key = jax.random.PRNGKey(seed=1234)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"model = pax_fiddle.Config(\n",
" patched_decoder.PatchedDecoderFinetuneModel,\n",
" name='patched_decoder_finetune',\n",
" core_layer_tpl=tfm.model_p,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### We will hold the transformer layers fixed while finetuning, while training all other components."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@pax_fiddle.auto_config\n",
"def build_learner() -> learners.Learner:\n",
" return pax_fiddle.Config(\n",
" learners.Learner,\n",
" name='learner',\n",
" loss_name='avg_qloss',\n",
" optimizer=optimizers.Adam(\n",
" epsilon=1e-7,\n",
" clip_threshold=1e2,\n",
" learning_rate=1e-2,\n",
" lr_schedule=pax_fiddle.Config(\n",
" schedules.Cosine,\n",
" initial_value=1e-3,\n",
" final_value=1e-4,\n",
" total_steps=40000,\n",
" ),\n",
" ema_decay=0.9999,\n",
" ),\n",
" # Linear probing i.e we hold the transformer layers fixed.\n",
" bprop_variable_exclusion=['.*/stacked_transformer_layer/.*'],\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"task_p = tasks_lib.SingleTask(\n",
" name='ts-learn',\n",
" model=model,\n",
" train=tasks_lib.SingleTask.Train(\n",
" learner=build_learner(),\n",
" ),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"task_p.model.ici_mesh_shape = [1, 1, 1]\n",
"task_p.model.mesh_axis_names = ['replica', 'data', 'mdl']\n",
"\n",
"DEVICES = np.array(jax.devices()).reshape([1, 1, 1])\n",
"MESH = jax.sharding.Mesh(DEVICES, ['replica', 'data', 'mdl'])\n",
"\n",
"num_devices = jax.local_device_count()\n",
"print(f'num_devices: {num_devices}')\n",
"print(f'device kind: {jax.local_devices()[0].device_kind}')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"jax_task = task_p\n",
"key, init_key = jax.random.split(key)\n",
"\n",
"# To correctly prepare a batch of data for model initialization (now that shape\n",
"# inference is merged), we take one devices*batch_size tensor tuple of data,\n",
"# slice out just one batch, then run the prepare_input_batch function over it.\n",
"\n",
"\n",
"def process_train_batch(batch):\n",
" past_ts = batch[0].reshape(batch_size * num_ts, -1)\n",
" actual_ts = batch[3].reshape(batch_size * num_ts, -1)\n",
" return NestedMap(input_ts=past_ts, actual_ts=actual_ts)\n",
"\n",
"\n",
"def process_eval_batch(batch):\n",
" past_ts = batch[0]\n",
" actual_ts = batch[3]\n",
" return NestedMap(input_ts=past_ts, actual_ts=actual_ts)\n",
"\n",
"\n",
"jax_model_states, _ = trainer_lib.initialize_model_state(\n",
" jax_task,\n",
" init_key,\n",
" process_train_batch(tbatch),\n",
" checkpoint_type=checkpoint_types.CheckpointType.GDA,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Setting the initial model weights to the pretrained TimesFM parameters."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"jax_model_states.mdl_vars['params']['core_layer'] = tfm._train_state.mdl_vars['params']\n",
"jax_vars = jax_model_states.mdl_vars\n",
"gc.collect()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Training loop"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"jax_task = task_p\n",
"\n",
"\n",
"def train_step(states, prng_key, inputs):\n",
" return trainer_lib.train_step_single_learner(\n",
" jax_task, states, prng_key, inputs\n",
" )\n",
"\n",
"\n",
"def eval_step(states, prng_key, inputs):\n",
" states = states.to_eval_state()\n",
" return trainer_lib.eval_step_single_learner(\n",
" jax_task, states, prng_key, inputs\n",
" )\n",
"\n",
"key, train_key, eval_key = jax.random.split(key, 3)\n",
"train_prng_seed = jax.random.split(train_key, num=jax.local_device_count())\n",
"eval_prng_seed = jax.random.split(eval_key, num=jax.local_device_count())\n",
"\n",
"p_train_step = jax.pmap(train_step, axis_name='batch')\n",
"p_eval_step = jax.pmap(eval_step, axis_name='batch')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"replicated_jax_states = trainer_lib.replicate_model_state(jax_model_states)\n",
"replicated_jax_vars = replicated_jax_states.mdl_vars"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"best_eval_loss = 1e7\n",
"step_count = 0\n",
"patience = 0\n",
"NUM_EPOCHS = 100\n",
"PATIENCE = 5\n",
"TRAIN_STEPS_PER_EVAL = 1000\n",
"CHECKPOINT_DIR='/home/senrajat_google_com/ettm1_finetune'"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def reshape_batch_for_pmap(batch, num_devices):\n",
" def _reshape(input_tensor):\n",
" bsize = input_tensor.shape[0]\n",
" residual_shape = list(input_tensor.shape[1:])\n",
" nbsize = bsize // num_devices\n",
" return jnp.reshape(input_tensor, [num_devices, nbsize] + residual_shape)\n",
"\n",
" return jax.tree.map(_reshape, batch)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"for epoch in range(NUM_EPOCHS):\n",
" print(f\"__________________Epoch: {epoch}__________________\", flush=True)\n",
" train_its = train_batches.as_numpy_iterator()\n",
" if patience >= PATIENCE:\n",
" print(\"Early stopping.\", flush=True)\n",
" break\n",
" for batch in tqdm(train_its):\n",
" train_losses = []\n",
" if patience >= PATIENCE:\n",
" print(\"Early stopping.\", flush=True)\n",
" break\n",
" tbatch = process_train_batch(batch)\n",
" tbatch = reshape_batch_for_pmap(tbatch, num_devices)\n",
" replicated_jax_states, step_fun_out = p_train_step(\n",
" replicated_jax_states, train_prng_seed, tbatch\n",
" )\n",
" train_losses.append(step_fun_out.loss[0])\n",
" if step_count % TRAIN_STEPS_PER_EVAL == 0:\n",
" print(\n",
" f\"Train loss at step {step_count}: {np.mean(train_losses)}\",\n",
" flush=True,\n",
" )\n",
" train_losses = []\n",
" print(\"Starting eval.\", flush=True)\n",
" val_its = val_batches.as_numpy_iterator()\n",
" eval_losses = []\n",
" for ev_batch in tqdm(val_its):\n",
" ebatch = process_eval_batch(ev_batch)\n",
" ebatch = reshape_batch_for_pmap(ebatch, num_devices)\n",
" _, step_fun_out = p_eval_step(\n",
" replicated_jax_states, eval_prng_seed, ebatch\n",
" )\n",
" eval_losses.append(step_fun_out.loss[0])\n",
" mean_loss = np.mean(eval_losses)\n",
" print(f\"Eval loss at step {step_count}: {mean_loss}\", flush=True)\n",
" if mean_loss < best_eval_loss or np.isnan(mean_loss):\n",
" best_eval_loss = mean_loss\n",
" print(\"Saving checkpoint.\")\n",
" jax_state_for_saving = py_utils.maybe_unreplicate_for_fully_replicated(\n",
" replicated_jax_states\n",
" )\n",
" checkpoints.save_checkpoint(\n",
" jax_state_for_saving, CHECKPOINT_DIR, overwrite=True\n",
" )\n",
" patience = 0\n",
" del jax_state_for_saving\n",
" gc.collect()\n",
" else:\n",
" patience += 1\n",
" print(f\"patience: {patience}\")\n",
" step_count += 1"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Loading and evaluating the best (according to validation loss) finetuned checkpoint"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_state = checkpoints.restore_checkpoint(jax_model_states, CHECKPOINT_DIR)\n",
"print(train_state.step)\n",
"tfm._train_state.mdl_vars['params'] = train_state.mdl_vars['params']['core_layer']\n",
"tfm.jit_decode()\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"mae_losses = []\n",
"for batch in tqdm(test_batches.as_numpy_iterator()):\n",
" past = batch[0]\n",
" actuals = batch[3]\n",
" _, forecasts = tfm.forecast(list(past), [0] * past.shape[0])\n",
" forecasts = forecasts[:, 0 : actuals.shape[1], 5]\n",
" mae_losses.append(np.abs(forecasts - actuals).mean())\n",
"\n",
"print(f\"MAE: {np.mean(mae_losses)}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## There is around a __9%__ reduction in MAE from finetuning."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "tfm_env_v3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.14"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
pyproject.toml
+17 -2
View File
@@ -1,10 +1,25 @@
[tool.poetry]
name = "timesfm"
# This project can be installed with `python3 -m pip install -e .` from the main directory.
[project]
name = "timesfm-jax"
packages = [
{ include = "*", from = "src" },
]
version = "0.0.1"
description = "Open weights time-series foundation model from Google Research."
version = "1.0.1"
dependencies = [
"jax==0.4.26",
"paxml==1.4.0",
"praxis==1.4.0",
"jaxlib==0.4.26",
"numpy==1.26.4",
"pandas==2.1.4",
"einshape==1.0.0",
"utilsforecast==0.1.10",
"huggingface_hub[cli]==0.23.0",
"scikit-learn==1.5.1",
]
authors = [
"Rajat Sen <senrajat@google.com>",
"Yichen Zhou <yichenzhou@google.com>",
src/__init__.py → src/timesfm/__init__.py
-3
View File
@@ -14,7 +14,4 @@
"""TimesFM init file."""
from __future__ import absolute_import
from .patched_decoder import PatchedTimeSeriesDecoder
from .timesfm import TimesFm, freq_map
experiments/long_horizon_benchmarks/data_loader.py → src/timesfm/data_loader.py
View File
src/patched_decoder.py → src/timesfm/patched_decoder.py
+105 -11
View File
@@ -24,6 +24,7 @@ import einshape as es
from jax import lax
import jax.numpy as jnp
from praxis import base_layer
from praxis import base_model
from praxis import layers
from praxis import pax_fiddle
from praxis import py_utils
@@ -49,6 +50,7 @@ DEFAULT_QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
# NestedMap keys
_INPUT_TS = "input_ts"
_TARGET_FUTURE = "actual_ts"
_INPUT_PADDING = "input_padding"
_OUTPUT_TS = "output_ts"
_FREQ = "freq"
@@ -324,15 +326,19 @@ 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)
input_padding = jnp.where(
jnp.abs(input_ts - PAD_VAL) < _TOLERANCE, 1, input_padding
)
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
)
patched_pads = jnp.where(
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)
concat_inputs = jnp.concatenate([patched_inputs, patched_pads], axis=-1)
@@ -404,6 +410,7 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
horizon_len: int,
output_patch_len: Optional[int] = None,
max_len: int = 512,
return_forecast_on_context: bool = False,
) -> tuple[JTensor, JTensor]:
"""Auto-regressive decoding without caching.
@@ -414,15 +421,19 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
output_patch_len: output length to be fetched from one step of
auto-regressive decoding.
max_len: maximum training context length.
return_forecast_on_context: whether to return the model forecast on the
context except the first input patch.
Returns:
Tuple of two forecasting results:
- Point (mean) output predictions as a tensor with shape B x H.
- Point (mean) output predictions as a tensor with shape B x H'.
- Full predictions (mean and quantiles) as a tensor with shape
B x H x (1 + # quantiles).
B x H' x (1 + # quantiles).
In particular, if return_forecast_on_context is True, H' is H plus
the forecastable context length, i.e. context_len - (first) patch_len.
"""
final_out = inputs[_INPUT_TS]
inp_time_len = final_out.shape[1]
context_len = final_out.shape[1]
paddings = inputs[_INPUT_PADDING]
if self.use_freq:
freq = inputs[_FREQ].astype(jnp.int32)
@@ -439,7 +450,7 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
num_decode_patches = (
horizon_len + output_patch_len - 1
) // output_patch_len
for _ in range(num_decode_patches):
for step_index in range(num_decode_patches):
current_padding = paddings[:, 0 : final_out.shape[1]]
input_ts = final_out[:, -max_len:]
input_padding = current_padding[:, -max_len:]
@@ -449,13 +460,96 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
freq=freq,
)
fprop_outputs = self(model_input)[_OUTPUT_TS]
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 = es.jax_einshape("bnph->b(np)h", new_full_ts)
full_outputs.append(new_full_ts)
# (full batch, last patch, output_patch_len, index of mean forecast = 0)
new_ts = fprop_outputs[:, -1, :output_patch_len, 0]
new_full_ts = fprop_outputs[:, -1, :output_patch_len, :]
# (full batch, last patch, output_patch_len, all output indices)
full_outputs.append(fprop_outputs[:, -1, :output_patch_len, :])
full_outputs.append(new_full_ts)
final_out = jnp.concatenate([final_out, new_ts], axis=-1)
return (
final_out[:, inp_time_len : inp_time_len + horizon_len],
jnp.concatenate(full_outputs, axis=1)[:, 0:horizon_len, :],
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), :
]
else:
# `full_outputs` indexing starts at the forecast horizon.
full_outputs = jnp.concatenate(full_outputs, axis=1)[:, 0:horizon_len, :]
return (full_outputs[:, :, 0], full_outputs)
class PatchedDecoderFinetuneModel(base_model.BaseModel):
"""Model class for finetuning patched time-series decoder.
Attributes:
core_layer_tpl: config for core layer.
freq: freq to finetune on.
"""
core_layer_tpl: LayerTpl = template_field(PatchedTimeSeriesDecoder)
freq: int = 0
def setup(self) -> None:
self.create_child("core_layer", self.core_layer_tpl)
def compute_predictions(self, input_batch: NestedMap) -> NestedMap:
input_ts = input_batch[_INPUT_TS]
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
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
)
freq = jnp.ones([input_ts.shape[0], 1], dtype=jnp.int32) * self.freq
new_input_batch = NestedMap(
input_ts=input_ts,
input_padding=input_padding,
freq=freq,
)
return self.core_layer(new_input_batch)
def _quantile_loss(
self, pred: JTensor, actual: JTensor, quantile: float
) -> JTensor:
"""Calculates quantile loss.
Args:
pred: B x T
actual: B x T
quantile: quantile at which loss is computed.
Returns:
per coordinate loss.
"""
dev = actual - pred
loss_first = dev * quantile
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]:
output_ts = prediction_output[_OUTPUT_TS]
actual_ts = input_batch[_TARGET_FUTURE]
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)
loss = loss.mean()
loss_weight = jnp.array(1.0, dtype=jnp.float32)
per_example_out = NestedMap()
return {"avg_qloss": (loss, loss_weight)}, per_example_out
experiments/long_horizon_benchmarks/time_features.py → src/timesfm/time_features.py
View File
src/timesfm.py → src/timesfm/timesfm.py
+270 -2
View File
@@ -14,6 +14,7 @@
"""TimesFM forecast API for inference."""
import collections
import logging
import multiprocessing
from os import path
@@ -21,11 +22,11 @@ import time
from typing import Any, Literal, Optional, 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 huggingface_hub import snapshot_download
from paxml import checkpoints
from paxml import tasks_lib
from praxis import base_hyperparams
@@ -35,12 +36,19 @@ from praxis import py_utils
from praxis import pytypes
from praxis.layers import normalizations
from praxis.layers import transformers
import patched_decoder
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
def process_group(key, group, value_name, forecast_context_len):
@@ -78,6 +86,20 @@ def freq_map(freq: str):
raise ValueError(f"Invalid frequency: {freq}")
# Per time series normalization: forward.
def _normalize(batch):
stats = [
(np.mean(x), np.where((w := np.std(x)) > _TOL, w, 1.0)) for x in batch
]
new_batch = [(x - stat[0]) / stat[1] for x, stat in zip(batch, stats)]
return new_batch, stats
# Per time series normalization: inverse.
def _renormalize(batch, stats):
return [x * stat[1] + stat[0] for x, stat in zip(batch, stats)]
class TimesFm:
"""TimesFM forecast API for inference.
@@ -157,6 +179,7 @@ class TimesFm:
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.mesh_shape = [1, self.num_devices, 1]
self.mesh_name = ["replica", "data", "mdl"]
@@ -277,6 +300,10 @@ class TimesFm:
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):
@@ -288,6 +315,7 @@ class TimesFm:
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,
@@ -397,6 +425,7 @@ class TimesFm:
freq: Sequence[int] | None = None,
window_size: int | None = None,
forecast_context_len: int | None = None,
return_forecast_on_context: bool = False,
) -> tuple[JTensor, JTensor]:
"""Forecasts on a list of time series.
@@ -409,6 +438,8 @@ class TimesFm:
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.
Returns:
A tuple for JTensors:
@@ -480,6 +511,9 @@ class TimesFm:
),
})
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
)
@@ -506,6 +540,240 @@ class TimesFm:
full_outputs = np.maximum(full_outputs, 0.0)
return mean_outputs, full_outputs
def forecast_with_covariates(
self,
inputs: list[Sequence[float]],
dynamic_numerical_covariates: (
dict[str, Sequence[Sequence[float]]] | None
) = None,
dynamic_categorical_covariates: (
dict[str, Sequence[Sequence[Category]]] | None
) = None,
static_numerical_covariates: dict[str, Sequence[float]] | None = None,
static_categorical_covariates: (
dict[str, Sequence[Category]] | None
) = None,
freq: Sequence[int] | None = None,
window_size: int | None = None,
forecast_context_len: int | None = None,
xreg_mode: XRegMode = "xreg + timesfm",
normalize_xreg_target_per_input: bool = True,
ridge: float = 0.0,
max_rows_per_col: int = 0,
force_on_cpu: bool = False,
):
"""Forecasts on a list of time series with covariates.
To optimize inference speed, avoid string valued categorical covariates.
Args:
inputs: A list of time series forecast contexts. Each context time series
should be in a format convertible to JTensor by `jnp.array`.
dynamic_numerical_covariates: A dict of dynamic numerical covariates.
dynamic_categorical_covariates: A dict of dynamic categorical covariates.
static_numerical_covariates: A dict of static numerical covariates.
static_categorical_covariates: A dict of static categorical covariates.
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.
xreg_mode: one of "xreg + timesfm" or "timesfm + xreg". "xreg + timesfm"
fits a model on the residuals of the TimesFM forecast. "timesfm + xreg"
fits a model on the targets then forecasts on the residuals via TimesFM.
normalize_xreg_target_per_input: whether to normalize the xreg target per
input in the given batch.
ridge: ridge penalty for the linear model.
max_rows_per_col: max number of rows per column for the linear model.
force_on_cpu: whether to force running on cpu for the linear model.
Returns:
A tuple of two lists. The first is the outputs of the model. The second is
the outputs of the xreg.
"""
# Verify and bookkeep covariates.
if not (
dynamic_numerical_covariates
or dynamic_categorical_covariates
or static_numerical_covariates
or static_categorical_covariates
):
raise ValueError(
"At least one of dynamic_numerical_covariates,"
" dynamic_categorical_covariates, static_numerical_covariates,"
" static_categorical_covariates must be set."
)
# Track the lengths of (1) each input, (2) the part that can be used in the
# linear model, and (3) the horizon.
input_lens, train_lens, test_lens = [], [], []
for i, input_ts in enumerate(inputs):
input_len = len(input_ts)
input_lens.append(input_len)
if xreg_mode == "timesfm + xreg":
# For fitting residuals, no TimesFM forecast on the first patch.
train_lens.append(max(0, input_len - self.input_patch_len))
elif xreg_mode == "xreg + timesfm":
train_lens.append(input_len)
else:
raise ValueError(f"Unsupported mode: {xreg_mode}")
if dynamic_numerical_covariates:
test_lens.append(
len(list(dynamic_numerical_covariates.values())[0][i]) - input_len
)
elif dynamic_categorical_covariates:
test_lens.append(
len(list(dynamic_categorical_covariates.values())[0][i]) - input_len
)
else:
test_lens.append(self.horizon_len)
if test_lens[-1] > self.horizon_len:
raise ValueError(
"Forecast requested longer horizon than the model definition "
f"supports: {test_lens[-1]} vs {self.horizon_len}."
)
# Prepare the covariates into train and test.
train_dynamic_numerical_covariates = collections.defaultdict(list)
test_dynamic_numerical_covariates = collections.defaultdict(list)
train_dynamic_categorical_covariates = collections.defaultdict(list)
test_dynamic_categorical_covariates = collections.defaultdict(list)
for covariates, train_covariates, test_covariates in (
(
dynamic_numerical_covariates,
train_dynamic_numerical_covariates,
test_dynamic_numerical_covariates,
),
(
dynamic_categorical_covariates,
train_dynamic_categorical_covariates,
test_dynamic_categorical_covariates,
),
):
if not covariates:
continue
for covariate_name, covariate_values in covariates.items():
for input_len, train_len, covariate_value in zip(
input_lens, train_lens, covariate_values
):
train_covariates[covariate_name].append(
covariate_value[(input_len - train_len) : input_len]
)
test_covariates[covariate_name].append(covariate_value[input_len:])
# Fit models.
if xreg_mode == "timesfm + xreg":
# Forecast via TimesFM then fit a model on the residuals.
mean_outputs, _ = self.forecast(
inputs,
freq,
window_size,
forecast_context_len,
return_forecast_on_context=True,
)
targets = [
(
np.array(input_ts)[-train_len:]
- mean_output[
(self._horizon_start - train_len) : self._horizon_start
]
)
for input_ts, mean_output, train_len in zip(
inputs, mean_outputs, train_lens
)
]
per_instance_stats = None
if normalize_xreg_target_per_input:
targets, per_instance_stats = _normalize(targets)
xregs = xreg_lib.BatchedInContextXRegLinear(
targets=targets,
train_lens=train_lens,
test_lens=test_lens,
train_dynamic_numerical_covariates=train_dynamic_numerical_covariates,
test_dynamic_numerical_covariates=test_dynamic_numerical_covariates,
train_dynamic_categorical_covariates=train_dynamic_categorical_covariates,
test_dynamic_categorical_covariates=test_dynamic_categorical_covariates,
static_numerical_covariates=static_numerical_covariates,
static_categorical_covariates=static_categorical_covariates,
).fit(
ridge=ridge,
one_hot_encoder_drop=None if ridge > 0 else "first",
max_rows_per_col=max_rows_per_col,
force_on_cpu=force_on_cpu,
debug_info=False,
assert_covariates=True,
assert_covariate_shapes=True,
)
if normalize_xreg_target_per_input:
xregs = _renormalize(xregs, per_instance_stats)
outputs = [
(
mean_output[
self._horizon_start : (self._horizon_start + test_len)
]
+ xreg
)
for mean_output, test_len, xreg in zip(mean_outputs, test_lens, xregs)
]
else:
# Fit a model on the targets then forecast on the residuals via TimesFM.
targets = [
np.array(input_ts)[-train_len:]
for input_ts, train_len in zip(inputs, train_lens)
]
per_instance_stats = None
if normalize_xreg_target_per_input:
targets, per_instance_stats = _normalize(targets)
xregs, xregs_on_context, _, _, _ = xreg_lib.BatchedInContextXRegLinear(
targets=targets,
train_lens=train_lens,
test_lens=test_lens,
train_dynamic_numerical_covariates=train_dynamic_numerical_covariates,
test_dynamic_numerical_covariates=test_dynamic_numerical_covariates,
train_dynamic_categorical_covariates=train_dynamic_categorical_covariates,
test_dynamic_categorical_covariates=test_dynamic_categorical_covariates,
static_numerical_covariates=static_numerical_covariates,
static_categorical_covariates=static_categorical_covariates,
).fit(
ridge=ridge,
one_hot_encoder_drop=None if ridge > 0 else "first",
max_rows_per_col=max_rows_per_col,
force_on_cpu=force_on_cpu,
debug_info=True,
assert_covariates=True,
assert_covariate_shapes=True,
)
mean_outputs, _ = self.forecast(
[
target - xreg_on_context
for target, xreg_on_context in zip(targets, xregs_on_context)
],
freq,
window_size,
forecast_context_len,
return_forecast_on_context=True,
)
outputs = [
(
mean_output[
self._horizon_start : (self._horizon_start + test_len)
]
+ xreg
)
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)
return outputs, xregs
def forecast_on_df(
self,
inputs: pd.DataFrame,
src/timesfm/xreg_lib.py
+532
View File
@@ -0,0 +1,532 @@
# 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.
"""Helper functions for in-context covariates and regression."""
import itertools
import math
from typing import Any, Iterable, Literal, Mapping, Sequence
import jax
import jax.numpy as jnp
import numpy as np
from sklearn import preprocessing
Category = int | str
_TOL = 1e-6
XRegMode = Literal["timesfm + xreg", "xreg + timesfm"]
def _unnest(nested: Sequence[Sequence[Any]]) -> np.ndarray:
return np.array(list(itertools.chain.from_iterable(nested)))
def _repeat(elements: Iterable[Any], counts: Iterable[int]) -> np.ndarray:
return np.array(
list(
itertools.chain.from_iterable(map(itertools.repeat, elements, counts))
)
)
def _to_padded_jax_array(x: np.ndarray) -> jax.Array:
if x.ndim == 1:
(i,) = x.shape
di = 2 ** math.ceil(math.log2(i)) - i
return jnp.pad(x, ((0, di),), mode="constant", constant_values=0.0)
elif x.ndim == 2:
i, j = x.shape
di = 2 ** math.ceil(math.log2(i)) - i
dj = 2 ** math.ceil(math.log2(j)) - j
return jnp.pad(x, ((0, di), (0, dj)), mode="constant", constant_values=0.0)
else:
raise ValueError(f"Unsupported array shape: {x.shape}")
class BatchedInContextXRegBase:
"""Helper class for in-context regression covariate formatting.
Attributes:
targets: List of targets (responses) of the in-context regression.
train_lens: List of lengths of each target vector from the context.
test_lens: List of lengths of each forecast horizon.
train_dynamic_numerical_covariates: Dict of covariate names mapping to the
dynamic numerical covariates of each forecast task on the context. Their
lengths should match the corresponding lengths in `train_lens`.
train_dynamic_categorical_covariates: Dict of covariate names mapping to the
dynamic categorical covariates of each forecast task on the context. Their
lengths should match the corresponding lengths in `train_lens`.
test_dynamic_numerical_covariates: Dict of covariate names mapping to the
dynamic numerical covariates of each forecast task on the horizon. Their
lengths should match the corresponding lengths in `test_lens`.
test_dynamic_categorical_covariates: Dict of covariate names mapping to the
dynamic categorical covariates of each forecast task on the horizon. Their
lengths should match the corresponding lengths in `test_lens`.
static_numerical_covariates: Dict of covariate names mapping to the static
numerical covariates of each forecast task.
static_categorical_covariates: Dict of covariate names mapping to the static
categorical covariates of each forecast task.
"""
def __init__(
self,
targets: Sequence[Sequence[float]],
train_lens: Sequence[int],
test_lens: Sequence[int],
train_dynamic_numerical_covariates: (
Mapping[str, Sequence[Sequence[float]]] | None
) = None,
train_dynamic_categorical_covariates: (
Mapping[str, Sequence[Sequence[Category]]] | None
) = None,
test_dynamic_numerical_covariates: (
Mapping[str, Sequence[Sequence[float]]] | None
) = None,
test_dynamic_categorical_covariates: (
Mapping[str, Sequence[Sequence[Category]]] | None
) = None,
static_numerical_covariates: Mapping[str, Sequence[float]] | None = None,
static_categorical_covariates: (
Mapping[str, Sequence[Category]] | None
) = None,
) -> None:
"""Initializes with the exogenous covariate inputs.
Here we use model fitting language to refer to the context as 'train' and
the horizon as 'test'. We assume batched inputs. To properly format the
request:
- `train_lens` represents the contexts in the batch. Targets and all train
dynamic covariates should have the same lengths as the corresponding
elements
in `train_lens`. Notice each `train_len` can be different from the exact
length of the corresponding context depending on how much of the context is
used for fitting the in-context model.
- `test_lens` represents the horizon lengths in the batch. All tesdt
dynamic
covariates should have the same lengths as the corresponding elements in
`test_lens`.
- Static covariates should be one for each input.
- For train and test dynamic covariates, they should have the same
covariate
names.
Pass an empty dict {} for a covariate type if it is not present.
Example:
Here is a set of valid inputs whose schema can be used for reference.
```
targets = [
[0.0, 0.1, 0.2],
[0.0, 0.1, 0.2, 0.3],
] # Two inputs in this batch.
train_lens = [3, 4]
test_lens = [2, 5] # Forecast horizons 2 and 5 respectively.
train_dynamic_numerical_covariates = {
"cov_1_dn": [[0.0, 0.5, 1.0], [0.0, 0.5, 1.0, 1.5]],
"cov_2_dn": [[0.0, 1.5, 1.0], [0.0, 1.5, 1.0, 2.5]],
} # Each train dynamic covariate has 3 and 4 elements respectively.
test_dynamic_numerical_covariates = {
"cov_1_dn": [[0.1, 0.6], [0.1, 0.6, 1.1, 1.6, 2.4]],
"cov_2_dn": [[0.1, 1.1], [0.1, 1.6, 1.1, 2.6, 10.0]],
} # Each test dynamic covariate has 2 and 5 elements respectively.
train_dynamic_categorical_covariates = {
"cov_1_dc": [[0, 1, 0], [0, 1, 2, 3]],
"cov_2_dc": [["good", "bad", "good"], ["good", "good", "bad",
"bad"]],
}
test_dynamic_categorical_covariates = {
"cov_1_dc": [[1, 0], [1, 0, 2, 3, 1]],
"cov_2_dc": [["bad", "good"], ["bad", "bad", "bad", "bad", "bad"]],
}
static_numerical_covariates = {
"cov_1_sn": [0.0, 3.0],
"cov_2_sn": [2.0, 1.0],
"cov_3_sn": [1.0, 2.0],
} # Each static covariate has 1 element for each input.
static_categorical_covariates = {
"cov_1_sc": ["apple", "orange"],
"cov_2_sc": [2, 3],
}
```
Args:
targets: List of targets (responses) of the in-context regression.
train_lens: List of lengths of each target vector from the context.
test_lens: List of lengths of each forecast horizon.
train_dynamic_numerical_covariates: Dict of covariate names mapping to the
dynamic numerical covariates of each forecast task on the context. Their
lengths should match the corresponding lengths in `train_lens`.
train_dynamic_categorical_covariates: Dict of covariate names mapping to
the dynamic categorical covariates of each forecast task on the context.
Their lengths should match the corresponding lengths in `train_lens`.
test_dynamic_numerical_covariates: Dict of covariate names mapping to the
dynamic numerical covariates of each forecast task on the horizon. Their
lengths should match the corresponding lengths in `test_lens`.
test_dynamic_categorical_covariates: Dict of covariate names mapping to
the dynamic categorical covariates of each forecast task on the horizon.
Their lengths should match the corresponding lengths in `test_lens`.
static_numerical_covariates: Dict of covariate names mapping to the static
numerical covariates of each forecast task.
static_categorical_covariates: Dict of covariate names mapping to the
static categorical covariates of each forecast task.
"""
self.targets = targets
self.train_lens = train_lens
self.test_lens = test_lens
self.train_dynamic_numerical_covariates = (
train_dynamic_numerical_covariates or {}
)
self.train_dynamic_categorical_covariates = (
train_dynamic_categorical_covariates or {}
)
self.test_dynamic_numerical_covariates = (
test_dynamic_numerical_covariates or {}
)
self.test_dynamic_categorical_covariates = (
test_dynamic_categorical_covariates or {}
)
self.static_numerical_covariates = static_numerical_covariates or {}
self.static_categorical_covariates = static_categorical_covariates or {}
def _assert_covariates(self, assert_covariate_shapes: bool = False) -> None:
"""Verifies the validity of the covariate inputs."""
# Check presence.
if (
self.train_dynamic_numerical_covariates
and not self.test_dynamic_numerical_covariates
) or (
not self.train_dynamic_numerical_covariates
and self.test_dynamic_numerical_covariates
):
raise ValueError(
"train_dynamic_numerical_covariates and"
" test_dynamic_numerical_covariates must be both present or both"
" absent."
)
if (
self.train_dynamic_categorical_covariates
and not self.test_dynamic_categorical_covariates
) or (
not self.train_dynamic_categorical_covariates
and self.test_dynamic_categorical_covariates
):
raise ValueError(
"train_dynamic_categorical_covariates and"
" test_dynamic_categorical_covariates must be both present or both"
" absent."
)
# Check keys.
for dict_a, dict_b, dict_a_name, dict_b_name in (
(
self.train_dynamic_numerical_covariates,
self.test_dynamic_numerical_covariates,
"train_dynamic_numerical_covariates",
"test_dynamic_numerical_covariates",
),
(
self.train_dynamic_categorical_covariates,
self.test_dynamic_categorical_covariates,
"train_dynamic_categorical_covariates",
"test_dynamic_categorical_covariates",
),
):
if w := set(dict_a.keys()) - set(dict_b.keys()):
raise ValueError(
f"{dict_a_name} has keys not present in {dict_b_name}: {w}"
)
if w := set(dict_b.keys()) - set(dict_a.keys()):
raise ValueError(
f"{dict_b_name} has keys not present in {dict_a_name}: {w}"
)
# Check shapes.
if assert_covariate_shapes:
if len(self.targets) != len(self.train_lens):
raise ValueError(
"targets and train_lens must have the same number of elements."
)
if len(self.train_lens) != len(self.test_lens):
raise ValueError(
"train_lens and test_lens must have the same number of elements."
)
for i, (target, train_len) in enumerate(
zip(self.targets, self.train_lens)
):
if len(target) != train_len:
raise ValueError(
f"targets[{i}] has length {len(target)} != expected {train_len}."
)
for key, values in self.static_numerical_covariates.items():
if len(values) != len(self.train_lens):
raise ValueError(
f"static_numerical_covariates has key {key} with number of"
f" examples {len(values)} != expected {len(self.train_lens)}."
)
for key, values in self.static_categorical_covariates.items():
if len(values) != len(self.train_lens):
raise ValueError(
f"static_categorical_covariates has key {key} with number of"
f" examples {len(values)} != expected {len(self.train_lens)}."
)
for lens, dict_cov, dict_cov_name in (
(
self.train_lens,
self.train_dynamic_numerical_covariates,
"train_dynamic_numerical_covariates",
),
(
self.train_lens,
self.train_dynamic_categorical_covariates,
"train_dynamic_categorical_covariates",
),
(
self.test_lens,
self.test_dynamic_numerical_covariates,
"test_dynamic_numerical_covariates",
),
(
self.test_lens,
self.test_dynamic_categorical_covariates,
"test_dynamic_categorical_covariates",
),
):
for key, cov_values in dict_cov.items():
if len(cov_values) != len(lens):
raise ValueError(
f"{dict_cov_name} has key {key} with number of examples"
f" {len(cov_values)} != expected {len(lens)}."
)
for i, cov_value in enumerate(cov_values):
if len(cov_value) != lens[i]:
raise ValueError(
f"{dict_cov_name} has key {key} with its {i}-th example"
f" length {len(cov_value)} != expected {lens[i]}."
)
def create_covariate_matrix(
self,
one_hot_encoder_drop: str | None = "first",
use_intercept: bool = True,
assert_covariates: bool = False,
assert_covariate_shapes: bool = False,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Creates target vector and covariate matrices for in context regression.
Here we use model fitting language to refer to the context as 'train' and
the horizon as 'test'.
Args:
one_hot_encoder_drop: Which drop strategy to use for the one hot encoder.
use_intercept: Whether to prepare an intercept (all 1) column in the
matrices.
assert_covariates: Whether to assert the validity of the covariate inputs.
assert_covariate_shapes: Whether to assert the shapes of the covariate
inputs when `assert_covariates` is True.
Returns:
A tuple of the target vector, the covariate matrix for the context, and
the covariate matrix for the horizon.
"""
if assert_covariates:
self._assert_covariates(assert_covariate_shapes)
x_train, x_test = [], []
# Numerical features.
for name in sorted(self.train_dynamic_numerical_covariates):
x_train.append(
_unnest(self.train_dynamic_numerical_covariates[name])[:, np.newaxis]
)
x_test.append(
_unnest(self.test_dynamic_numerical_covariates[name])[:, np.newaxis]
)
for covs in self.static_numerical_covariates.values():
x_train.append(_repeat(covs, self.train_lens)[:, np.newaxis])
x_test.append(_repeat(covs, self.test_lens)[:, np.newaxis])
if x_train:
x_train = np.concatenate(x_train, axis=1)
x_test = np.concatenate(x_test, axis=1)
# Normalize for robustness.
x_mean = np.mean(x_train, axis=0, keepdims=True)
x_std = np.where(
(w := np.std(x_train, axis=0, keepdims=True)) > _TOL, w, 1.0
)
x_train = [(x_train - x_mean) / x_std]
x_test = [(x_test - x_mean) / x_std]
# Categorical features. Encode one by one.
one_hot_encoder = preprocessing.OneHotEncoder(
drop=one_hot_encoder_drop,
sparse=False,
handle_unknown="ignore",
)
for name in sorted(self.train_dynamic_categorical_covariates.keys()):
ohe_train = _unnest(self.train_dynamic_categorical_covariates[name])[
:, np.newaxis
]
ohe_test = _unnest(self.test_dynamic_categorical_covariates[name])[
:, np.newaxis
]
x_train.append(np.array(one_hot_encoder.fit_transform(ohe_train)))
x_test.append(np.array(one_hot_encoder.transform(ohe_test)))
for covs in self.static_categorical_covariates.values():
ohe = one_hot_encoder.fit_transform(np.array(covs)[:, np.newaxis])
x_train.append(_repeat(ohe, self.train_lens))
x_test.append(_repeat(ohe, self.test_lens))
x_train = np.concatenate(x_train, axis=1)
x_test = np.concatenate(x_test, axis=1)
if use_intercept:
x_train = np.pad(x_train, ((0, 0), (1, 0)), constant_values=1.0)
x_test = np.pad(x_test, ((0, 0), (1, 0)), constant_values=1.0)
return _unnest(self.targets), x_train, x_test
def fit(self) -> Any:
raise NotImplementedError("Fit is not implemented.")
class BatchedInContextXRegLinear(BatchedInContextXRegBase):
"""Linear in-context regression model."""
def fit(
self,
ridge: float = 0.0,
one_hot_encoder_drop: str | None = "first",
use_intercept: bool = True,
force_on_cpu: bool = False,
max_rows_per_col: int = 0,
max_rows_per_col_sample_seed: int = 42,
debug_info: bool = False,
assert_covariates: bool = False,
assert_covariate_shapes: bool = False,
) -> (
list[np.ndarray]
| tuple[
list[np.ndarray], list[np.ndarray], jax.Array, jax.Array, jax.Array
]
):
"""Fits a linear model for in-context regression.
Args:
ridge: A non-negative value for specifying the ridge regression penalty.
If 0 is provided, fallback to ordinary least squares. Note this penalty
is added to the normalized covariate matrix.
one_hot_encoder_drop: Which drop strategy to use for the one hot encoder.
use_intercept: Whether to prepare an intercept (all 1) column in the
matrices.
force_on_cpu: Whether to force execution on cpu for accelerator machines.
max_rows_per_col: How many rows to subsample per column. 0 for no
subsampling. This is for speeding up model fitting.
max_rows_per_col_sample_seed: The seed for the subsampling if needed by
`max_rows_per_col`.
debug_info: Whether to return debug info.
assert_covariates: Whether to assert the validity of the covariate inputs.
assert_covariate_shapes: Whether to assert the shapes of the covariate
inputs when `assert_covariates` is True.
Returns:
If `debug_info` is False:
The linear fits on the horizon.
If `debug_info` is True:
A tuple of:
- the linear fits on the horizon,
- the linear fits on the context,
- the flattened target vector,
- the covariate matrix for the context, and
- the covariate matrix for the horizon.
"""
flat_targets, x_train_raw, x_test = self.create_covariate_matrix(
one_hot_encoder_drop=one_hot_encoder_drop,
use_intercept=use_intercept,
assert_covariates=assert_covariates,
assert_covariate_shapes=assert_covariate_shapes,
)
x_train = x_train_raw.copy()
if max_rows_per_col:
nrows, ncols = x_train.shape
if nrows > (w := ncols * max_rows_per_col):
subsample = jax.random.choice(
jax.random.PRNGKey(max_rows_per_col_sample_seed),
nrows,
(w,),
replace=False,
)
x_train = x_train[subsample]
flat_targets = flat_targets[subsample]
device = jax.devices("cpu")[0] if force_on_cpu else None
# Runs jitted version of the solvers which are quicker at the cost of
# running jitting during the first time calling. Re-jitting happens whenever
# new (padded) shapes are encountered.
# Ocassionally it helps with the speed and the accuracy if we force single
# thread execution on cpu for accelerator machines:
# 1. Avoid moving data to accelarator memory.
# 2. Avoid precision loss if any.
with jax.default_device(device):
x_train_raw = _to_padded_jax_array(x_train_raw)
x_train = _to_padded_jax_array(x_train)
flat_targets = _to_padded_jax_array(flat_targets)
x_test = _to_padded_jax_array(x_test)
beta_hat = (
jnp.linalg.pinv(
x_train.T @ x_train + ridge * jnp.eye(x_train.shape[1]),
hermitian=True,
)
@ x_train.T
@ flat_targets
)
y_hat = x_test @ beta_hat
y_hat_context = x_train_raw @ beta_hat if debug_info else None
outputs = []
outputs_context = []
# Reconstruct the ragged 2-dim batched forecasts from flattened linear fits.
train_index, test_index = 0, 0
for train_index_delta, test_index_delta in zip(
self.train_lens, self.test_lens
):
outputs.append(
np.array(y_hat[test_index : (test_index + test_index_delta)])
)
if debug_info:
outputs_context.append(
np.array(
y_hat_context[train_index : (train_index + train_index_delta)]
)
)
train_index += train_index_delta
test_index += test_index_delta
if debug_info:
return outputs, outputs_context, flat_targets, x_train, x_test
else:
return outputs