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

Merge pull request #89 from google-research/rajat_dev

Browse Source 
standardizing styles across all files.
This commit is contained in:
Yichen Zhou
2024-07-09 10:17:53 -07:00
committed by GitHub
parent 1facd4eaab 954b7ff4ac
commit 28dcfa669b
11 changed files with 734 additions and 782 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+16 -1
View File
@@ -164,4 +164,19 @@ forecast_df = tfm.forecast_on_df(
## Finetuning ## Finetuning
We have provided an example of finetuning the model on a new dataset in `notebooks/finetuning.ipynb`. 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.
experiments/baselines/timegpt_pipeline.py
+183 -181
View File
@@ -31,191 +31,191 @@ from nixtla import NixtlaClient
def get_seasonality(freq: str) -> int: def get_seasonality(freq: str) -> int:
return _get_seasonality(freq, seasonalities={"D": 7}) return _get_seasonality(freq, seasonalities={"D": 7})
def maybe_convert_col_to_datetime(df: pd.DataFrame, col_name: str) -> pd.DataFrame: def maybe_convert_col_to_datetime(df: pd.DataFrame,
if not pd.api.types.is_datetime64_any_dtype(df[col_name]): col_name: str) -> pd.DataFrame:
df = df.copy() if not pd.api.types.is_datetime64_any_dtype(df[col_name]):
df[col_name] = pd.to_datetime(df[col_name]) df = df.copy()
return df df[col_name] = pd.to_datetime(df[col_name])
return df
def zero_pad_time_series(df, freq, min_length=36): def zero_pad_time_series(df, freq, min_length=36):
"""If time_series length is less than min_length, front pad it with zeros.""" """If time_series length is less than min_length, front pad it with zeros."""
# 1. Calculate required padding for each unique_id # 1. Calculate required padding for each unique_id
value_counts = df["unique_id"].value_counts() value_counts = df["unique_id"].value_counts()
to_pad = value_counts[value_counts < min_length].index to_pad = value_counts[value_counts < min_length].index
# 2. Create a new DataFrame to hold padded data # 2. Create a new DataFrame to hold padded data
padded_data = [] padded_data = []
for unique_id in to_pad: for unique_id in to_pad:
# 2a. Filter data for the specific unique_id # 2a. Filter data for the specific unique_id
subset = df[df["unique_id"] == unique_id] subset = df[df["unique_id"] == unique_id]
if len(subset) > min_length: if len(subset) > min_length:
padded_data.append(subset) padded_data.append(subset)
else: else:
# 2b. Determine earliest date and calculate padding dates # 2b. Determine earliest date and calculate padding dates
start_date = subset["ds"].min() start_date = subset["ds"].min()
padding_dates = pd.date_range( padding_dates = pd.date_range(
end=start_date, end=start_date,
periods=min_length - len(subset) + 1, periods=min_length - len(subset) + 1,
freq=freq, # 'MS' for month start freq=freq, # 'MS' for month start
)[ )[:-1] # Exclude the start_date itself
:-1
] # Exclude the start_date itself
# 2c. Create padding data # 2c. Create padding data
padding_df = pd.DataFrame( padding_df = pd.DataFrame({
{"ds": padding_dates, "unique_id": unique_id, "y": 0} # Zero padding "ds": padding_dates,
) "unique_id": unique_id,
"y": 0
} # Zero padding
)
# 2d. Combine original and padding data, and append to the list # 2d. Combine original and padding data, and append to the list
padded_data.append(pd.concat([padding_df, subset]).sort_values("ds")) padded_data.append(pd.concat([padding_df, subset]).sort_values("ds"))
# 3. Combine all padded data and original data (unchanged) # 3. Combine all padded data and original data (unchanged)
result_df = pd.concat(padded_data + [df[~df["unique_id"].isin(to_pad)]]) result_df = pd.concat(padded_data + [df[~df["unique_id"].isin(to_pad)]])
return result_df return result_df
class Forecaster: class Forecaster:
"""Borrowed from """Borrowed from
https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models. https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models.
""" """
def forecast( def forecast(
self, self,
df: pd.DataFrame, df: pd.DataFrame,
h: int, h: int,
freq: str, freq: str,
) -> pd.DataFrame: ) -> pd.DataFrame:
raise NotImplementedError raise NotImplementedError
def cross_validation( def cross_validation(
self, self,
df: pd.DataFrame, df: pd.DataFrame,
h: int, h: int,
freq: str, freq: str,
n_windows: int = 1, n_windows: int = 1,
step_size: int | None = None, step_size: int | None = None,
) -> pd.DataFrame: ) -> pd.DataFrame:
df = maybe_convert_col_to_datetime(df, "ds") df = maybe_convert_col_to_datetime(df, "ds")
# mlforecast cv code # mlforecast cv code
results = [] results = []
sort_idxs = maybe_compute_sort_indices(df, "unique_id", "ds") sort_idxs = maybe_compute_sort_indices(df, "unique_id", "ds")
if sort_idxs is not None: if sort_idxs is not None:
df = take_rows(df, sort_idxs) df = take_rows(df, sort_idxs)
splits = backtest_splits( splits = backtest_splits(
df, df,
n_windows=n_windows, n_windows=n_windows,
h=h, h=h,
id_col="unique_id", id_col="unique_id",
time_col="ds", time_col="ds",
freq=pd.tseries.frequencies.to_offset(freq), freq=pd.tseries.frequencies.to_offset(freq),
step_size=h if step_size is None else step_size, step_size=h if step_size is None else step_size,
) )
for _, (cutoffs, train, valid) in tqdm(enumerate(splits)): for _, (cutoffs, train, valid) in tqdm(enumerate(splits)):
if len(valid.columns) > 3: if len(valid.columns) > 3:
raise NotImplementedError( raise NotImplementedError(
"Cross validation with exogenous variables is not yet supported." "Cross validation with exogenous variables is not yet supported.")
) y_pred = self.forecast(
y_pred = self.forecast( df=train,
df=train, h=h,
h=h, freq=freq,
freq=freq, )
) y_pred = join(y_pred, cutoffs, on="unique_id", how="left")
y_pred = join(y_pred, cutoffs, on="unique_id", how="left") result = join(
result = join( valid[["unique_id", "ds", "y"]],
valid[["unique_id", "ds", "y"]], y_pred,
y_pred, on=["unique_id", "ds"],
on=["unique_id", "ds"], )
) if result.shape[0] < valid.shape[0]:
if result.shape[0] < valid.shape[0]: raise ValueError(
raise ValueError( "Cross validation result produced less results than expected. "
"Cross validation result produced less results than expected. " "Please verify that the frequency parameter (freq) matches your series' "
"Please verify that the frequency parameter (freq) matches your series' " "and that there aren't any missing periods.")
"and that there aren't any missing periods." results.append(result)
) out = vertical_concat(results)
results.append(result) out = drop_index_if_pandas(out)
out = vertical_concat(results) first_out_cols = ["unique_id", "ds", "cutoff", "y"]
out = drop_index_if_pandas(out) remaining_cols = [c for c in out.columns if c not in first_out_cols]
first_out_cols = ["unique_id", "ds", "cutoff", "y"] fcst_cv_df = out[first_out_cols + remaining_cols]
remaining_cols = [c for c in out.columns if c not in first_out_cols] return fcst_cv_df
fcst_cv_df = out[first_out_cols + remaining_cols]
return fcst_cv_df
class TimeGPT(Forecaster): class TimeGPT(Forecaster):
"""Borrowed from """Borrowed from
https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models. https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models.
We modify the class to take care of edge cases. We modify the class to take care of edge cases.
""" """
def __init__( def __init__(
self, self,
api_key: str | None = None, api_key: str | None = None,
base_url: Optional[str] = None, base_url: Optional[str] = None,
max_retries: int = 1, max_retries: int = 1,
model: str = "timegpt-1", model: str = "timegpt-1",
alias: str = "TimeGPT", alias: str = "TimeGPT",
): ):
self.api_key = api_key self.api_key = api_key
self.base_url = base_url self.base_url = base_url
self.max_retries = max_retries self.max_retries = max_retries
self.model = model self.model = model
self.alias = alias self.alias = alias
def _get_client(self) -> NixtlaClient: def _get_client(self) -> NixtlaClient:
if self.api_key is None: if self.api_key is None:
api_key = os.environ["NIXTLA_API_KEY"] api_key = os.environ["NIXTLA_API_KEY"]
else: else:
api_key = self.api_key api_key = self.api_key
return NixtlaClient( return NixtlaClient(
api_key=api_key, api_key=api_key,
base_url=self.base_url, base_url=self.base_url,
max_retries=self.max_retries, 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)
def forecast( fcst_df = pd.concat(all_fcst_df)
self, fcst_df["ds"] = pd.to_datetime(fcst_df["ds"])
df: pd.DataFrame, replace_dict = {}
h: int, for col in fcst_df.columns:
freq: str, if col.startswith("TimeGPT"):
level: List = [90.0], replace_dict[col] = col.replace("TimeGPT", self.alias)
chunk_size: Optional[int] = None, fcst_df = fcst_df.rename(columns=replace_dict)
) -> pd.DataFrame: return fcst_df
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( def run_timegpt(
@@ -227,25 +227,27 @@ def run_timegpt(
dataset: str, dataset: str,
model: str = "timegpt-1", model: str = "timegpt-1",
) -> Tuple[pd.DataFrame, float, str]: ) -> Tuple[pd.DataFrame, float, str]:
os.environ["NIXTLA_ID_AS_COL"] = "true" os.environ["NIXTLA_ID_AS_COL"] = "true"
model = TimeGPT(model="timegpt-1", alias=model) model = TimeGPT(model="timegpt-1", alias=model)
padded_train_df = zero_pad_time_series(train_df, freq) padded_train_df = zero_pad_time_series(train_df, freq)
init_time = time() init_time = time()
# For these datasets the API fails if we do not chunk. # For these datasets the API fails if we do not chunk.
if dataset in ["m5", "m4_quarterly"]: if dataset in ["m5", "m4_quarterly"]:
chunk_size = 5000 chunk_size = 5000
else: else:
chunk_size = None chunk_size = None
fcsts_df = model.forecast( fcsts_df = model.forecast(df=padded_train_df,
df=padded_train_df, h=horizon, level=level, freq=freq, chunk_size=chunk_size h=horizon,
) level=level,
total_time = time() - init_time freq=freq,
# In case levels are not returned we replace the levels with the mean predictions. chunk_size=chunk_size)
# Note that this does not affect the results table as we only compare on point total_time = time() - init_time
# forecastign metrics. # In case levels are not returned we replace the levels with the mean predictions.
for lvl in level: # Note that this does not affect the results table as we only compare on point
if f"{model.alias}-lo-{lvl}" not in fcsts_df.columns: # forecastign metrics.
fcsts_df[f"{model.alias}-lo-{lvl}"] = fcsts_df[model.alias] for lvl in level:
if f"{model.alias}-hi-{lvl}" not in fcsts_df.columns: if f"{model.alias}-lo-{lvl}" not in fcsts_df.columns:
fcsts_df[f"{model.alias}-hi-{lvl}"] = fcsts_df[model.alias] fcsts_df[f"{model.alias}-lo-{lvl}"] = fcsts_df[model.alias]
return fcsts_df, total_time, 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/extended_benchmarks/run_timegpt.py
+35 -38
View File
@@ -11,7 +11,6 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""Evaluation script for timegpt.""" """Evaluation script for timegpt."""
import os import os
@@ -25,7 +24,6 @@ from ..baselines.timegpt_pipeline import run_timegpt
from .utils import ExperimentHandler from .utils import ExperimentHandler
dataset_names = [ dataset_names = [
"m1_monthly", "m1_monthly",
"m1_quarterly", "m1_quarterly",
@@ -63,46 +61,45 @@ _MODEL_NAME = flags.DEFINE_string(
) )
_SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory") _SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory")
QUANTILES = list(np.arange(1, 10) / 10.0) QUANTILES = list(np.arange(1, 10) / 10.0)
def main(): def main():
results_list = [] results_list = []
run_id = np.random.randint(100000) run_id = np.random.randint(100000)
model_name = _MODEL_NAME.value model_name = _MODEL_NAME.value
for dataset in dataset_names: for dataset in dataset_names:
print(f"Evaluating model {model_name} on dataset {dataset}", flush=True) print(f"Evaluating model {model_name} on dataset {dataset}", flush=True)
exp = ExperimentHandler(dataset, quantiles=QUANTILES) exp = ExperimentHandler(dataset, quantiles=QUANTILES)
train_df = exp.train_df train_df = exp.train_df
horizon = exp.horizon horizon = exp.horizon
seasonality = exp.seasonality seasonality = exp.seasonality
freq = exp.freq freq = exp.freq
level = exp.level level = exp.level
fcsts_df, total_time, model_name = run_timegpt( fcsts_df, total_time, model_name = run_timegpt(
train_df=train_df, train_df=train_df,
horizon=exp.horizon, horizon=exp.horizon,
model=model_name, model=model_name,
seasonality=seasonality, seasonality=seasonality,
freq=freq, freq=freq,
dataset=dataset, dataset=dataset,
level=level, level=level,
) )
time_df = pd.DataFrame({"time": [total_time], "model": model_name}) time_df = pd.DataFrame({"time": [total_time], "model": model_name})
fcsts_df = exp.fcst_from_level_to_quantiles(fcsts_df, model_name) fcsts_df = exp.fcst_from_level_to_quantiles(fcsts_df, model_name)
results = exp.evaluate_from_predictions( results = exp.evaluate_from_predictions(models=[model_name],
models=[model_name], fcsts_df=fcsts_df, times_df=time_df fcsts_df=fcsts_df,
) times_df=time_df)
print(results, flush=True) print(results, flush=True)
results_list.append(results) results_list.append(results)
results_full = pd.concat(results_list) results_full = pd.concat(results_list)
save_path = os.path.join(_SAVE_DIR.value, str(run_id)) save_path = os.path.join(_SAVE_DIR.value, str(run_id))
print(f"Saving results to {save_path}", flush=True) print(f"Saving results to {save_path}", flush=True)
os.makedirs(save_path, exist_ok=True) os.makedirs(save_path, exist_ok=True)
results_full.to_csv(f"{save_path}/results.csv") results_full.to_csv(f"{save_path}/results.csv")
if __name__ == "__main__": if __name__ == "__main__":
FLAGS = flags.FLAGS FLAGS = flags.FLAGS
FLAGS(sys.argv) FLAGS(sys.argv)
main() main()
experiments/extended_benchmarks/run_timesfm.py
+5 -9
View File
@@ -11,7 +11,6 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""Evaluation script for timesfm.""" """Evaluation script for timesfm."""
import os import os
@@ -26,7 +25,6 @@ import timesfm
from .utils import ExperimentHandler from .utils import ExperimentHandler
dataset_names = [ dataset_names = [
"m1_monthly", "m1_monthly",
"m1_quarterly", "m1_quarterly",
@@ -74,16 +72,14 @@ context_dict = {
"m4_yearly": 64, "m4_yearly": 64,
} }
_MODEL_PATH = flags.DEFINE_string( _MODEL_PATH = flags.DEFINE_string("model_path", "/home/timesfm_q10_20240501",
"model_path", "/home/timesfm_q10_20240501", "Path to model" "Path to model")
)
_BATCH_SIZE = flags.DEFINE_integer("batch_size", 64, "Batch size") _BATCH_SIZE = flags.DEFINE_integer("batch_size", 64, "Batch size")
_HORIZON = flags.DEFINE_integer("horizon", 128, "Horizon") _HORIZON = flags.DEFINE_integer("horizon", 128, "Horizon")
_BACKEND = flags.DEFINE_string("backend", "gpu", "Backend") _BACKEND = flags.DEFINE_string("backend", "gpu", "Backend")
_NUM_JOBS = flags.DEFINE_integer("num_jobs", 1, "Number of jobs") _NUM_JOBS = flags.DEFINE_integer("num_jobs", 1, "Number of jobs")
_SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory") _SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory")
QUANTILES = list(np.arange(1, 10) / 10.0) QUANTILES = list(np.arange(1, 10) / 10.0)
@@ -127,9 +123,9 @@ def main():
) )
total_time = time.time() - init_time total_time = time.time() - init_time
time_df = pd.DataFrame({"time": [total_time], "model": model_name}) time_df = pd.DataFrame({"time": [total_time], "model": model_name})
results = exp.evaluate_from_predictions( results = exp.evaluate_from_predictions(models=[model_name],
models=[model_name], fcsts_df=fcsts_df, times_df=time_df fcsts_df=fcsts_df,
) times_df=time_df)
print(results, flush=True) print(results, flush=True)
results_list.append(results) results_list.append(results)
results_full = pd.concat(results_list) results_full = pd.concat(results_list)
experiments/extended_benchmarks/utils.py
+24 -36
View File
@@ -11,7 +11,6 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""Forked from https://github.com/Nixtla/nixtla/blob/main/experiments/amazon-chronos/src/utils.py.""" """Forked from https://github.com/Nixtla/nixtla/blob/main/experiments/amazon-chronos/src/utils.py."""
from functools import partial from functools import partial
@@ -46,11 +45,9 @@ def quantile_loss(
target_col: str = "y", target_col: str = "y",
) -> pd.DataFrame: ) -> pd.DataFrame:
delta_y = df[models].sub(df[target_col], axis=0) delta_y = df[models].sub(df[target_col], axis=0)
res = ( res = (np.maximum(q * delta_y,
np.maximum(q * delta_y, (q - 1) * delta_y) (q - 1) * delta_y).groupby(df[id_col],
.groupby(df[id_col], observed=True) observed=True).mean())
.mean()
)
res.index.name = id_col res.index.name = id_col
res = res.reset_index() res = res.reset_index()
return res return res
@@ -66,10 +63,8 @@ class ExperimentHandler:
models_dir: str = "./models", models_dir: str = "./models",
): ):
if dataset not in gluonts_datasets: if dataset not in gluonts_datasets:
raise Exception( raise Exception(f"dataset {dataset} not found in gluonts "
f"dataset {dataset} not found in gluonts " f"available datasets: {', '.join(gluonts_datasets)}")
f"available datasets: {', '.join(gluonts_datasets)}"
)
self.dataset = dataset self.dataset = dataset
self.quantiles = quantiles self.quantiles = quantiles
self.level = self._transform_quantiles_to_levels(quantiles) self.level = self._transform_quantiles_to_levels(quantiles)
@@ -80,10 +75,8 @@ class ExperimentHandler:
gluonts_dataset = get_dataset(self.dataset) gluonts_dataset = get_dataset(self.dataset)
self.horizon = gluonts_dataset.metadata.prediction_length self.horizon = gluonts_dataset.metadata.prediction_length
if self.horizon is None: if self.horizon is None:
raise Exception( raise Exception(f"horizon not found for dataset {self.dataset} "
f"horizon not found for dataset {self.dataset} " "experiment cannot be run")
"experiment cannot be run"
)
self.freq = gluonts_dataset.metadata.freq self.freq = gluonts_dataset.metadata.freq
# get_seasonality() returns 1 for freq='D', override this to 7. This significantly improves the accuracy of # get_seasonality() returns 1 for freq='D', override this to 7. This significantly improves the accuracy of
# statistical models on datasets like m5/nn5_daily. The models like AutoARIMA/AutoETS can still set # statistical models on datasets like m5/nn5_daily. The models like AutoARIMA/AutoETS can still set
@@ -122,9 +115,8 @@ class ExperimentHandler:
@staticmethod @staticmethod
def _transform_quantiles_to_levels(quantiles: List[float]) -> List[int]: def _transform_quantiles_to_levels(quantiles: List[float]) -> List[int]:
level = [ level = [int(100 - 200 * q) for q in quantiles if q < 0.5
int(100 - 200 * q) for q in quantiles if q < 0.5 ] # in this case mean=mediain
] # in this case mean=mediain
level = sorted(list(set(level))) level = sorted(list(set(level)))
return level return level
@@ -153,9 +145,8 @@ class ExperimentHandler:
last_n: int | None = None, last_n: int | None = None,
) -> pd.DataFrame: ) -> pd.DataFrame:
with multiprocessing.Pool(os.cpu_count()) as pool: # Create a process pool with multiprocessing.Pool(os.cpu_count()) as pool: # Create a process pool
results = pool.map( results = pool.map(parallel_transform, zip(gluonts_dataset,
parallel_transform, zip(gluonts_dataset, repeat(last_n)) repeat(last_n)))
)
df = pd.concat(results) df = pd.concat(results)
df = df.reset_index(drop=True) df = df.reset_index(drop=True)
return df return df
@@ -177,9 +168,8 @@ class ExperimentHandler:
def save_dataframe(self, df: pd.DataFrame, file_name: str): def save_dataframe(self, df: pd.DataFrame, file_name: str):
df.to_csv(f"{self.results_dir}/{file_name}", index=False) df.to_csv(f"{self.results_dir}/{file_name}", index=False)
def save_results( def save_results(self, fcst_df: pd.DataFrame, total_time: float,
self, fcst_df: pd.DataFrame, total_time: float, model_name: str model_name: str):
):
self.save_dataframe( self.save_dataframe(
fcst_df, fcst_df,
f"{model_name}-{self.dataset}-fcst.csv", f"{model_name}-{self.dataset}-fcst.csv",
@@ -215,23 +205,21 @@ class ExperimentHandler:
times_df = [] times_df = []
for model in models: for model in models:
fcst_method_df = pd.read_csv( fcst_method_df = pd.read_csv(
f"{self.results_dir}/{model}-{self.dataset}-fcst.csv" f"{self.results_dir}/{model}-{self.dataset}-fcst.csv").set_index(
).set_index(["unique_id", "ds"]) ["unique_id", "ds"])
fcsts_df.append(fcst_method_df) fcsts_df.append(fcst_method_df)
time_method_df = pd.read_csv( time_method_df = pd.read_csv(
f"{self.results_dir}/{model}-{self.dataset}-time.csv" f"{self.results_dir}/{model}-{self.dataset}-time.csv")
)
times_df.append(time_method_df) times_df.append(time_method_df)
fcsts_df = pd.concat(fcsts_df, axis=1).reset_index() fcsts_df = pd.concat(fcsts_df, axis=1).reset_index()
fcsts_df["ds"] = pd.to_datetime(fcsts_df["ds"]) fcsts_df["ds"] = pd.to_datetime(fcsts_df["ds"])
times_df = pd.concat(times_df) times_df = pd.concat(times_df)
return self.evaluate_from_predictions( return self.evaluate_from_predictions(models=models,
models=models, fcsts_df=fcsts_df, times_df=times_df fcsts_df=fcsts_df,
) times_df=times_df)
def evaluate_from_predictions( def evaluate_from_predictions(self, models: List[str], fcsts_df: pd.DataFrame,
self, models: List[str], fcsts_df: pd.DataFrame, times_df: pd.DataFrame times_df: pd.DataFrame) -> pd.DataFrame:
) -> pd.DataFrame:
test_df = self.test_df test_df = self.test_df
train_df = self.train_df train_df = self.train_df
test_df = test_df.merge(fcsts_df, how="left") test_df = test_df.merge(fcsts_df, how="left")
@@ -262,9 +250,9 @@ class ExperimentHandler:
eval_prob_df["metric"] = "scaled_crps" eval_prob_df["metric"] = "scaled_crps"
eval_df = pd.concat([eval_df, eval_prob_df]).reset_index(drop=True) eval_df = pd.concat([eval_df, eval_prob_df]).reset_index(drop=True)
eval_df = eval_df.groupby("metric").mean(numeric_only=True).reset_index() eval_df = eval_df.groupby("metric").mean(numeric_only=True).reset_index()
eval_df = eval_df.melt( eval_df = eval_df.melt(id_vars="metric",
id_vars="metric", value_name="value", var_name="model" value_name="value",
) var_name="model")
times_df.insert(0, "metric", "time") times_df.insert(0, "metric", "time")
times_df = times_df.rename(columns={"time": "value"}) times_df = times_df.rename(columns={"time": "value"})
eval_df = pd.concat([eval_df, times_df]) eval_df = pd.concat([eval_df, times_df])
experiments/long_horizon_benchmarks/run_eval.py
+21 -34
View File
@@ -11,7 +11,6 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""Eval pipeline.""" """Eval pipeline."""
import json import json
@@ -28,38 +27,28 @@ import torch
import tqdm import tqdm
from timesfm import data_loader from timesfm import data_loader
FLAGS = flags.FLAGS FLAGS = flags.FLAGS
_BATCH_SIZE = flags.DEFINE_integer( _BATCH_SIZE = flags.DEFINE_integer("batch_size", 64,
"batch_size", 64, "Batch size for the randomly sampled batch" "Batch size for the randomly sampled batch")
)
_DATASET = flags.DEFINE_string("dataset", "etth1", "The name of the dataset.") _DATASET = flags.DEFINE_string("dataset", "etth1", "The name of the dataset.")
_MODEL_PATH = flags.DEFINE_string( _MODEL_PATH = flags.DEFINE_string("model_path", "./timesfm_q10_20240501",
"model_path", "./timesfm_q10_20240501", "The name of the dataset." "The name of the dataset.")
) _DATETIME_COL = flags.DEFINE_string("datetime_col", "date",
_DATETIME_COL = flags.DEFINE_string( "Column having datetime.")
"datetime_col", "date", "Column having datetime." _NUM_COV_COLS = flags.DEFINE_list("num_cov_cols", None,
) "Column having numerical features.")
_NUM_COV_COLS = flags.DEFINE_list( _CAT_COV_COLS = flags.DEFINE_list("cat_cov_cols", None,
"num_cov_cols", None, "Column having numerical features." "Column having categorical features.")
)
_CAT_COV_COLS = flags.DEFINE_list(
"cat_cov_cols", None, "Column having categorical features."
)
_TS_COLS = flags.DEFINE_list("ts_cols", None, "Columns of time-series features") _TS_COLS = flags.DEFINE_list("ts_cols", None, "Columns of time-series features")
_NORMALIZE = flags.DEFINE_bool( _NORMALIZE = flags.DEFINE_bool("normalize", True,
"normalize", True, "normalize data for eval or not" "normalize data for eval or not")
) _CONTEXT_LEN = flags.DEFINE_integer("context_len", 512,
_CONTEXT_LEN = flags.DEFINE_integer( "Length of the context window")
"context_len", 512, "Length of the context window"
)
_PRED_LEN = flags.DEFINE_integer("pred_len", 96, "prediction length.") _PRED_LEN = flags.DEFINE_integer("pred_len", 96, "prediction length.")
_BACKEND = flags.DEFINE_string("backend", "gpu", "backend to use") _BACKEND = flags.DEFINE_string("backend", "gpu", "backend to use")
_RESULTS_DIR = flags.DEFINE_string( _RESULTS_DIR = flags.DEFINE_string("results_dir", "./results/long_horizon",
"results_dir", "./results/long_horizon", "results directory" "results directory")
)
DATA_DICT = { DATA_DICT = {
"ettm2": { "ettm2": {
@@ -176,9 +165,8 @@ def eval():
holiday=False, holiday=False,
permute=False, permute=False,
) )
eval_itr = dtl.tf_dataset( eval_itr = dtl.tf_dataset(mode="test",
mode="test", shift=_PRED_LEN.value shift=_PRED_LEN.value).as_numpy_iterator()
).as_numpy_iterator()
model_path = _MODEL_PATH.value model_path = _MODEL_PATH.value
if model_path.startswith("amazon"): if model_path.startswith("amazon"):
model = chronos.ChronosPipeline.from_pretrained( model = chronos.ChronosPipeline.from_pretrained(
@@ -213,10 +201,9 @@ def eval():
for batch in tqdm.tqdm(eval_itr): for batch in tqdm.tqdm(eval_itr):
past = batch[0] past = batch[0]
actuals = batch[3] actuals = batch[3]
forecasts = get_forecasts( forecasts = get_forecasts(model_path, model, past, int_freq,
model_path, model, past, int_freq, _PRED_LEN.value _PRED_LEN.value)
) forecasts = forecasts[:, 0:actuals.shape[1]]
forecasts = forecasts[:, 0 : actuals.shape[1]]
mae_run_losses.append(_mae(forecasts, actuals).sum()) mae_run_losses.append(_mae(forecasts, actuals).sum())
mse_run_losses.append(_mse(forecasts, actuals).sum()) mse_run_losses.append(_mse(forecasts, actuals).sum())
smape_run_losses.append(_smape(forecasts, actuals).sum()) smape_run_losses.append(_smape(forecasts, actuals).sum())
src/timesfm/__init__.py
-1
View File
@@ -11,7 +11,6 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""TimesFM init file.""" """TimesFM init file."""
from .timesfm import TimesFm, freq_map from .timesfm import TimesFm, freq_map
src/timesfm/data_loader.py
+8 -14
View File
@@ -11,13 +11,11 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""TF dataloaders for general timeseries datasets. """TF dataloaders for general timeseries datasets.
The expected input format is csv file with a datetime index. The expected input format is csv file with a datetime index.
""" """
from absl import logging from absl import logging
import numpy as np import numpy as np
import pandas as pd import pandas as pd
@@ -79,9 +77,8 @@ class TimeSeriesdata(object):
self.data_df['ccol'] = np.zeros(self.data_df.shape[0]) self.data_df['ccol'] = np.zeros(self.data_df.shape[0])
cat_cov_cols = ['ccol'] cat_cov_cols = ['ccol']
self.data_df.fillna(0, inplace=True) self.data_df.fillna(0, inplace=True)
self.data_df.set_index( self.data_df.set_index(pd.DatetimeIndex(self.data_df[datetime_col]),
pd.DatetimeIndex(self.data_df[datetime_col]), inplace=True inplace=True)
)
self.num_cov_cols = num_cov_cols self.num_cov_cols = num_cov_cols
self.cat_cov_cols = cat_cov_cols self.cat_cov_cols = cat_cov_cols
self.ts_cols = ts_cols self.ts_cols = ts_cols
@@ -94,18 +91,16 @@ class TimeSeriesdata(object):
data_df_idx[-1] + pd.Timedelta(1, freq=freq), data_df_idx[-1] + pd.Timedelta(1, freq=freq),
periods=pred_len + 1, periods=pred_len + 1,
freq=freq, freq=freq,
) ))
)
self.time_df = time_features.TimeCovariates( self.time_df = time_features.TimeCovariates(
date_index, holiday=holiday date_index, holiday=holiday).get_covariates()
).get_covariates()
self.hist_len = hist_len self.hist_len = hist_len
self.pred_len = pred_len self.pred_len = pred_len
self.batch_size = batch_size self.batch_size = batch_size
self.freq = freq self.freq = freq
self.normalize = normalize self.normalize = normalize
self.data_mat = self.data_df[self.ts_cols].to_numpy().transpose() self.data_mat = self.data_df[self.ts_cols].to_numpy().transpose()
self.data_mat = self.data_mat[:, 0 : self.test_range[1]] self.data_mat = self.data_mat[:, 0:self.test_range[1]]
self.time_mat = self.time_df.to_numpy().transpose() self.time_mat = self.time_df.to_numpy().transpose()
self.num_feat_mat = self.data_df[num_cov_cols].to_numpy().transpose() self.num_feat_mat = self.data_df[num_cov_cols].to_numpy().transpose()
self.cat_feat_mat, self.cat_sizes = self._get_cat_cols(cat_cov_cols) self.cat_feat_mat, self.cat_sizes = self._get_cat_cols(cat_cov_cols)
@@ -135,7 +130,7 @@ class TimeSeriesdata(object):
def _normalize_data(self): def _normalize_data(self):
self.scaler = StandardScaler() self.scaler = StandardScaler()
train_mat = self.data_mat[:, self.train_range[0] : self.train_range[1]] train_mat = self.data_mat[:, self.train_range[0]:self.train_range[1]]
self.scaler = self.scaler.fit(train_mat.transpose()) self.scaler = self.scaler.fit(train_mat.transpose())
self.data_mat = self.scaler.transform(self.data_mat.transpose()).transpose() self.data_mat = self.scaler.transform(self.data_mat.transpose()).transpose()
@@ -253,9 +248,8 @@ class TimeSeriesdata(object):
gen_fn = self.train_gen gen_fn = self.train_gen
else: else:
gen_fn = lambda: self.test_val_gen(mode, shift) gen_fn = lambda: self.test_val_gen(mode, shift)
output_types = tuple( output_types = tuple([tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 +
[tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 + [tf.int32] * 2 [tf.int32] * 2)
)
dataset = tf.data.Dataset.from_generator(gen_fn, output_types) dataset = tf.data.Dataset.from_generator(gen_fn, output_types)
dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE) dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
return dataset return dataset
src/timesfm/patched_decoder.py
+318 -323
View File
@@ -11,7 +11,6 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""Pax ML model for patched time-series decoder. """Pax ML model for patched time-series decoder.
The file implements Residual MLPs, Patched Decoder layers and PAX ML models. 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 stochastics
from praxis.layers import transformers from praxis.layers import transformers
# PAX shortcuts # PAX shortcuts
NestedMap = py_utils.NestedMap NestedMap = py_utils.NestedMap
JTensor = pytypes.JTensor JTensor = pytypes.JTensor
@@ -44,7 +42,6 @@ JTensor = pytypes.JTensor
LayerTpl = pax_fiddle.Config[base_layer.BaseLayer] LayerTpl = pax_fiddle.Config[base_layer.BaseLayer]
template_field = base_layer.template_field template_field = base_layer.template_field
PAD_VAL = 1123581321.0 PAD_VAL = 1123581321.0
DEFAULT_QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] DEFAULT_QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
@@ -57,36 +54,35 @@ _FREQ = "freq"
_OUTPUT_TOKENS = "output_tokens" _OUTPUT_TOKENS = "output_tokens"
_STATS = "stats" _STATS = "stats"
# Small numerical value. # Small numerical value.
_TOLERANCE = 1e-7 _TOLERANCE = 1e-7
def _shift_padded_seq(mask: JTensor, seq: JTensor) -> JTensor: def _shift_padded_seq(mask: JTensor, seq: JTensor) -> JTensor:
"""Shifts rows of seq based on the first 0 in each row of the mask.""" """Shifts rows of seq based on the first 0 in each row of the mask."""
num = seq.shape[1] num = seq.shape[1]
# Find the index of the first 0 in each row of the mask # Find the index of the first 0 in each row of the mask
first_zero_idx = jnp.argmin(mask, axis=1) first_zero_idx = jnp.argmin(mask, axis=1)
# Create a range array for indexing # Create a range array for indexing
idx_range = jnp.arange(num) idx_range = jnp.arange(num)
def shift_row(carry, x): def shift_row(carry, x):
seq_row, shift = x seq_row, shift = x
shifted_idx = (idx_range - shift) % num shifted_idx = (idx_range - shift) % num
shifted_row = seq_row[shifted_idx] shifted_row = seq_row[shifted_idx]
return carry, shifted_row return carry, shifted_row
# Use lax.scan to shift each row of seq based on the corresponding # Use lax.scan to shift each row of seq based on the corresponding
# first_zero_idx. # first_zero_idx.
_, shifted_seq = lax.scan(shift_row, None, (seq, first_zero_idx)) _, shifted_seq = lax.scan(shift_row, None, (seq, first_zero_idx))
return shifted_seq return shifted_seq
class ResidualBlock(base_layer.BaseLayer): class ResidualBlock(base_layer.BaseLayer):
"""Simple feedforward block with residual connection. """Simple feedforward block with residual connection.
Attributes: Attributes:
input_dims: input dimension. input_dims: input dimension.
@@ -99,67 +95,68 @@ class ResidualBlock(base_layer.BaseLayer):
act_tpl: config for activation in hidden layer. act_tpl: config for activation in hidden layer.
""" """
input_dims: int = 0 input_dims: int = 0
hidden_dims: int = 0 hidden_dims: int = 0
output_dims: int = 0 output_dims: int = 0
dropout_prob: float = 0.0 dropout_prob: float = 0.0
layer_norm: bool = False layer_norm: bool = False
dropout_tpl: LayerTpl = template_field(stochastics.Dropout) dropout_tpl: LayerTpl = template_field(stochastics.Dropout)
ln_tpl: LayerTpl = template_field(normalizations.LayerNorm) ln_tpl: LayerTpl = template_field(normalizations.LayerNorm)
act_tpl: LayerTpl = template_field(activations.Swish) act_tpl: LayerTpl = template_field(activations.Swish)
def setup(self): def setup(self):
lnorm_tpl = self.ln_tpl.clone() lnorm_tpl = self.ln_tpl.clone()
lnorm_tpl.dim = self.output_dims lnorm_tpl.dim = self.output_dims
self.create_child("ln_layer", lnorm_tpl) self.create_child("ln_layer", lnorm_tpl)
dropout_tpl = self.dropout_tpl.clone() dropout_tpl = self.dropout_tpl.clone()
dropout_tpl.keep_prob = 1.0 - self.dropout_prob dropout_tpl.keep_prob = 1.0 - self.dropout_prob
self.create_child("dropout", dropout_tpl) self.create_child("dropout", dropout_tpl)
self.create_child( self.create_child(
"hidden_layer", "hidden_layer",
pax_fiddle.Config( pax_fiddle.Config(
linears.FeedForward, linears.FeedForward,
input_dims=self.input_dims, input_dims=self.input_dims,
output_dims=self.hidden_dims, output_dims=self.hidden_dims,
activation_tpl=self.act_tpl.clone(), activation_tpl=self.act_tpl.clone(),
), ),
) )
self.create_child( self.create_child(
"output_layer", "output_layer",
pax_fiddle.Config( pax_fiddle.Config(
linears.FeedForward, linears.FeedForward,
input_dims=self.hidden_dims, input_dims=self.hidden_dims,
output_dims=self.output_dims, output_dims=self.output_dims,
activation_tpl=pax_fiddle.Config(activations.Identity), activation_tpl=pax_fiddle.Config(activations.Identity),
), ),
) )
self.create_child( self.create_child(
"residual_layer", "residual_layer",
pax_fiddle.Config( pax_fiddle.Config(
linears.FeedForward, linears.FeedForward,
input_dims=self.input_dims, input_dims=self.input_dims,
output_dims=self.output_dims, output_dims=self.output_dims,
activation_tpl=pax_fiddle.Config(activations.Identity), activation_tpl=pax_fiddle.Config(activations.Identity),
), ),
) )
def __call__(self, inputs: JTensor) -> JTensor: def __call__(self, inputs: JTensor) -> JTensor:
hidden = self.hidden_layer(inputs) hidden = self.hidden_layer(inputs)
output = self.output_layer(hidden) output = self.output_layer(hidden)
output = self.dropout(output) output = self.dropout(output)
residual = self.residual_layer(inputs) residual = self.residual_layer(inputs)
if self.layer_norm: if self.layer_norm:
return self.ln_layer(output + residual) return self.ln_layer(output + residual)
else: else:
return output + residual return output + residual
def _masked_mean_std(inputs: JTensor, padding: JTensor) -> Tuple[JTensor, JTensor]: def _masked_mean_std(inputs: JTensor,
"""Calculates mean and standard deviation of arr across axis 1. padding: JTensor) -> Tuple[JTensor, JTensor]:
"""Calculates mean and standard deviation of arr across axis 1.
It should exclude values where pad is 1. It should exclude values where pad is 1.
@@ -171,48 +168,48 @@ def _masked_mean_std(inputs: JTensor, padding: JTensor) -> Tuple[JTensor, JTenso
A tuple containing the mean and standard deviation of arr. We return the A tuple containing the mean and standard deviation of arr. We return the
statistics of the first patch with more than three non-padded values. statistics of the first patch with more than three non-padded values.
""" """
# Selecting the first pad with more than 3 unpadded values. # Selecting the first pad with more than 3 unpadded values.
pad_sum = jnp.sum(1 - padding, axis=2) pad_sum = jnp.sum(1 - padding, axis=2)
def _get_patch_index(arr: JTensor): def _get_patch_index(arr: JTensor):
indices = jnp.argmax(arr >= 3, axis=1) indices = jnp.argmax(arr >= 3, axis=1)
row_sum = (arr >= 3).sum(axis=1) row_sum = (arr >= 3).sum(axis=1)
return jnp.where(row_sum == 0, arr.shape[1] - 1, indices) return jnp.where(row_sum == 0, arr.shape[1] - 1, indices)
patch_indices = _get_patch_index(pad_sum) patch_indices = _get_patch_index(pad_sum)
bidxs = jnp.arange(inputs.shape[0]) bidxs = jnp.arange(inputs.shape[0])
arr = inputs[bidxs, patch_indices, :] arr = inputs[bidxs, patch_indices, :]
pad = padding[bidxs, patch_indices, :] pad = padding[bidxs, patch_indices, :]
# Create a mask where P is 0 # Create a mask where P is 0
mask = 1 - pad mask = 1 - pad
# Calculate the number of valid elements # Calculate the number of valid elements
num_valid_elements = jnp.sum(mask, axis=1) num_valid_elements = jnp.sum(mask, axis=1)
num_valid_elements = jnp.where(num_valid_elements == 0, 1, num_valid_elements) num_valid_elements = jnp.where(num_valid_elements == 0, 1, num_valid_elements)
# Calculate the masked sum and squared sum of M # Calculate the masked sum and squared sum of M
masked_sum = jnp.sum(arr * mask, axis=1) 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 # Calculate the masked mean and standard deviation
masked_mean = masked_sum / num_valid_elements masked_mean = masked_sum / num_valid_elements
masked_var = masked_squared_sum / num_valid_elements - masked_mean**2 masked_var = masked_squared_sum / num_valid_elements - masked_mean**2
masked_var = jnp.where(masked_var < 0.0, 0.0, masked_var) masked_var = jnp.where(masked_var < 0.0, 0.0, masked_var)
masked_std = jnp.sqrt(masked_var) masked_std = jnp.sqrt(masked_var)
return masked_mean, masked_std return masked_mean, masked_std
def _create_quantiles() -> list[float]: def _create_quantiles() -> list[float]:
"""Returns the quantiles for forecasting.""" """Returns the quantiles for forecasting."""
return DEFAULT_QUANTILES return DEFAULT_QUANTILES
class PatchedTimeSeriesDecoder(base_layer.BaseLayer): class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
"""Patch decoder layer for time-series foundation model. """Patch decoder layer for time-series foundation model.
Attributes: Attributes:
patch_len: length of input patches. patch_len: length of input patches.
@@ -231,137 +228,137 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
number of output logits. D is model dimension. number of output logits. D is model dimension.
""" """
patch_len: int = 0 patch_len: int = 0
horizon_len: int = 0 horizon_len: int = 0
model_dims: int = 0 model_dims: int = 0
hidden_dims: int = 0 hidden_dims: int = 0
quantiles: list[float] = dataclasses.field(default_factory=_create_quantiles) quantiles: list[float] = dataclasses.field(default_factory=_create_quantiles)
residual_block_tpl: LayerTpl = template_field(ResidualBlock) residual_block_tpl: LayerTpl = template_field(ResidualBlock)
stacked_transformer_params_tpl: LayerTpl = template_field( stacked_transformer_params_tpl: LayerTpl = template_field(
transformers.StackedTransformer transformers.StackedTransformer)
use_freq: bool = True
def setup(self) -> None:
"""Construct the model."""
num_outputs = len(self.quantiles) + 1
stl = self.stacked_transformer_params_tpl.clone()
stl.model_dims = self.model_dims
stl.hidden_dims = self.hidden_dims
stl.mask_self_attention = True
self.create_child("stacked_transformer_layer", stl)
input_resl = self.residual_block_tpl.clone()
ff_in_dims = 2 * self.patch_len
input_resl.input_dims = ff_in_dims
input_resl.hidden_dims = self.hidden_dims
input_resl.output_dims = self.model_dims
self.create_child(
"input_ff_layer",
input_resl,
) )
use_freq: bool = True
def setup(self) -> None: horizon_resl = self.residual_block_tpl.clone()
"""Construct the model.""" horizon_resl.input_dims = self.model_dims
num_outputs = len(self.quantiles) + 1 horizon_resl.hidden_dims = self.hidden_dims
horizon_resl.output_dims = self.horizon_len * num_outputs
self.create_child(
"horizon_ff_layer",
horizon_resl,
)
stl = self.stacked_transformer_params_tpl.clone() self.create_child(
stl.model_dims = self.model_dims "position_emb",
stl.hidden_dims = self.hidden_dims pax_fiddle.Config(layers.PositionalEmbedding,
stl.mask_self_attention = True embedding_dims=self.model_dims),
)
self.create_child("stacked_transformer_layer", stl) if self.use_freq:
self.create_child(
"freq_emb",
pax_fiddle.Config(
embedding_softmax.Embedding,
num_classes=3,
input_dims=self.model_dims,
),
)
input_resl = self.residual_block_tpl.clone() def transform_decode_state(
ff_in_dims = 2 * self.patch_len self, transform_fn: base_layer.DecodeStateTransformFn) -> None:
input_resl.input_dims = ff_in_dims """Transforms all decode state variables based on transform_fn."""
input_resl.hidden_dims = self.hidden_dims self.stacked_transformer_layer.transform_decode_state(transform_fn)
input_resl.output_dims = self.model_dims
self.create_child(
"input_ff_layer",
input_resl,
)
horizon_resl = self.residual_block_tpl.clone() def _forward_transform(
horizon_resl.input_dims = self.model_dims self, inputs: JTensor,
horizon_resl.hidden_dims = self.hidden_dims patched_pads: JTensor) -> Tuple[JTensor, Tuple[JTensor, JTensor]]:
horizon_resl.output_dims = self.horizon_len * num_outputs """Input is of shape [B, N, P]."""
self.create_child( mu, sigma = _masked_mean_std(inputs, patched_pads)
"horizon_ff_layer", sigma = jnp.where(sigma < _TOLERANCE, 1.0, sigma)
horizon_resl, # Normalize each patch.
) outputs = (inputs - mu[:, None, None]) / sigma[:, None, None]
outputs = jnp.where(
jnp.abs(inputs - PAD_VAL) < _TOLERANCE, PAD_VAL, outputs)
return outputs, (mu, sigma)
self.create_child( def _reverse_transform(self, outputs: JTensor,
"position_emb", stats: Tuple[JTensor, JTensor]) -> JTensor:
pax_fiddle.Config( """Output is of shape [B, N, P, Q]."""
layers.PositionalEmbedding, embedding_dims=self.model_dims mu, sigma = stats
), return outputs * sigma[:, None, None, None] + mu[:, None, None, None]
)
if self.use_freq: def _preprocess_input(
self.create_child( self,
"freq_emb", input_ts: JTensor,
pax_fiddle.Config( input_padding: JTensor,
embedding_softmax.Embedding, pos_emb: Optional[JTensor] = None,
num_classes=3, ) -> Tuple[JTensor, JTensor, Optional[Tuple[JTensor, JTensor]], JTensor]:
input_dims=self.model_dims, """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, 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)
model_input = self.input_ff_layer(concat_inputs)
# A patch should not be padded even if there is at least one zero.
patched_padding = jnp.min(patched_pads, axis=-1)
def transform_decode_state( if pos_emb is None:
self, transform_fn: base_layer.DecodeStateTransformFn position_emb = self.position_emb(seq_length=model_input.shape[1])
) -> None: else:
"""Transforms all decode state variables based on transform_fn.""" position_emb = pos_emb
self.stacked_transformer_layer.transform_decode_state(transform_fn) if self.do_eval:
if position_emb.shape[0] != model_input.shape[0]:
position_emb = jnp.repeat(position_emb, model_input.shape[0], axis=0)
position_emb = _shift_padded_seq(patched_padding, position_emb)
model_input += position_emb
def _forward_transform( return model_input, patched_padding, stats, patched_inputs
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)
return outputs, (mu, sigma)
def _reverse_transform( def _postprocess_output(
self, outputs: JTensor, stats: Tuple[JTensor, JTensor] self,
) -> JTensor: model_output: JTensor,
"""Output is of shape [B, N, P, Q].""" num_outputs: int,
mu, sigma = stats stats: Tuple[JTensor, JTensor],
return outputs * sigma[:, None, None, None] + mu[:, None, None, None] ) -> JTensor:
"""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)
return self._reverse_transform(output_ts, stats)
def _preprocess_input( def __call__(self, inputs: NestedMap) -> NestedMap:
self, """PatchTST call.
input_ts: JTensor,
input_padding: JTensor,
pos_emb: Optional[JTensor] = None,
) -> Tuple[JTensor, JTensor, Optional[Tuple[JTensor, JTensor]], JTensor]:
"""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, 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)
model_input = self.input_ff_layer(concat_inputs)
# A patch should not be padded even if there is at least one zero.
patched_padding = jnp.min(patched_pads, axis=-1)
if pos_emb is None:
position_emb = self.position_emb(seq_length=model_input.shape[1])
else:
position_emb = pos_emb
if self.do_eval:
if position_emb.shape[0] != model_input.shape[0]:
position_emb = jnp.repeat(position_emb, model_input.shape[0], axis=0)
position_emb = _shift_padded_seq(patched_padding, position_emb)
model_input += position_emb
return model_input, patched_padding, stats, patched_inputs
def _postprocess_output(
self,
model_output: JTensor,
num_outputs: int,
stats: Tuple[JTensor, JTensor],
) -> JTensor:
"""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
)
return self._reverse_transform(output_ts, stats)
def __call__(self, inputs: NestedMap) -> NestedMap:
"""PatchTST call.
Args: Args:
inputs: A NestedMap containing (1) input_ts: input sequence of shape [B, inputs: A NestedMap containing (1) input_ts: input sequence of shape [B,
@@ -374,32 +371,34 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
(2) 'output_ts' of shape [B, N, H, Q] (2) 'output_ts' of shape [B, N, H, Q]
(3) 'stats' a Tuple of statistics for renormalization. (3) 'stats' a Tuple of statistics for renormalization.
""" """
input_ts, input_padding = inputs[_INPUT_TS], inputs[_INPUT_PADDING] input_ts, input_padding = inputs[_INPUT_TS], inputs[_INPUT_PADDING]
num_outputs = len(self.quantiles) + 1 num_outputs = len(self.quantiles) + 1
model_input, patched_padding, stats, _ = self._preprocess_input( model_input, patched_padding, stats, _ = self._preprocess_input(
input_ts=input_ts, input_ts=input_ts,
input_padding=input_padding, input_padding=input_padding,
) )
if self.use_freq: if self.use_freq:
freq = inputs[_FREQ].astype(jnp.int32) freq = inputs[_FREQ].astype(jnp.int32)
f_emb = self.freq_emb(freq) # B x 1 x D f_emb = self.freq_emb(freq) # B x 1 x D
f_emb = jnp.repeat(f_emb, model_input.shape[1], axis=1) f_emb = jnp.repeat(f_emb, model_input.shape[1], axis=1)
model_input += f_emb model_input += f_emb
model_output = self.stacked_transformer_layer(model_input, patched_padding) model_output = self.stacked_transformer_layer(model_input, patched_padding)
output_ts = self._postprocess_output(model_output, num_outputs, stats) output_ts = self._postprocess_output(model_output, num_outputs, stats)
return NestedMap( return NestedMap({
{_OUTPUT_TOKENS: model_output, _OUTPUT_TS: output_ts, _STATS: stats} _OUTPUT_TOKENS: model_output,
) _OUTPUT_TS: output_ts,
_STATS: stats
})
def decode( def decode(
self, self,
inputs: NestedMap, inputs: NestedMap,
horizon_len: int, horizon_len: int,
output_patch_len: Optional[int] = None, output_patch_len: Optional[int] = None,
max_len: int = 512, max_len: int = 512,
) -> tuple[JTensor, JTensor]: ) -> tuple[JTensor, JTensor]:
"""Auto-regressive decoding without caching. """Auto-regressive decoding without caching.
Args: Args:
inputs: input time-series and paddings. Time-series shape B x C, padding inputs: input time-series and paddings. Time-series shape B x C, padding
@@ -415,83 +414,80 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
- Full predictions (mean and quantiles) as a tensor with shape - Full predictions (mean and quantiles) as a tensor with shape
B x H x (1 + # quantiles). B x H x (1 + # quantiles).
""" """
final_out = inputs[_INPUT_TS] final_out = inputs[_INPUT_TS]
inp_time_len = final_out.shape[1] inp_time_len = final_out.shape[1]
paddings = inputs[_INPUT_PADDING] paddings = inputs[_INPUT_PADDING]
if self.use_freq: if self.use_freq:
freq = inputs[_FREQ].astype(jnp.int32) freq = inputs[_FREQ].astype(jnp.int32)
else: else:
freq = jnp.zeros([final_out.shape[0], 1], dtype=jnp.int32) freq = jnp.zeros([final_out.shape[0], 1], dtype=jnp.int32)
full_outputs = [] full_outputs = []
if paddings.shape[1] != final_out.shape[1] + horizon_len: if paddings.shape[1] != final_out.shape[1] + horizon_len:
raise ValueError( raise ValueError(
"Length of paddings must match length of input + horizon_len:" "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:
if output_patch_len is None: output_patch_len = self.horizon_len
output_patch_len = self.horizon_len num_decode_patches = (horizon_len + output_patch_len -
num_decode_patches = (horizon_len + output_patch_len - 1) // output_patch_len 1) // output_patch_len
for _ in range(num_decode_patches): for _ 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_ts = final_out[:, -max_len:]
input_padding = current_padding[:, -max_len:] input_padding = current_padding[:, -max_len:]
model_input = NestedMap( model_input = NestedMap(
input_ts=input_ts, input_ts=input_ts,
input_padding=input_padding, input_padding=input_padding,
freq=freq, freq=freq,
) )
fprop_outputs = self(model_input)[_OUTPUT_TS] fprop_outputs = self(model_input)[_OUTPUT_TS]
# (full batch, last patch, output_patch_len, index of mean forecast = 0) # (full batch, last patch, output_patch_len, index of mean forecast = 0)
new_ts = fprop_outputs[:, -1, :output_patch_len, 0] new_ts = fprop_outputs[:, -1, :output_patch_len, 0]
# (full batch, last patch, output_patch_len, all output indices) # (full batch, last patch, output_patch_len, all output indices)
full_outputs.append(fprop_outputs[:, -1, :output_patch_len, :]) full_outputs.append(fprop_outputs[:, -1, :output_patch_len, :])
final_out = jnp.concatenate([final_out, new_ts], axis=-1) final_out = jnp.concatenate([final_out, new_ts], axis=-1)
return ( return (
final_out[:, inp_time_len : inp_time_len + horizon_len], final_out[:, inp_time_len:inp_time_len + horizon_len],
jnp.concatenate(full_outputs, axis=1)[:, 0:horizon_len, :], jnp.concatenate(full_outputs, axis=1)[:, 0:horizon_len, :],
) )
class PatchedDecoderFinetuneModel(base_model.BaseModel): class PatchedDecoderFinetuneModel(base_model.BaseModel):
"""Model class for finetuning patched time-series decoder. """Model class for finetuning patched time-series decoder.
Attributes: Attributes:
core_layer_tpl: config for core layer. core_layer_tpl: config for core layer.
freq: freq to finetune on. freq: freq to finetune on.
""" """
core_layer_tpl: LayerTpl = template_field(PatchedTimeSeriesDecoder) core_layer_tpl: LayerTpl = template_field(PatchedTimeSeriesDecoder)
freq: int = 0 freq: int = 0
def setup(self) -> None: def setup(self) -> None:
self.create_child("core_layer", self.core_layer_tpl) self.create_child("core_layer", self.core_layer_tpl)
def compute_predictions(self, input_batch: NestedMap) -> NestedMap: def compute_predictions(self, input_batch: NestedMap) -> NestedMap:
input_ts = input_batch[_INPUT_TS] input_ts = input_batch[_INPUT_TS]
input_padding = jnp.zeros_like(input_ts) input_padding = jnp.zeros_like(input_ts)
context_len = input_ts.shape[1] context_len = input_ts.shape[1]
input_patch_len = self.core_layer_tpl.patch_len input_patch_len = self.core_layer_tpl.patch_len
context_pad = ( context_pad = ((context_len + input_patch_len - 1) //
(context_len + input_patch_len - 1) // input_patch_len input_patch_len) * input_patch_len - context_len
) * input_patch_len - context_len
input_ts = jnp.pad(input_ts, [(0, 0), (context_pad, 0)]) input_ts = jnp.pad(input_ts, [(0, 0), (context_pad, 0)])
input_padding = jnp.pad( input_padding = jnp.pad(input_padding, [(0, 0), (context_pad, 0)],
input_padding, [(0, 0), (context_pad, 0)], constant_values=1 constant_values=1)
) freq = jnp.ones([input_ts.shape[0], 1], dtype=jnp.int32) * self.freq
freq = jnp.ones([input_ts.shape[0], 1], dtype=jnp.int32) * self.freq new_input_batch = NestedMap(
new_input_batch = NestedMap( input_ts=input_ts,
input_ts=input_ts, input_padding=input_padding,
input_padding=input_padding, freq=freq,
freq=freq, )
) return self.core_layer(new_input_batch)
return self.core_layer(new_input_batch)
def _quantile_loss( def _quantile_loss(self, pred: JTensor, actual: JTensor,
self, pred: JTensor, actual: JTensor, quantile: float quantile: float) -> JTensor:
) -> JTensor: """Calculates quantile loss.
"""Calculates quantile loss.
Args: Args:
pred: B x T pred: B x T
@@ -501,21 +497,20 @@ class PatchedDecoderFinetuneModel(base_model.BaseModel):
Returns: Returns:
per coordinate loss. per coordinate loss.
""" """
dev = actual - pred dev = actual - pred
loss_first = dev * quantile loss_first = dev * quantile
loss_second = -dev * (1.0 - quantile) loss_second = -dev * (1.0 - quantile)
return 2 * jnp.where(loss_first >= 0, loss_first, loss_second) return 2 * jnp.where(loss_first >= 0, loss_first, loss_second)
def compute_loss( def compute_loss(self, prediction_output: NestedMap,
self, prediction_output: NestedMap, input_batch: NestedMap input_batch: NestedMap) -> Tuple[NestedMap, NestedMap]:
) -> Tuple[NestedMap, NestedMap]: output_ts = prediction_output[_OUTPUT_TS]
output_ts = prediction_output[_OUTPUT_TS] actual_ts = input_batch[_TARGET_FUTURE]
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)
loss = jnp.square(pred_ts[:, :, 0] - actual_ts) for i, quantile in enumerate(self.core_layer.quantiles):
for i, quantile in enumerate(self.core_layer.quantiles): loss += self._quantile_loss(pred_ts[:, :, i + 1], actual_ts, quantile)
loss += self._quantile_loss(pred_ts[:, :, i + 1], actual_ts, quantile) loss = loss.mean()
loss = loss.mean() loss_weight = jnp.array(1.0, dtype=jnp.float32)
loss_weight = jnp.array(1.0, dtype=jnp.float32) per_example_out = NestedMap()
per_example_out = NestedMap() return {"avg_qloss": (loss, loss_weight)}, per_example_out
return {"avg_qloss": (loss, loss_weight)}, per_example_out
src/timesfm/time_features.py
+14 -14
View File
@@ -11,7 +11,6 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""Directory to extract time covariates. """Directory to extract time covariates.
Extract time covariates from datetime. Extract time covariates from datetime.
@@ -36,7 +35,6 @@ from pandas.tseries.offsets import Easter
from sklearn.preprocessing import StandardScaler from sklearn.preprocessing import StandardScaler
from tqdm import tqdm from tqdm import tqdm
# This is 183 to cover half a year (in both directions), also for leap years # This is 183 to cover half a year (in both directions), also for leap years
# + 17 as Eastern can be between March, 22 - April, 25 # + 17 as Eastern can be between March, 22 - April, 25
MAX_WINDOW = 183 + 17 MAX_WINDOW = 183 + 17
@@ -50,9 +48,8 @@ def _distance_to_holiday(holiday):
index - pd.Timedelta(days=MAX_WINDOW), index - pd.Timedelta(days=MAX_WINDOW),
index + pd.Timedelta(days=MAX_WINDOW), index + pd.Timedelta(days=MAX_WINDOW),
) )
assert ( assert (len(holiday_date) != 0 # pylint: disable=g-explicit-length-test
len(holiday_date) != 0 # pylint: disable=g-explicit-length-test ), f"No closest holiday for the date index {index} found."
), f"No closest holiday for the date index {index} found."
# It sometimes returns two dates if it is exactly half a year after the # It sometimes returns two dates if it is exactly half a year after the
# holiday. In this case, the smaller distance (182 days) is returned. # holiday. In this case, the smaller distance (182 days) is returned.
return (index - holiday_date[0]).days return (index - holiday_date[0]).days
@@ -60,16 +57,19 @@ def _distance_to_holiday(holiday):
return _distance_to_day return _distance_to_day
EasterSunday = Holiday( EasterSunday = Holiday("Easter Sunday",
"Easter Sunday", month=1, day=1, offset=[Easter(), Day(0)] month=1,
) day=1,
offset=[Easter(), Day(0)])
NewYearsDay = Holiday("New Years Day", month=1, day=1) NewYearsDay = Holiday("New Years Day", month=1, day=1)
SuperBowl = Holiday( SuperBowl = Holiday("Superbowl",
"Superbowl", month=2, day=1, offset=DateOffset(weekday=SU(1)) month=2,
) day=1,
MothersDay = Holiday( offset=DateOffset(weekday=SU(1)))
"Mothers Day", month=5, day=1, offset=DateOffset(weekday=SU(2)) MothersDay = Holiday("Mothers Day",
) month=5,
day=1,
offset=DateOffset(weekday=SU(2)))
IndependenceDay = Holiday("Independence Day", month=7, day=4) IndependenceDay = Holiday("Independence Day", month=7, day=4)
ChristmasEve = Holiday("Christmas", month=12, day=24) ChristmasEve = Holiday("Christmas", month=12, day=24)
ChristmasDay = Holiday("Christmas", month=12, day=25) ChristmasDay = Holiday("Christmas", month=12, day=25)
src/timesfm/timesfm.py
+110 -131
View File
@@ -11,7 +11,6 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""TimesFM forecast API for inference.""" """TimesFM forecast API for inference."""
import logging import logging
@@ -52,23 +51,16 @@ def moving_average(arr, window_size):
"""Calculates the moving average using NumPy's convolution function.""" """Calculates the moving average using NumPy's convolution function."""
# Pad with zeros to handle initial window positions # Pad with zeros to handle initial window positions
arr_padded = np.pad(arr, (window_size - 1, 0), "constant") arr_padded = np.pad(arr, (window_size - 1, 0), "constant")
smoothed_arr = ( smoothed_arr = (np.convolve(arr_padded, np.ones(window_size), "valid") /
np.convolve(arr_padded, np.ones(window_size), "valid") / window_size window_size)
)
return [smoothed_arr, arr - smoothed_arr] return [smoothed_arr, arr - smoothed_arr]
def freq_map(freq: str): def freq_map(freq: str):
"""Returns the frequency map for the given frequency string.""" """Returns the frequency map for the given frequency string."""
freq = str.upper(freq) freq = str.upper(freq)
if ( if (freq.endswith("H") or freq.endswith("T") or freq.endswith("MIN") or
freq.endswith("H") freq.endswith("D") or freq.endswith("B") or freq.endswith("U")):
or freq.endswith("T")
or freq.endswith("MIN")
or freq.endswith("D")
or freq.endswith("B")
or freq.endswith("U")
):
return 0 return 0
elif freq.endswith(("W", "M", "MS")): elif freq.endswith(("W", "M", "MS")):
return 1 return 1
@@ -179,9 +171,7 @@ class TimesFm:
num_layers=num_layers, num_layers=num_layers,
transformer_layer_params_tpl=pax_fiddle.Config( transformer_layer_params_tpl=pax_fiddle.Config(
transformers.Transformer, transformers.Transformer,
ln_tpl=pax_fiddle.Config( ln_tpl=pax_fiddle.Config(normalizations.RmsNorm,),
normalizations.RmsNorm,
),
), ),
), ),
) )
@@ -199,34 +189,38 @@ class TimesFm:
def _get_sample_inputs(self): def _get_sample_inputs(self):
return { return {
"input_ts": jnp.zeros( "input_ts":
( jnp.zeros(
self.per_core_batch_size, (
self.context_len + self.output_patch_len, self.per_core_batch_size,
self.context_len + self.output_patch_len,
),
dtype=jnp.float32,
), ),
dtype=jnp.float32, "input_padding":
), jnp.zeros(
"input_padding": jnp.zeros( (
( self.per_core_batch_size,
self.per_core_batch_size, self.context_len + self.output_patch_len,
self.context_len + self.output_patch_len, ),
dtype=jnp.float32,
), ),
dtype=jnp.float32, "freq":
), jnp.zeros(
"freq": jnp.zeros( (
( self.per_core_batch_size,
self.per_core_batch_size, 1,
1, ),
dtype=jnp.int32,
), ),
dtype=jnp.int32,
),
} }
def load_from_checkpoint( def load_from_checkpoint(
self, self,
checkpoint_path: Optional[str] = None, checkpoint_path: Optional[str] = None,
repo_id: str = "google/timesfm-1.0-200m", repo_id: str = "google/timesfm-1.0-200m",
checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.FLAX, checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.
FLAX,
step: int | None = None, step: int | None = None,
) -> None: ) -> None:
"""Loads a checkpoint and compiles the decoder. """Loads a checkpoint and compiles the decoder.
@@ -246,8 +240,7 @@ class TimesFm:
start_time = time.time() start_time = time.time()
self._model = instantiate(self.model_p) self._model = instantiate(self.model_p)
var_weight_hparams = self._model.abstract_init_with_metadata( var_weight_hparams = self._model.abstract_init_with_metadata(
self._get_sample_inputs(), do_eval=True self._get_sample_inputs(), do_eval=True)
)
train_state_partition_specs = tasks_lib.create_state_partition_specs( train_state_partition_specs = tasks_lib.create_state_partition_specs(
var_weight_hparams, var_weight_hparams,
mesh_shape=self.mesh_shape, mesh_shape=self.mesh_shape,
@@ -261,8 +254,7 @@ class TimesFm:
learners=None, learners=None,
) )
self._logging( self._logging(
f"Constructed model weights in {time.time() - start_time:.2f} seconds." f"Constructed model weights in {time.time() - start_time:.2f} seconds.")
)
# Load the model weights. # Load the model weights.
self._logging(f"Restoring checkpoint from {checkpoint_path}.") self._logging(f"Restoring checkpoint from {checkpoint_path}.")
@@ -275,12 +267,12 @@ class TimesFm:
step=step, step=step,
) )
self._logging( self._logging(
f"Restored checkpoint in {time.time() - start_time:.2f} seconds." f"Restored checkpoint in {time.time() - start_time:.2f} seconds.")
)
self.jit_decode() self.jit_decode()
def jit_decode(self): def jit_decode(self):
"""Jitting decoding function.""" """Jitting decoding function."""
# Initialize and jit the decode fn. # Initialize and jit the decode fn.
def _decode(inputs): def _decode(inputs):
assert self._model is not None assert self._model is not None
@@ -310,34 +302,36 @@ class TimesFm:
with base_layer.JaxContext.new_context(hparams=self._eval_context): with base_layer.JaxContext.new_context(hparams=self._eval_context):
_ = self._pmapped_decode( _ = self._pmapped_decode(
NestedMap({ NestedMap({
"input_ts": jnp.zeros( "input_ts":
( jnp.zeros(
self.num_devices, (
self.per_core_batch_size, self.num_devices,
self.context_len, self.per_core_batch_size,
self.context_len,
),
dtype=jnp.float32,
), ),
dtype=jnp.float32, "input_padding":
), jnp.zeros(
"input_padding": jnp.zeros( (
( self.num_devices,
self.num_devices, self.per_core_batch_size,
self.per_core_batch_size, self.context_len + self.horizon_len,
self.context_len + self.horizon_len, ),
dtype=jnp.float32,
), ),
dtype=jnp.float32, "date_features":
), None,
"date_features": None, "freq":
"freq": jnp.zeros( jnp.zeros(
(self.num_devices, self.per_core_batch_size, 1), (self.num_devices, self.per_core_batch_size, 1),
dtype=jnp.int32, dtype=jnp.int32,
), ),
}) }))
)
self._logging(f"Jitted decoding in {time.time() - start_time:.2f} seconds.") self._logging(f"Jitted decoding in {time.time() - start_time:.2f} seconds.")
def _preprocess( def _preprocess(self, inputs: Sequence[np.array],
self, inputs: Sequence[np.array], freq: Sequence[int] freq: Sequence[int]) -> tuple[np.array, np.array, int]:
) -> tuple[np.array, np.array, int]:
"""Formats and pads raw inputs to feed into the model. """Formats and pads raw inputs to feed into the model.
This function both pads each time series to match the context length, and This function both pads each time series to match the context length, and
@@ -358,24 +352,21 @@ class TimesFm:
input_ts, input_padding, inp_freq = [], [], [] input_ts, input_padding, inp_freq = [], [], []
pmap_pad = ( pmap_pad = ((len(inputs) - 1) // self.global_batch_size +
(len(inputs) - 1) // self.global_batch_size + 1 1) * self.global_batch_size - len(inputs)
) * self.global_batch_size - len(inputs)
for i, ts in enumerate(inputs): for i, ts in enumerate(inputs):
input_len = ts.shape[0] input_len = ts.shape[0]
padding = np.zeros(shape=(input_len + self.horizon_len,), dtype=float) padding = np.zeros(shape=(input_len + self.horizon_len,), dtype=float)
if input_len < self.context_len: if input_len < self.context_len:
num_front_pad = self.context_len - input_len num_front_pad = self.context_len - input_len
ts = np.concatenate( ts = np.concatenate([np.zeros(shape=(num_front_pad,), dtype=float), ts],
[np.zeros(shape=(num_front_pad,), dtype=float), ts], axis=0 axis=0)
)
padding = np.concatenate( padding = np.concatenate(
[np.ones(shape=(num_front_pad,), dtype=float), padding], axis=0 [np.ones(shape=(num_front_pad,), dtype=float), padding], axis=0)
)
elif input_len > self.context_len: elif input_len > self.context_len:
ts = ts[-self.context_len :] ts = ts[-self.context_len:]
padding = padding[-(self.context_len + self.horizon_len) :] padding = padding[-(self.context_len + self.horizon_len):]
input_ts.append(ts) input_ts.append(ts)
input_padding.append(padding) input_padding.append(padding)
@@ -425,8 +416,7 @@ class TimesFm:
if not self._train_state or not self._model: if not self._train_state or not self._model:
raise ValueError( raise ValueError(
"Checkpoint not loaded. Call `load_from_checkpoint` before" "Checkpoint not loaded. Call `load_from_checkpoint` before"
" `forecast`." " `forecast`.")
)
if forecast_context_len is None: if forecast_context_len is None:
forecast_context_len = self.context_len forecast_context_len = self.context_len
inputs = [np.array(ts)[-forecast_context_len:] for ts in inputs] inputs = [np.array(ts)[-forecast_context_len:] for ts in inputs]
@@ -448,47 +438,45 @@ class TimesFm:
full_outputs = [] full_outputs = []
assert input_ts.shape[0] % self.global_batch_size == 0 assert input_ts.shape[0] % self.global_batch_size == 0
for i in range(input_ts.shape[0] // self.global_batch_size): for i in range(input_ts.shape[0] // self.global_batch_size):
input_ts_in = jnp.array( input_ts_in = jnp.array(input_ts[i * self.global_batch_size:(i + 1) *
input_ts[ self.global_batch_size])
i * self.global_batch_size : (i + 1) * self.global_batch_size
]
)
input_padding_in = jnp.array( input_padding_in = jnp.array(
input_padding[ input_padding[i * self.global_batch_size:(i + 1) *
i * self.global_batch_size : (i + 1) * self.global_batch_size self.global_batch_size],)
],
)
inp_freq_in = jnp.array( inp_freq_in = jnp.array(
inp_freq[ inp_freq[i * self.global_batch_size:(i + 1) *
i * self.global_batch_size : (i + 1) * self.global_batch_size, : self.global_batch_size, :],
],
dtype=jnp.int32, dtype=jnp.int32,
) )
pmapped_inputs = NestedMap({ pmapped_inputs = NestedMap({
"input_ts": es.jax_einshape( "input_ts":
"(db)...->db...", es.jax_einshape(
input_ts_in, "(db)...->db...",
d=self.num_devices, input_ts_in,
), d=self.num_devices,
"input_padding": es.jax_einshape( ),
"(db)...->db...", "input_padding":
input_padding_in, es.jax_einshape(
d=self.num_devices, "(db)...->db...",
), input_padding_in,
"date_features": None, d=self.num_devices,
"freq": es.jax_einshape( ),
"(db)...->db...", "date_features":
inp_freq_in, None,
d=self.num_devices, "freq":
), es.jax_einshape(
"(db)...->db...",
inp_freq_in,
d=self.num_devices,
),
}) })
mean_output, full_output = self._pmapped_decode(pmapped_inputs) mean_output, full_output = self._pmapped_decode(pmapped_inputs)
mean_output = es.jax_einshape( mean_output = es.jax_einshape("db...->(db)...",
"db...->(db)...", mean_output, d=self.num_devices mean_output,
) d=self.num_devices)
full_output = es.jax_einshape( full_output = es.jax_einshape("db...->(db)...",
"db...->(db)...", full_output, d=self.num_devices full_output,
) d=self.num_devices)
mean_output = np.array(mean_output) mean_output = np.array(mean_output)
full_output = np.array(full_output) full_output = np.array(full_output)
mean_outputs.append(mean_output) mean_outputs.append(mean_output)
@@ -539,14 +527,10 @@ class TimesFm:
Returns: Returns:
Future forecasts dataframe. Future forecasts dataframe.
""" """
if not ( if not ("unique_id" in inputs.columns and "ds" in inputs.columns and
"unique_id" in inputs.columns value_name in inputs.columns):
and "ds" in inputs.columns
and value_name in inputs.columns
):
raise ValueError( raise ValueError(
f"DataFrame must have unique_id, ds and {value_name} columns." f"DataFrame must have unique_id, ds and {value_name} columns.")
)
if not forecast_context_len: if not forecast_context_len:
forecast_context_len = self.context_len forecast_context_len = self.context_len
logging.info("Preprocessing dataframe.") logging.info("Preprocessing dataframe.")
@@ -571,17 +555,15 @@ class TimesFm:
with multiprocessing.Pool(processes=num_jobs) as pool: with multiprocessing.Pool(processes=num_jobs) as pool:
results = pool.starmap( results = pool.starmap(
process_group, process_group,
[ [(key, group, value_name, forecast_context_len)
(key, group, value_name, forecast_context_len) for key, group in df_sorted.groupby("unique_id")],
for key, group in df_sorted.groupby("unique_id")
],
) )
new_inputs, uids = zip(*results) new_inputs, uids = zip(*results)
print("Finished preprocessing dataframe.") print("Finished preprocessing dataframe.")
freq_inps = [freq_map(freq)] * len(new_inputs) freq_inps = [freq_map(freq)] * len(new_inputs)
_, full_forecast = self.forecast( _, full_forecast = self.forecast(new_inputs,
new_inputs, freq=freq_inps, window_size=window_size freq=freq_inps,
) window_size=window_size)
print("Finished forecasting.") print("Finished forecasting.")
fcst_df = make_future_dataframe( fcst_df = make_future_dataframe(
uids=uids, uids=uids,
@@ -589,16 +571,13 @@ class TimesFm:
h=self.horizon_len, h=self.horizon_len,
freq=freq, freq=freq,
) )
fcst_df[model_name] = full_forecast[:, 0 : self.horizon_len, 0].reshape( fcst_df[model_name] = full_forecast[:, 0:self.horizon_len, 0].reshape(-1, 1)
-1, 1
)
if self._model.quantiles is not None: if self._model.quantiles is not None:
for i, q in enumerate(self._model.quantiles): for i, q in enumerate(self._model.quantiles):
q_col = f"{model_name}-q-{q}" q_col = f"{model_name}-q-{q}"
fcst_df[q_col] = full_forecast[:, 0 : self.horizon_len, 1 + i].reshape( fcst_df[q_col] = full_forecast[:, 0:self.horizon_len,
-1, 1 1 + i].reshape(-1, 1)
)
if q == 0.5: if q == 0.5:
fcst_df[model_name] = fcst_df[q_col] fcst_df[model_name] = fcst_df[q_col]
logging.info("Finished creating output dataframe.") logging.info("Finished creating output dataframe.")