Merge pull request #89 from google-research/rajat_dev
standardizing styles across all files.
This commit is contained in:
@@ -164,4 +164,19 @@ forecast_df = tfm.forecast_on_df(
|
||||
|
||||
## 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.
|
||||
@@ -31,191 +31,191 @@ from nixtla import NixtlaClient
|
||||
|
||||
|
||||
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:
|
||||
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 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
|
||||
"""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 = []
|
||||
# 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
|
||||
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
|
||||
)
|
||||
# 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"))
|
||||
# 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
|
||||
# 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
|
||||
"""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 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
|
||||
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
|
||||
"""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 __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 _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,
|
||||
)
|
||||
|
||||
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
|
||||
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(
|
||||
@@ -227,25 +227,27 @@ def run_timegpt(
|
||||
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
|
||||
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
|
||||
|
||||
@@ -11,7 +11,6 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Evaluation script for timegpt."""
|
||||
|
||||
import os
|
||||
@@ -25,7 +24,6 @@ from ..baselines.timegpt_pipeline import run_timegpt
|
||||
|
||||
from .utils import ExperimentHandler
|
||||
|
||||
|
||||
dataset_names = [
|
||||
"m1_monthly",
|
||||
"m1_quarterly",
|
||||
@@ -63,46 +61,45 @@ _MODEL_NAME = flags.DEFINE_string(
|
||||
)
|
||||
_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")
|
||||
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()
|
||||
FLAGS = flags.FLAGS
|
||||
FLAGS(sys.argv)
|
||||
main()
|
||||
|
||||
@@ -11,7 +11,6 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Evaluation script for timesfm."""
|
||||
|
||||
import os
|
||||
@@ -26,7 +25,6 @@ import timesfm
|
||||
|
||||
from .utils import ExperimentHandler
|
||||
|
||||
|
||||
dataset_names = [
|
||||
"m1_monthly",
|
||||
"m1_quarterly",
|
||||
@@ -74,16 +72,14 @@ context_dict = {
|
||||
"m4_yearly": 64,
|
||||
}
|
||||
|
||||
_MODEL_PATH = flags.DEFINE_string(
|
||||
"model_path", "/home/timesfm_q10_20240501", "Path to model"
|
||||
)
|
||||
_MODEL_PATH = flags.DEFINE_string("model_path", "/home/timesfm_q10_20240501",
|
||||
"Path to model")
|
||||
_BATCH_SIZE = flags.DEFINE_integer("batch_size", 64, "Batch size")
|
||||
_HORIZON = flags.DEFINE_integer("horizon", 128, "Horizon")
|
||||
_BACKEND = flags.DEFINE_string("backend", "gpu", "Backend")
|
||||
_NUM_JOBS = flags.DEFINE_integer("num_jobs", 1, "Number of jobs")
|
||||
_SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory")
|
||||
|
||||
|
||||
QUANTILES = list(np.arange(1, 10) / 10.0)
|
||||
|
||||
|
||||
@@ -127,9 +123,9 @@ def main():
|
||||
)
|
||||
total_time = time.time() - init_time
|
||||
time_df = pd.DataFrame({"time": [total_time], "model": model_name})
|
||||
results = exp.evaluate_from_predictions(
|
||||
models=[model_name], fcsts_df=fcsts_df, times_df=time_df
|
||||
)
|
||||
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)
|
||||
|
||||
@@ -11,7 +11,6 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Forked from https://github.com/Nixtla/nixtla/blob/main/experiments/amazon-chronos/src/utils.py."""
|
||||
|
||||
from functools import partial
|
||||
@@ -46,11 +45,9 @@ def quantile_loss(
|
||||
target_col: str = "y",
|
||||
) -> pd.DataFrame:
|
||||
delta_y = df[models].sub(df[target_col], axis=0)
|
||||
res = (
|
||||
np.maximum(q * delta_y, (q - 1) * delta_y)
|
||||
.groupby(df[id_col], observed=True)
|
||||
.mean()
|
||||
)
|
||||
res = (np.maximum(q * delta_y,
|
||||
(q - 1) * delta_y).groupby(df[id_col],
|
||||
observed=True).mean())
|
||||
res.index.name = id_col
|
||||
res = res.reset_index()
|
||||
return res
|
||||
@@ -66,10 +63,8 @@ class ExperimentHandler:
|
||||
models_dir: str = "./models",
|
||||
):
|
||||
if dataset not in gluonts_datasets:
|
||||
raise Exception(
|
||||
f"dataset {dataset} not found in gluonts "
|
||||
f"available datasets: {', '.join(gluonts_datasets)}"
|
||||
)
|
||||
raise Exception(f"dataset {dataset} not found in gluonts "
|
||||
f"available datasets: {', '.join(gluonts_datasets)}")
|
||||
self.dataset = dataset
|
||||
self.quantiles = quantiles
|
||||
self.level = self._transform_quantiles_to_levels(quantiles)
|
||||
@@ -80,10 +75,8 @@ class ExperimentHandler:
|
||||
gluonts_dataset = get_dataset(self.dataset)
|
||||
self.horizon = gluonts_dataset.metadata.prediction_length
|
||||
if self.horizon is None:
|
||||
raise Exception(
|
||||
f"horizon not found for dataset {self.dataset} "
|
||||
"experiment cannot be run"
|
||||
)
|
||||
raise Exception(f"horizon not found for dataset {self.dataset} "
|
||||
"experiment cannot be run")
|
||||
self.freq = gluonts_dataset.metadata.freq
|
||||
# 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
|
||||
@@ -122,9 +115,8 @@ class ExperimentHandler:
|
||||
|
||||
@staticmethod
|
||||
def _transform_quantiles_to_levels(quantiles: List[float]) -> List[int]:
|
||||
level = [
|
||||
int(100 - 200 * q) for q in quantiles if q < 0.5
|
||||
] # in this case mean=mediain
|
||||
level = [int(100 - 200 * q) for q in quantiles if q < 0.5
|
||||
] # in this case mean=mediain
|
||||
level = sorted(list(set(level)))
|
||||
return level
|
||||
|
||||
@@ -153,9 +145,8 @@ class ExperimentHandler:
|
||||
last_n: int | None = None,
|
||||
) -> pd.DataFrame:
|
||||
with multiprocessing.Pool(os.cpu_count()) as pool: # Create a process pool
|
||||
results = pool.map(
|
||||
parallel_transform, zip(gluonts_dataset, repeat(last_n))
|
||||
)
|
||||
results = pool.map(parallel_transform, zip(gluonts_dataset,
|
||||
repeat(last_n)))
|
||||
df = pd.concat(results)
|
||||
df = df.reset_index(drop=True)
|
||||
return df
|
||||
@@ -177,9 +168,8 @@ class ExperimentHandler:
|
||||
def save_dataframe(self, df: pd.DataFrame, file_name: str):
|
||||
df.to_csv(f"{self.results_dir}/{file_name}", index=False)
|
||||
|
||||
def save_results(
|
||||
self, fcst_df: pd.DataFrame, total_time: float, model_name: str
|
||||
):
|
||||
def save_results(self, fcst_df: pd.DataFrame, total_time: float,
|
||||
model_name: str):
|
||||
self.save_dataframe(
|
||||
fcst_df,
|
||||
f"{model_name}-{self.dataset}-fcst.csv",
|
||||
@@ -215,23 +205,21 @@ class ExperimentHandler:
|
||||
times_df = []
|
||||
for model in models:
|
||||
fcst_method_df = pd.read_csv(
|
||||
f"{self.results_dir}/{model}-{self.dataset}-fcst.csv"
|
||||
).set_index(["unique_id", "ds"])
|
||||
f"{self.results_dir}/{model}-{self.dataset}-fcst.csv").set_index(
|
||||
["unique_id", "ds"])
|
||||
fcsts_df.append(fcst_method_df)
|
||||
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)
|
||||
fcsts_df = pd.concat(fcsts_df, axis=1).reset_index()
|
||||
fcsts_df["ds"] = pd.to_datetime(fcsts_df["ds"])
|
||||
times_df = pd.concat(times_df)
|
||||
return self.evaluate_from_predictions(
|
||||
models=models, fcsts_df=fcsts_df, times_df=times_df
|
||||
)
|
||||
return self.evaluate_from_predictions(models=models,
|
||||
fcsts_df=fcsts_df,
|
||||
times_df=times_df)
|
||||
|
||||
def evaluate_from_predictions(
|
||||
self, models: List[str], fcsts_df: pd.DataFrame, times_df: pd.DataFrame
|
||||
) -> pd.DataFrame:
|
||||
def evaluate_from_predictions(self, models: List[str], fcsts_df: pd.DataFrame,
|
||||
times_df: pd.DataFrame) -> pd.DataFrame:
|
||||
test_df = self.test_df
|
||||
train_df = self.train_df
|
||||
test_df = test_df.merge(fcsts_df, how="left")
|
||||
@@ -262,9 +250,9 @@ class ExperimentHandler:
|
||||
eval_prob_df["metric"] = "scaled_crps"
|
||||
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.melt(
|
||||
id_vars="metric", value_name="value", var_name="model"
|
||||
)
|
||||
eval_df = eval_df.melt(id_vars="metric",
|
||||
value_name="value",
|
||||
var_name="model")
|
||||
times_df.insert(0, "metric", "time")
|
||||
times_df = times_df.rename(columns={"time": "value"})
|
||||
eval_df = pd.concat([eval_df, times_df])
|
||||
|
||||
@@ -11,7 +11,6 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Eval pipeline."""
|
||||
|
||||
import json
|
||||
@@ -28,38 +27,28 @@ import torch
|
||||
import tqdm
|
||||
from timesfm import data_loader
|
||||
|
||||
|
||||
FLAGS = flags.FLAGS
|
||||
|
||||
_BATCH_SIZE = flags.DEFINE_integer(
|
||||
"batch_size", 64, "Batch size for the randomly sampled batch"
|
||||
)
|
||||
_BATCH_SIZE = flags.DEFINE_integer("batch_size", 64,
|
||||
"Batch size for the randomly sampled batch")
|
||||
_DATASET = flags.DEFINE_string("dataset", "etth1", "The name of the dataset.")
|
||||
_MODEL_PATH = flags.DEFINE_string(
|
||||
"model_path", "./timesfm_q10_20240501", "The name of the dataset."
|
||||
)
|
||||
_DATETIME_COL = flags.DEFINE_string(
|
||||
"datetime_col", "date", "Column having datetime."
|
||||
)
|
||||
_NUM_COV_COLS = flags.DEFINE_list(
|
||||
"num_cov_cols", None, "Column having numerical features."
|
||||
)
|
||||
_CAT_COV_COLS = flags.DEFINE_list(
|
||||
"cat_cov_cols", None, "Column having categorical features."
|
||||
)
|
||||
_MODEL_PATH = flags.DEFINE_string("model_path", "./timesfm_q10_20240501",
|
||||
"The name of the dataset.")
|
||||
_DATETIME_COL = flags.DEFINE_string("datetime_col", "date",
|
||||
"Column having datetime.")
|
||||
_NUM_COV_COLS = flags.DEFINE_list("num_cov_cols", None,
|
||||
"Column having numerical 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")
|
||||
_NORMALIZE = flags.DEFINE_bool(
|
||||
"normalize", True, "normalize data for eval or not"
|
||||
)
|
||||
_CONTEXT_LEN = flags.DEFINE_integer(
|
||||
"context_len", 512, "Length of the context window"
|
||||
)
|
||||
_NORMALIZE = flags.DEFINE_bool("normalize", True,
|
||||
"normalize data for eval or not")
|
||||
_CONTEXT_LEN = flags.DEFINE_integer("context_len", 512,
|
||||
"Length of the context window")
|
||||
_PRED_LEN = flags.DEFINE_integer("pred_len", 96, "prediction length.")
|
||||
_BACKEND = flags.DEFINE_string("backend", "gpu", "backend to use")
|
||||
_RESULTS_DIR = flags.DEFINE_string(
|
||||
"results_dir", "./results/long_horizon", "results directory"
|
||||
)
|
||||
|
||||
_RESULTS_DIR = flags.DEFINE_string("results_dir", "./results/long_horizon",
|
||||
"results directory")
|
||||
|
||||
DATA_DICT = {
|
||||
"ettm2": {
|
||||
@@ -176,9 +165,8 @@ def eval():
|
||||
holiday=False,
|
||||
permute=False,
|
||||
)
|
||||
eval_itr = dtl.tf_dataset(
|
||||
mode="test", shift=_PRED_LEN.value
|
||||
).as_numpy_iterator()
|
||||
eval_itr = dtl.tf_dataset(mode="test",
|
||||
shift=_PRED_LEN.value).as_numpy_iterator()
|
||||
model_path = _MODEL_PATH.value
|
||||
if model_path.startswith("amazon"):
|
||||
model = chronos.ChronosPipeline.from_pretrained(
|
||||
@@ -213,10 +201,9 @@ def eval():
|
||||
for batch in tqdm.tqdm(eval_itr):
|
||||
past = batch[0]
|
||||
actuals = batch[3]
|
||||
forecasts = get_forecasts(
|
||||
model_path, model, past, int_freq, _PRED_LEN.value
|
||||
)
|
||||
forecasts = forecasts[:, 0 : actuals.shape[1]]
|
||||
forecasts = get_forecasts(model_path, model, past, int_freq,
|
||||
_PRED_LEN.value)
|
||||
forecasts = forecasts[:, 0:actuals.shape[1]]
|
||||
mae_run_losses.append(_mae(forecasts, actuals).sum())
|
||||
mse_run_losses.append(_mse(forecasts, actuals).sum())
|
||||
smape_run_losses.append(_smape(forecasts, actuals).sum())
|
||||
|
||||
@@ -11,7 +11,6 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""TimesFM init file."""
|
||||
|
||||
from .timesfm import TimesFm, freq_map
|
||||
|
||||
@@ -11,13 +11,11 @@
|
||||
# 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.
|
||||
|
||||
"""TF dataloaders for general timeseries datasets.
|
||||
|
||||
The expected input format is csv file with a datetime index.
|
||||
"""
|
||||
|
||||
|
||||
from absl import logging
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
@@ -79,9 +77,8 @@ class TimeSeriesdata(object):
|
||||
self.data_df['ccol'] = np.zeros(self.data_df.shape[0])
|
||||
cat_cov_cols = ['ccol']
|
||||
self.data_df.fillna(0, inplace=True)
|
||||
self.data_df.set_index(
|
||||
pd.DatetimeIndex(self.data_df[datetime_col]), inplace=True
|
||||
)
|
||||
self.data_df.set_index(pd.DatetimeIndex(self.data_df[datetime_col]),
|
||||
inplace=True)
|
||||
self.num_cov_cols = num_cov_cols
|
||||
self.cat_cov_cols = cat_cov_cols
|
||||
self.ts_cols = ts_cols
|
||||
@@ -94,18 +91,16 @@ class TimeSeriesdata(object):
|
||||
data_df_idx[-1] + pd.Timedelta(1, freq=freq),
|
||||
periods=pred_len + 1,
|
||||
freq=freq,
|
||||
)
|
||||
)
|
||||
))
|
||||
self.time_df = time_features.TimeCovariates(
|
||||
date_index, holiday=holiday
|
||||
).get_covariates()
|
||||
date_index, holiday=holiday).get_covariates()
|
||||
self.hist_len = hist_len
|
||||
self.pred_len = pred_len
|
||||
self.batch_size = batch_size
|
||||
self.freq = freq
|
||||
self.normalize = normalize
|
||||
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.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)
|
||||
@@ -135,7 +130,7 @@ class TimeSeriesdata(object):
|
||||
|
||||
def _normalize_data(self):
|
||||
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.data_mat = self.scaler.transform(self.data_mat.transpose()).transpose()
|
||||
|
||||
@@ -253,9 +248,8 @@ class TimeSeriesdata(object):
|
||||
gen_fn = self.train_gen
|
||||
else:
|
||||
gen_fn = lambda: self.test_val_gen(mode, shift)
|
||||
output_types = tuple(
|
||||
[tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 + [tf.int32] * 2
|
||||
)
|
||||
output_types = tuple([tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 +
|
||||
[tf.int32] * 2)
|
||||
dataset = tf.data.Dataset.from_generator(gen_fn, output_types)
|
||||
dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
|
||||
return dataset
|
||||
|
||||
+318
-323
@@ -11,7 +11,6 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Pax ML model for patched time-series decoder.
|
||||
|
||||
The file implements Residual MLPs, Patched Decoder layers and PAX ML models.
|
||||
@@ -36,7 +35,6 @@ from praxis.layers import normalizations
|
||||
from praxis.layers import stochastics
|
||||
from praxis.layers import transformers
|
||||
|
||||
|
||||
# PAX shortcuts
|
||||
NestedMap = py_utils.NestedMap
|
||||
JTensor = pytypes.JTensor
|
||||
@@ -44,7 +42,6 @@ JTensor = pytypes.JTensor
|
||||
LayerTpl = pax_fiddle.Config[base_layer.BaseLayer]
|
||||
template_field = base_layer.template_field
|
||||
|
||||
|
||||
PAD_VAL = 1123581321.0
|
||||
DEFAULT_QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
|
||||
|
||||
@@ -57,36 +54,35 @@ _FREQ = "freq"
|
||||
_OUTPUT_TOKENS = "output_tokens"
|
||||
_STATS = "stats"
|
||||
|
||||
|
||||
# Small numerical value.
|
||||
_TOLERANCE = 1e-7
|
||||
|
||||
|
||||
def _shift_padded_seq(mask: JTensor, seq: JTensor) -> JTensor:
|
||||
"""Shifts rows of seq based on the first 0 in each row of the mask."""
|
||||
num = seq.shape[1]
|
||||
"""Shifts rows of seq based on the first 0 in each row of the mask."""
|
||||
num = seq.shape[1]
|
||||
|
||||
# Find the index of the first 0 in each row of the mask
|
||||
first_zero_idx = jnp.argmin(mask, axis=1)
|
||||
# Find the index of the first 0 in each row of the mask
|
||||
first_zero_idx = jnp.argmin(mask, axis=1)
|
||||
|
||||
# Create a range array for indexing
|
||||
idx_range = jnp.arange(num)
|
||||
# Create a range array for indexing
|
||||
idx_range = jnp.arange(num)
|
||||
|
||||
def shift_row(carry, x):
|
||||
seq_row, shift = x
|
||||
shifted_idx = (idx_range - shift) % num
|
||||
shifted_row = seq_row[shifted_idx]
|
||||
return carry, shifted_row
|
||||
def shift_row(carry, x):
|
||||
seq_row, shift = x
|
||||
shifted_idx = (idx_range - shift) % num
|
||||
shifted_row = seq_row[shifted_idx]
|
||||
return carry, shifted_row
|
||||
|
||||
# Use lax.scan to shift each row of seq based on the corresponding
|
||||
# first_zero_idx.
|
||||
_, shifted_seq = lax.scan(shift_row, None, (seq, first_zero_idx))
|
||||
# Use lax.scan to shift each row of seq based on the corresponding
|
||||
# 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):
|
||||
"""Simple feedforward block with residual connection.
|
||||
"""Simple feedforward block with residual connection.
|
||||
|
||||
Attributes:
|
||||
input_dims: input dimension.
|
||||
@@ -99,67 +95,68 @@ class ResidualBlock(base_layer.BaseLayer):
|
||||
act_tpl: config for activation in hidden layer.
|
||||
"""
|
||||
|
||||
input_dims: int = 0
|
||||
hidden_dims: int = 0
|
||||
output_dims: int = 0
|
||||
dropout_prob: float = 0.0
|
||||
layer_norm: bool = False
|
||||
dropout_tpl: LayerTpl = template_field(stochastics.Dropout)
|
||||
ln_tpl: LayerTpl = template_field(normalizations.LayerNorm)
|
||||
act_tpl: LayerTpl = template_field(activations.Swish)
|
||||
input_dims: int = 0
|
||||
hidden_dims: int = 0
|
||||
output_dims: int = 0
|
||||
dropout_prob: float = 0.0
|
||||
layer_norm: bool = False
|
||||
dropout_tpl: LayerTpl = template_field(stochastics.Dropout)
|
||||
ln_tpl: LayerTpl = template_field(normalizations.LayerNorm)
|
||||
act_tpl: LayerTpl = template_field(activations.Swish)
|
||||
|
||||
def setup(self):
|
||||
lnorm_tpl = self.ln_tpl.clone()
|
||||
lnorm_tpl.dim = self.output_dims
|
||||
self.create_child("ln_layer", lnorm_tpl)
|
||||
def setup(self):
|
||||
lnorm_tpl = self.ln_tpl.clone()
|
||||
lnorm_tpl.dim = self.output_dims
|
||||
self.create_child("ln_layer", lnorm_tpl)
|
||||
|
||||
dropout_tpl = self.dropout_tpl.clone()
|
||||
dropout_tpl.keep_prob = 1.0 - self.dropout_prob
|
||||
self.create_child("dropout", dropout_tpl)
|
||||
dropout_tpl = self.dropout_tpl.clone()
|
||||
dropout_tpl.keep_prob = 1.0 - self.dropout_prob
|
||||
self.create_child("dropout", dropout_tpl)
|
||||
|
||||
self.create_child(
|
||||
"hidden_layer",
|
||||
pax_fiddle.Config(
|
||||
linears.FeedForward,
|
||||
input_dims=self.input_dims,
|
||||
output_dims=self.hidden_dims,
|
||||
activation_tpl=self.act_tpl.clone(),
|
||||
),
|
||||
)
|
||||
self.create_child(
|
||||
"hidden_layer",
|
||||
pax_fiddle.Config(
|
||||
linears.FeedForward,
|
||||
input_dims=self.input_dims,
|
||||
output_dims=self.hidden_dims,
|
||||
activation_tpl=self.act_tpl.clone(),
|
||||
),
|
||||
)
|
||||
|
||||
self.create_child(
|
||||
"output_layer",
|
||||
pax_fiddle.Config(
|
||||
linears.FeedForward,
|
||||
input_dims=self.hidden_dims,
|
||||
output_dims=self.output_dims,
|
||||
activation_tpl=pax_fiddle.Config(activations.Identity),
|
||||
),
|
||||
)
|
||||
self.create_child(
|
||||
"output_layer",
|
||||
pax_fiddle.Config(
|
||||
linears.FeedForward,
|
||||
input_dims=self.hidden_dims,
|
||||
output_dims=self.output_dims,
|
||||
activation_tpl=pax_fiddle.Config(activations.Identity),
|
||||
),
|
||||
)
|
||||
|
||||
self.create_child(
|
||||
"residual_layer",
|
||||
pax_fiddle.Config(
|
||||
linears.FeedForward,
|
||||
input_dims=self.input_dims,
|
||||
output_dims=self.output_dims,
|
||||
activation_tpl=pax_fiddle.Config(activations.Identity),
|
||||
),
|
||||
)
|
||||
self.create_child(
|
||||
"residual_layer",
|
||||
pax_fiddle.Config(
|
||||
linears.FeedForward,
|
||||
input_dims=self.input_dims,
|
||||
output_dims=self.output_dims,
|
||||
activation_tpl=pax_fiddle.Config(activations.Identity),
|
||||
),
|
||||
)
|
||||
|
||||
def __call__(self, inputs: JTensor) -> JTensor:
|
||||
hidden = self.hidden_layer(inputs)
|
||||
output = self.output_layer(hidden)
|
||||
output = self.dropout(output)
|
||||
residual = self.residual_layer(inputs)
|
||||
if self.layer_norm:
|
||||
return self.ln_layer(output + residual)
|
||||
else:
|
||||
return output + residual
|
||||
def __call__(self, inputs: JTensor) -> JTensor:
|
||||
hidden = self.hidden_layer(inputs)
|
||||
output = self.output_layer(hidden)
|
||||
output = self.dropout(output)
|
||||
residual = self.residual_layer(inputs)
|
||||
if self.layer_norm:
|
||||
return self.ln_layer(output + residual)
|
||||
else:
|
||||
return output + residual
|
||||
|
||||
|
||||
def _masked_mean_std(inputs: JTensor, padding: JTensor) -> Tuple[JTensor, JTensor]:
|
||||
"""Calculates mean and standard deviation of arr across axis 1.
|
||||
def _masked_mean_std(inputs: JTensor,
|
||||
padding: JTensor) -> Tuple[JTensor, JTensor]:
|
||||
"""Calculates mean and standard deviation of arr across axis 1.
|
||||
|
||||
It should exclude values where pad is 1.
|
||||
|
||||
@@ -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
|
||||
statistics of the first patch with more than three non-padded values.
|
||||
"""
|
||||
# Selecting the first pad with more than 3 unpadded values.
|
||||
pad_sum = jnp.sum(1 - padding, axis=2)
|
||||
# Selecting the first pad with more than 3 unpadded values.
|
||||
pad_sum = jnp.sum(1 - padding, axis=2)
|
||||
|
||||
def _get_patch_index(arr: JTensor):
|
||||
indices = jnp.argmax(arr >= 3, axis=1)
|
||||
row_sum = (arr >= 3).sum(axis=1)
|
||||
return jnp.where(row_sum == 0, arr.shape[1] - 1, indices)
|
||||
def _get_patch_index(arr: JTensor):
|
||||
indices = jnp.argmax(arr >= 3, axis=1)
|
||||
row_sum = (arr >= 3).sum(axis=1)
|
||||
return jnp.where(row_sum == 0, arr.shape[1] - 1, indices)
|
||||
|
||||
patch_indices = _get_patch_index(pad_sum)
|
||||
bidxs = jnp.arange(inputs.shape[0])
|
||||
patch_indices = _get_patch_index(pad_sum)
|
||||
bidxs = jnp.arange(inputs.shape[0])
|
||||
|
||||
arr = inputs[bidxs, patch_indices, :]
|
||||
pad = padding[bidxs, patch_indices, :]
|
||||
arr = inputs[bidxs, patch_indices, :]
|
||||
pad = padding[bidxs, patch_indices, :]
|
||||
|
||||
# Create a mask where P is 0
|
||||
mask = 1 - pad
|
||||
# Create a mask where P is 0
|
||||
mask = 1 - pad
|
||||
|
||||
# Calculate the number of valid elements
|
||||
num_valid_elements = jnp.sum(mask, axis=1)
|
||||
# Calculate the number of valid elements
|
||||
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
|
||||
masked_sum = jnp.sum(arr * mask, axis=1)
|
||||
masked_squared_sum = jnp.sum((arr * mask) ** 2, axis=1)
|
||||
# Calculate the masked sum and squared sum of M
|
||||
masked_sum = jnp.sum(arr * mask, axis=1)
|
||||
masked_squared_sum = jnp.sum((arr * mask)**2, axis=1)
|
||||
|
||||
# Calculate the masked mean and standard deviation
|
||||
masked_mean = masked_sum / num_valid_elements
|
||||
masked_var = masked_squared_sum / num_valid_elements - masked_mean**2
|
||||
masked_var = jnp.where(masked_var < 0.0, 0.0, masked_var)
|
||||
masked_std = jnp.sqrt(masked_var)
|
||||
# Calculate the masked mean and standard deviation
|
||||
masked_mean = masked_sum / num_valid_elements
|
||||
masked_var = masked_squared_sum / num_valid_elements - masked_mean**2
|
||||
masked_var = jnp.where(masked_var < 0.0, 0.0, masked_var)
|
||||
masked_std = jnp.sqrt(masked_var)
|
||||
|
||||
return masked_mean, masked_std
|
||||
return masked_mean, masked_std
|
||||
|
||||
|
||||
def _create_quantiles() -> list[float]:
|
||||
"""Returns the quantiles for forecasting."""
|
||||
return DEFAULT_QUANTILES
|
||||
"""Returns the quantiles for forecasting."""
|
||||
return DEFAULT_QUANTILES
|
||||
|
||||
|
||||
class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
|
||||
"""Patch decoder layer for time-series foundation model.
|
||||
"""Patch decoder layer for time-series foundation model.
|
||||
|
||||
Attributes:
|
||||
patch_len: length of input patches.
|
||||
@@ -231,137 +228,137 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
|
||||
number of output logits. D is model dimension.
|
||||
"""
|
||||
|
||||
patch_len: int = 0
|
||||
horizon_len: int = 0
|
||||
model_dims: int = 0
|
||||
hidden_dims: int = 0
|
||||
quantiles: list[float] = dataclasses.field(default_factory=_create_quantiles)
|
||||
residual_block_tpl: LayerTpl = template_field(ResidualBlock)
|
||||
stacked_transformer_params_tpl: LayerTpl = template_field(
|
||||
transformers.StackedTransformer
|
||||
patch_len: int = 0
|
||||
horizon_len: int = 0
|
||||
model_dims: int = 0
|
||||
hidden_dims: int = 0
|
||||
quantiles: list[float] = dataclasses.field(default_factory=_create_quantiles)
|
||||
residual_block_tpl: LayerTpl = template_field(ResidualBlock)
|
||||
stacked_transformer_params_tpl: LayerTpl = template_field(
|
||||
transformers.StackedTransformer)
|
||||
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:
|
||||
"""Construct the model."""
|
||||
num_outputs = len(self.quantiles) + 1
|
||||
horizon_resl = self.residual_block_tpl.clone()
|
||||
horizon_resl.input_dims = self.model_dims
|
||||
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()
|
||||
stl.model_dims = self.model_dims
|
||||
stl.hidden_dims = self.hidden_dims
|
||||
stl.mask_self_attention = True
|
||||
self.create_child(
|
||||
"position_emb",
|
||||
pax_fiddle.Config(layers.PositionalEmbedding,
|
||||
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()
|
||||
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,
|
||||
)
|
||||
def transform_decode_state(
|
||||
self, transform_fn: base_layer.DecodeStateTransformFn) -> None:
|
||||
"""Transforms all decode state variables based on transform_fn."""
|
||||
self.stacked_transformer_layer.transform_decode_state(transform_fn)
|
||||
|
||||
horizon_resl = self.residual_block_tpl.clone()
|
||||
horizon_resl.input_dims = self.model_dims
|
||||
horizon_resl.hidden_dims = self.hidden_dims
|
||||
horizon_resl.output_dims = self.horizon_len * num_outputs
|
||||
self.create_child(
|
||||
"horizon_ff_layer",
|
||||
horizon_resl,
|
||||
)
|
||||
def _forward_transform(
|
||||
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)
|
||||
|
||||
self.create_child(
|
||||
"position_emb",
|
||||
pax_fiddle.Config(
|
||||
layers.PositionalEmbedding, embedding_dims=self.model_dims
|
||||
),
|
||||
)
|
||||
def _reverse_transform(self, outputs: JTensor,
|
||||
stats: Tuple[JTensor, JTensor]) -> JTensor:
|
||||
"""Output is of shape [B, N, P, Q]."""
|
||||
mu, sigma = stats
|
||||
return outputs * sigma[:, None, None, None] + mu[:, None, None, None]
|
||||
|
||||
if self.use_freq:
|
||||
self.create_child(
|
||||
"freq_emb",
|
||||
pax_fiddle.Config(
|
||||
embedding_softmax.Embedding,
|
||||
num_classes=3,
|
||||
input_dims=self.model_dims,
|
||||
),
|
||||
)
|
||||
def _preprocess_input(
|
||||
self,
|
||||
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)
|
||||
|
||||
def transform_decode_state(
|
||||
self, transform_fn: base_layer.DecodeStateTransformFn
|
||||
) -> None:
|
||||
"""Transforms all decode state variables based on transform_fn."""
|
||||
self.stacked_transformer_layer.transform_decode_state(transform_fn)
|
||||
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
|
||||
|
||||
def _forward_transform(
|
||||
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)
|
||||
return model_input, patched_padding, stats, patched_inputs
|
||||
|
||||
def _reverse_transform(
|
||||
self, outputs: JTensor, stats: Tuple[JTensor, JTensor]
|
||||
) -> JTensor:
|
||||
"""Output is of shape [B, N, P, Q]."""
|
||||
mu, sigma = stats
|
||||
return outputs * sigma[:, None, None, None] + mu[:, None, None, None]
|
||||
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 _preprocess_input(
|
||||
self,
|
||||
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.
|
||||
def __call__(self, inputs: NestedMap) -> NestedMap:
|
||||
"""PatchTST call.
|
||||
|
||||
Args:
|
||||
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]
|
||||
(3) 'stats' a Tuple of statistics for renormalization.
|
||||
"""
|
||||
input_ts, input_padding = inputs[_INPUT_TS], inputs[_INPUT_PADDING]
|
||||
num_outputs = len(self.quantiles) + 1
|
||||
model_input, patched_padding, stats, _ = self._preprocess_input(
|
||||
input_ts=input_ts,
|
||||
input_padding=input_padding,
|
||||
)
|
||||
if self.use_freq:
|
||||
freq = inputs[_FREQ].astype(jnp.int32)
|
||||
f_emb = self.freq_emb(freq) # B x 1 x D
|
||||
f_emb = jnp.repeat(f_emb, model_input.shape[1], axis=1)
|
||||
model_input += f_emb
|
||||
model_output = self.stacked_transformer_layer(model_input, patched_padding)
|
||||
input_ts, input_padding = inputs[_INPUT_TS], inputs[_INPUT_PADDING]
|
||||
num_outputs = len(self.quantiles) + 1
|
||||
model_input, patched_padding, stats, _ = self._preprocess_input(
|
||||
input_ts=input_ts,
|
||||
input_padding=input_padding,
|
||||
)
|
||||
if self.use_freq:
|
||||
freq = inputs[_FREQ].astype(jnp.int32)
|
||||
f_emb = self.freq_emb(freq) # B x 1 x D
|
||||
f_emb = jnp.repeat(f_emb, model_input.shape[1], axis=1)
|
||||
model_input += f_emb
|
||||
model_output = self.stacked_transformer_layer(model_input, patched_padding)
|
||||
|
||||
output_ts = self._postprocess_output(model_output, num_outputs, stats)
|
||||
return NestedMap(
|
||||
{_OUTPUT_TOKENS: model_output, _OUTPUT_TS: output_ts, _STATS: stats}
|
||||
)
|
||||
output_ts = self._postprocess_output(model_output, num_outputs, stats)
|
||||
return NestedMap({
|
||||
_OUTPUT_TOKENS: model_output,
|
||||
_OUTPUT_TS: output_ts,
|
||||
_STATS: stats
|
||||
})
|
||||
|
||||
def decode(
|
||||
self,
|
||||
inputs: NestedMap,
|
||||
horizon_len: int,
|
||||
output_patch_len: Optional[int] = None,
|
||||
max_len: int = 512,
|
||||
) -> tuple[JTensor, JTensor]:
|
||||
"""Auto-regressive decoding without caching.
|
||||
def decode(
|
||||
self,
|
||||
inputs: NestedMap,
|
||||
horizon_len: int,
|
||||
output_patch_len: Optional[int] = None,
|
||||
max_len: int = 512,
|
||||
) -> tuple[JTensor, JTensor]:
|
||||
"""Auto-regressive decoding without caching.
|
||||
|
||||
Args:
|
||||
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
|
||||
B x H x (1 + # quantiles).
|
||||
"""
|
||||
final_out = inputs[_INPUT_TS]
|
||||
inp_time_len = final_out.shape[1]
|
||||
paddings = inputs[_INPUT_PADDING]
|
||||
if self.use_freq:
|
||||
freq = inputs[_FREQ].astype(jnp.int32)
|
||||
else:
|
||||
freq = jnp.zeros([final_out.shape[0], 1], dtype=jnp.int32)
|
||||
full_outputs = []
|
||||
if paddings.shape[1] != final_out.shape[1] + horizon_len:
|
||||
raise ValueError(
|
||||
"Length of paddings must match length of input + horizon_len:"
|
||||
f" {paddings.shape[1]} != {final_out.shape[1]} + {horizon_len}"
|
||||
)
|
||||
if output_patch_len is None:
|
||||
output_patch_len = self.horizon_len
|
||||
num_decode_patches = (horizon_len + output_patch_len - 1) // output_patch_len
|
||||
for _ 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:]
|
||||
model_input = NestedMap(
|
||||
input_ts=input_ts,
|
||||
input_padding=input_padding,
|
||||
freq=freq,
|
||||
)
|
||||
fprop_outputs = self(model_input)[_OUTPUT_TS]
|
||||
# (full batch, last patch, output_patch_len, index of mean forecast = 0)
|
||||
new_ts = fprop_outputs[:, -1, :output_patch_len, 0]
|
||||
# (full batch, last patch, output_patch_len, all output indices)
|
||||
full_outputs.append(fprop_outputs[:, -1, :output_patch_len, :])
|
||||
final_out = jnp.concatenate([final_out, new_ts], axis=-1)
|
||||
final_out = inputs[_INPUT_TS]
|
||||
inp_time_len = final_out.shape[1]
|
||||
paddings = inputs[_INPUT_PADDING]
|
||||
if self.use_freq:
|
||||
freq = inputs[_FREQ].astype(jnp.int32)
|
||||
else:
|
||||
freq = jnp.zeros([final_out.shape[0], 1], dtype=jnp.int32)
|
||||
full_outputs = []
|
||||
if paddings.shape[1] != final_out.shape[1] + horizon_len:
|
||||
raise ValueError(
|
||||
"Length of paddings must match length of input + horizon_len:"
|
||||
f" {paddings.shape[1]} != {final_out.shape[1]} + {horizon_len}")
|
||||
if output_patch_len is None:
|
||||
output_patch_len = self.horizon_len
|
||||
num_decode_patches = (horizon_len + output_patch_len -
|
||||
1) // output_patch_len
|
||||
for _ 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:]
|
||||
model_input = NestedMap(
|
||||
input_ts=input_ts,
|
||||
input_padding=input_padding,
|
||||
freq=freq,
|
||||
)
|
||||
fprop_outputs = self(model_input)[_OUTPUT_TS]
|
||||
# (full batch, last patch, output_patch_len, index of mean forecast = 0)
|
||||
new_ts = fprop_outputs[:, -1, :output_patch_len, 0]
|
||||
# (full batch, last patch, output_patch_len, all output indices)
|
||||
full_outputs.append(fprop_outputs[:, -1, :output_patch_len, :])
|
||||
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, :],
|
||||
)
|
||||
return (
|
||||
final_out[:, inp_time_len:inp_time_len + horizon_len],
|
||||
jnp.concatenate(full_outputs, axis=1)[:, 0:horizon_len, :],
|
||||
)
|
||||
|
||||
|
||||
class PatchedDecoderFinetuneModel(base_model.BaseModel):
|
||||
"""Model class for finetuning patched time-series decoder.
|
||||
"""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
|
||||
core_layer_tpl: LayerTpl = template_field(PatchedTimeSeriesDecoder)
|
||||
freq: int = 0
|
||||
|
||||
def setup(self) -> None:
|
||||
self.create_child("core_layer", self.core_layer_tpl)
|
||||
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
|
||||
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)
|
||||
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.
|
||||
def _quantile_loss(self, pred: JTensor, actual: JTensor,
|
||||
quantile: float) -> JTensor:
|
||||
"""Calculates quantile loss.
|
||||
|
||||
Args:
|
||||
pred: B x T
|
||||
@@ -501,21 +497,20 @@ class PatchedDecoderFinetuneModel(base_model.BaseModel):
|
||||
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)
|
||||
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
|
||||
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
|
||||
|
||||
@@ -11,7 +11,6 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Directory to extract time covariates.
|
||||
|
||||
Extract time covariates from datetime.
|
||||
@@ -36,7 +35,6 @@ from pandas.tseries.offsets import Easter
|
||||
from sklearn.preprocessing import StandardScaler
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
# 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
|
||||
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),
|
||||
)
|
||||
assert (
|
||||
len(holiday_date) != 0 # pylint: disable=g-explicit-length-test
|
||||
), f"No closest holiday for the date index {index} found."
|
||||
assert (len(holiday_date) != 0 # pylint: disable=g-explicit-length-test
|
||||
), f"No closest holiday for the date index {index} found."
|
||||
# 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.
|
||||
return (index - holiday_date[0]).days
|
||||
@@ -60,16 +57,19 @@ def _distance_to_holiday(holiday):
|
||||
return _distance_to_day
|
||||
|
||||
|
||||
EasterSunday = Holiday(
|
||||
"Easter Sunday", month=1, day=1, offset=[Easter(), Day(0)]
|
||||
)
|
||||
EasterSunday = Holiday("Easter Sunday",
|
||||
month=1,
|
||||
day=1,
|
||||
offset=[Easter(), Day(0)])
|
||||
NewYearsDay = Holiday("New Years Day", month=1, day=1)
|
||||
SuperBowl = Holiday(
|
||||
"Superbowl", month=2, day=1, offset=DateOffset(weekday=SU(1))
|
||||
)
|
||||
MothersDay = Holiday(
|
||||
"Mothers Day", month=5, day=1, offset=DateOffset(weekday=SU(2))
|
||||
)
|
||||
SuperBowl = Holiday("Superbowl",
|
||||
month=2,
|
||||
day=1,
|
||||
offset=DateOffset(weekday=SU(1)))
|
||||
MothersDay = Holiday("Mothers Day",
|
||||
month=5,
|
||||
day=1,
|
||||
offset=DateOffset(weekday=SU(2)))
|
||||
IndependenceDay = Holiday("Independence Day", month=7, day=4)
|
||||
ChristmasEve = Holiday("Christmas", month=12, day=24)
|
||||
ChristmasDay = Holiday("Christmas", month=12, day=25)
|
||||
|
||||
+110
-131
@@ -11,7 +11,6 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""TimesFM forecast API for inference."""
|
||||
|
||||
import logging
|
||||
@@ -52,23 +51,16 @@ def moving_average(arr, window_size):
|
||||
"""Calculates the moving average using NumPy's convolution function."""
|
||||
# Pad with zeros to handle initial window positions
|
||||
arr_padded = np.pad(arr, (window_size - 1, 0), "constant")
|
||||
smoothed_arr = (
|
||||
np.convolve(arr_padded, np.ones(window_size), "valid") / window_size
|
||||
)
|
||||
smoothed_arr = (np.convolve(arr_padded, np.ones(window_size), "valid") /
|
||||
window_size)
|
||||
return [smoothed_arr, arr - smoothed_arr]
|
||||
|
||||
|
||||
def freq_map(freq: str):
|
||||
"""Returns the frequency map for the given frequency string."""
|
||||
freq = str.upper(freq)
|
||||
if (
|
||||
freq.endswith("H")
|
||||
or freq.endswith("T")
|
||||
or freq.endswith("MIN")
|
||||
or freq.endswith("D")
|
||||
or freq.endswith("B")
|
||||
or freq.endswith("U")
|
||||
):
|
||||
if (freq.endswith("H") or freq.endswith("T") or freq.endswith("MIN") or
|
||||
freq.endswith("D") or freq.endswith("B") or freq.endswith("U")):
|
||||
return 0
|
||||
elif freq.endswith(("W", "M", "MS")):
|
||||
return 1
|
||||
@@ -179,9 +171,7 @@ class TimesFm:
|
||||
num_layers=num_layers,
|
||||
transformer_layer_params_tpl=pax_fiddle.Config(
|
||||
transformers.Transformer,
|
||||
ln_tpl=pax_fiddle.Config(
|
||||
normalizations.RmsNorm,
|
||||
),
|
||||
ln_tpl=pax_fiddle.Config(normalizations.RmsNorm,),
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -199,34 +189,38 @@ class TimesFm:
|
||||
|
||||
def _get_sample_inputs(self):
|
||||
return {
|
||||
"input_ts": jnp.zeros(
|
||||
(
|
||||
self.per_core_batch_size,
|
||||
self.context_len + self.output_patch_len,
|
||||
"input_ts":
|
||||
jnp.zeros(
|
||||
(
|
||||
self.per_core_batch_size,
|
||||
self.context_len + self.output_patch_len,
|
||||
),
|
||||
dtype=jnp.float32,
|
||||
),
|
||||
dtype=jnp.float32,
|
||||
),
|
||||
"input_padding": jnp.zeros(
|
||||
(
|
||||
self.per_core_batch_size,
|
||||
self.context_len + self.output_patch_len,
|
||||
"input_padding":
|
||||
jnp.zeros(
|
||||
(
|
||||
self.per_core_batch_size,
|
||||
self.context_len + self.output_patch_len,
|
||||
),
|
||||
dtype=jnp.float32,
|
||||
),
|
||||
dtype=jnp.float32,
|
||||
),
|
||||
"freq": jnp.zeros(
|
||||
(
|
||||
self.per_core_batch_size,
|
||||
1,
|
||||
"freq":
|
||||
jnp.zeros(
|
||||
(
|
||||
self.per_core_batch_size,
|
||||
1,
|
||||
),
|
||||
dtype=jnp.int32,
|
||||
),
|
||||
dtype=jnp.int32,
|
||||
),
|
||||
}
|
||||
|
||||
def load_from_checkpoint(
|
||||
self,
|
||||
checkpoint_path: Optional[str] = None,
|
||||
repo_id: str = "google/timesfm-1.0-200m",
|
||||
checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.FLAX,
|
||||
checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.
|
||||
FLAX,
|
||||
step: int | None = None,
|
||||
) -> None:
|
||||
"""Loads a checkpoint and compiles the decoder.
|
||||
@@ -246,8 +240,7 @@ class TimesFm:
|
||||
start_time = time.time()
|
||||
self._model = instantiate(self.model_p)
|
||||
var_weight_hparams = self._model.abstract_init_with_metadata(
|
||||
self._get_sample_inputs(), do_eval=True
|
||||
)
|
||||
self._get_sample_inputs(), do_eval=True)
|
||||
train_state_partition_specs = tasks_lib.create_state_partition_specs(
|
||||
var_weight_hparams,
|
||||
mesh_shape=self.mesh_shape,
|
||||
@@ -261,8 +254,7 @@ class TimesFm:
|
||||
learners=None,
|
||||
)
|
||||
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.
|
||||
self._logging(f"Restoring checkpoint from {checkpoint_path}.")
|
||||
@@ -275,12 +267,12 @@ class TimesFm:
|
||||
step=step,
|
||||
)
|
||||
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()
|
||||
|
||||
|
||||
def jit_decode(self):
|
||||
"""Jitting decoding function."""
|
||||
|
||||
# Initialize and jit the decode fn.
|
||||
def _decode(inputs):
|
||||
assert self._model is not None
|
||||
@@ -310,34 +302,36 @@ class TimesFm:
|
||||
with base_layer.JaxContext.new_context(hparams=self._eval_context):
|
||||
_ = self._pmapped_decode(
|
||||
NestedMap({
|
||||
"input_ts": jnp.zeros(
|
||||
(
|
||||
self.num_devices,
|
||||
self.per_core_batch_size,
|
||||
self.context_len,
|
||||
"input_ts":
|
||||
jnp.zeros(
|
||||
(
|
||||
self.num_devices,
|
||||
self.per_core_batch_size,
|
||||
self.context_len,
|
||||
),
|
||||
dtype=jnp.float32,
|
||||
),
|
||||
dtype=jnp.float32,
|
||||
),
|
||||
"input_padding": jnp.zeros(
|
||||
(
|
||||
self.num_devices,
|
||||
self.per_core_batch_size,
|
||||
self.context_len + self.horizon_len,
|
||||
"input_padding":
|
||||
jnp.zeros(
|
||||
(
|
||||
self.num_devices,
|
||||
self.per_core_batch_size,
|
||||
self.context_len + self.horizon_len,
|
||||
),
|
||||
dtype=jnp.float32,
|
||||
),
|
||||
dtype=jnp.float32,
|
||||
),
|
||||
"date_features": None,
|
||||
"freq": jnp.zeros(
|
||||
(self.num_devices, self.per_core_batch_size, 1),
|
||||
dtype=jnp.int32,
|
||||
),
|
||||
})
|
||||
)
|
||||
"date_features":
|
||||
None,
|
||||
"freq":
|
||||
jnp.zeros(
|
||||
(self.num_devices, self.per_core_batch_size, 1),
|
||||
dtype=jnp.int32,
|
||||
),
|
||||
}))
|
||||
self._logging(f"Jitted decoding in {time.time() - start_time:.2f} seconds.")
|
||||
|
||||
def _preprocess(
|
||||
self, inputs: Sequence[np.array], freq: Sequence[int]
|
||||
) -> tuple[np.array, np.array, int]:
|
||||
def _preprocess(self, inputs: Sequence[np.array],
|
||||
freq: Sequence[int]) -> tuple[np.array, np.array, int]:
|
||||
"""Formats and pads raw inputs to feed into the model.
|
||||
|
||||
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 = [], [], []
|
||||
|
||||
pmap_pad = (
|
||||
(len(inputs) - 1) // self.global_batch_size + 1
|
||||
) * self.global_batch_size - len(inputs)
|
||||
pmap_pad = ((len(inputs) - 1) // self.global_batch_size +
|
||||
1) * self.global_batch_size - len(inputs)
|
||||
|
||||
for i, ts in enumerate(inputs):
|
||||
input_len = ts.shape[0]
|
||||
padding = np.zeros(shape=(input_len + self.horizon_len,), dtype=float)
|
||||
if input_len < self.context_len:
|
||||
num_front_pad = self.context_len - input_len
|
||||
ts = np.concatenate(
|
||||
[np.zeros(shape=(num_front_pad,), dtype=float), ts], axis=0
|
||||
)
|
||||
ts = np.concatenate([np.zeros(shape=(num_front_pad,), dtype=float), ts],
|
||||
axis=0)
|
||||
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:
|
||||
ts = ts[-self.context_len :]
|
||||
padding = padding[-(self.context_len + self.horizon_len) :]
|
||||
ts = ts[-self.context_len:]
|
||||
padding = padding[-(self.context_len + self.horizon_len):]
|
||||
|
||||
input_ts.append(ts)
|
||||
input_padding.append(padding)
|
||||
@@ -425,8 +416,7 @@ class TimesFm:
|
||||
if not self._train_state or not self._model:
|
||||
raise ValueError(
|
||||
"Checkpoint not loaded. Call `load_from_checkpoint` before"
|
||||
" `forecast`."
|
||||
)
|
||||
" `forecast`.")
|
||||
if forecast_context_len is None:
|
||||
forecast_context_len = self.context_len
|
||||
inputs = [np.array(ts)[-forecast_context_len:] for ts in inputs]
|
||||
@@ -448,47 +438,45 @@ class TimesFm:
|
||||
full_outputs = []
|
||||
assert input_ts.shape[0] % self.global_batch_size == 0
|
||||
for i in range(input_ts.shape[0] // self.global_batch_size):
|
||||
input_ts_in = jnp.array(
|
||||
input_ts[
|
||||
i * self.global_batch_size : (i + 1) * self.global_batch_size
|
||||
]
|
||||
)
|
||||
input_ts_in = jnp.array(input_ts[i * self.global_batch_size:(i + 1) *
|
||||
self.global_batch_size])
|
||||
input_padding_in = jnp.array(
|
||||
input_padding[
|
||||
i * self.global_batch_size : (i + 1) * self.global_batch_size
|
||||
],
|
||||
)
|
||||
input_padding[i * self.global_batch_size:(i + 1) *
|
||||
self.global_batch_size],)
|
||||
inp_freq_in = jnp.array(
|
||||
inp_freq[
|
||||
i * self.global_batch_size : (i + 1) * self.global_batch_size, :
|
||||
],
|
||||
inp_freq[i * self.global_batch_size:(i + 1) *
|
||||
self.global_batch_size, :],
|
||||
dtype=jnp.int32,
|
||||
)
|
||||
pmapped_inputs = NestedMap({
|
||||
"input_ts": es.jax_einshape(
|
||||
"(db)...->db...",
|
||||
input_ts_in,
|
||||
d=self.num_devices,
|
||||
),
|
||||
"input_padding": es.jax_einshape(
|
||||
"(db)...->db...",
|
||||
input_padding_in,
|
||||
d=self.num_devices,
|
||||
),
|
||||
"date_features": None,
|
||||
"freq": es.jax_einshape(
|
||||
"(db)...->db...",
|
||||
inp_freq_in,
|
||||
d=self.num_devices,
|
||||
),
|
||||
"input_ts":
|
||||
es.jax_einshape(
|
||||
"(db)...->db...",
|
||||
input_ts_in,
|
||||
d=self.num_devices,
|
||||
),
|
||||
"input_padding":
|
||||
es.jax_einshape(
|
||||
"(db)...->db...",
|
||||
input_padding_in,
|
||||
d=self.num_devices,
|
||||
),
|
||||
"date_features":
|
||||
None,
|
||||
"freq":
|
||||
es.jax_einshape(
|
||||
"(db)...->db...",
|
||||
inp_freq_in,
|
||||
d=self.num_devices,
|
||||
),
|
||||
})
|
||||
mean_output, full_output = self._pmapped_decode(pmapped_inputs)
|
||||
mean_output = es.jax_einshape(
|
||||
"db...->(db)...", mean_output, d=self.num_devices
|
||||
)
|
||||
full_output = es.jax_einshape(
|
||||
"db...->(db)...", full_output, d=self.num_devices
|
||||
)
|
||||
mean_output = es.jax_einshape("db...->(db)...",
|
||||
mean_output,
|
||||
d=self.num_devices)
|
||||
full_output = es.jax_einshape("db...->(db)...",
|
||||
full_output,
|
||||
d=self.num_devices)
|
||||
mean_output = np.array(mean_output)
|
||||
full_output = np.array(full_output)
|
||||
mean_outputs.append(mean_output)
|
||||
@@ -539,14 +527,10 @@ class TimesFm:
|
||||
Returns:
|
||||
Future forecasts dataframe.
|
||||
"""
|
||||
if not (
|
||||
"unique_id" in inputs.columns
|
||||
and "ds" in inputs.columns
|
||||
and value_name in inputs.columns
|
||||
):
|
||||
if not ("unique_id" in inputs.columns and "ds" in inputs.columns and
|
||||
value_name in inputs.columns):
|
||||
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:
|
||||
forecast_context_len = self.context_len
|
||||
logging.info("Preprocessing dataframe.")
|
||||
@@ -571,17 +555,15 @@ class TimesFm:
|
||||
with multiprocessing.Pool(processes=num_jobs) as pool:
|
||||
results = pool.starmap(
|
||||
process_group,
|
||||
[
|
||||
(key, group, value_name, forecast_context_len)
|
||||
for key, group in df_sorted.groupby("unique_id")
|
||||
],
|
||||
[(key, group, value_name, forecast_context_len)
|
||||
for key, group in df_sorted.groupby("unique_id")],
|
||||
)
|
||||
new_inputs, uids = zip(*results)
|
||||
print("Finished preprocessing dataframe.")
|
||||
freq_inps = [freq_map(freq)] * len(new_inputs)
|
||||
_, full_forecast = self.forecast(
|
||||
new_inputs, freq=freq_inps, window_size=window_size
|
||||
)
|
||||
_, full_forecast = self.forecast(new_inputs,
|
||||
freq=freq_inps,
|
||||
window_size=window_size)
|
||||
print("Finished forecasting.")
|
||||
fcst_df = make_future_dataframe(
|
||||
uids=uids,
|
||||
@@ -589,16 +571,13 @@ class TimesFm:
|
||||
h=self.horizon_len,
|
||||
freq=freq,
|
||||
)
|
||||
fcst_df[model_name] = full_forecast[:, 0 : self.horizon_len, 0].reshape(
|
||||
-1, 1
|
||||
)
|
||||
fcst_df[model_name] = full_forecast[:, 0:self.horizon_len, 0].reshape(-1, 1)
|
||||
|
||||
if self._model.quantiles is not None:
|
||||
for i, q in enumerate(self._model.quantiles):
|
||||
q_col = f"{model_name}-q-{q}"
|
||||
fcst_df[q_col] = full_forecast[:, 0 : self.horizon_len, 1 + i].reshape(
|
||||
-1, 1
|
||||
)
|
||||
fcst_df[q_col] = full_forecast[:, 0:self.horizon_len,
|
||||
1 + i].reshape(-1, 1)
|
||||
if q == 0.5:
|
||||
fcst_df[model_name] = fcst_df[q_col]
|
||||
logging.info("Finished creating output dataframe.")
|
||||
|
||||
Reference in New Issue
Block a user