No public description

PiperOrigin-RevId: 631148246 No public description PiperOrigin-RevId: 629465861 No public description PiperOrigin-RevId: 629538727 No public description PiperOrigin-RevId: 629803430 No public description PiperOrigin-RevId: 629838772 No public description PiperOrigin-RevId: 629875841 No public description PiperOrigin-RevId: 630194303 No public description PiperOrigin-RevId: 630212078 No public description PiperOrigin-RevId: 630215811 No public description PiperOrigin-RevId: 630402195 No public description PiperOrigin-RevId: 630452190 No public description PiperOrigin-RevId: 630511573 No public description PiperOrigin-RevId: 630706657 No public description PiperOrigin-RevId: 630712978 No public description PiperOrigin-RevId: 630735187
2024-04-30 10:58:13 -07:00
commit 81dc60c086
19 changed files with 2776 additions and 0 deletions
@@ -0,0 +1,146 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Evaluation script for timesfm."""
+
+import os
+import sys
+import time
+
+from absl import flags
+import numpy as np
+import pandas as pd
+from paxml import checkpoints
+import timesfm
+
+from .utils import ExperimentHandler
+
+
+dataset_names = [
+    "m1_monthly",
+    "m1_quarterly",
+    "m1_yearly",
+    "m3_monthly",
+    "m3_other",
+    "m3_quarterly",
+    "m3_yearly",
+    "m4_quarterly",
+    "m4_yearly",
+    "tourism_monthly",
+    "tourism_quarterly",
+    "tourism_yearly",
+    "nn5_daily_without_missing",
+    "m5",
+    "nn5_weekly",
+    "traffic",
+    "weather",
+    "dominick",
+    "australian_electricity_demand",
+    "car_parts_without_missing",
+    "cif_2016",
+    "covid_deaths",
+    "ercot",
+    "ett_small_15min",
+    "ett_small_1h",
+    "exchange_rate",
+    "fred_md",
+    "hospital",
+]
+
+context_dict = {
+    "cif_2016": 32,
+    "tourism_yearly": 64,
+    "covid_deaths": 64,
+    "tourism_quarterly": 64,
+    "tourism_monthly": 64,
+    "m1_monthly": 64,
+    "m1_quarterly": 64,
+    "m1_yearly": 64,
+    "m3_monthly": 64,
+    "m3_other": 64,
+    "m3_quarterly": 64,
+    "m3_yearly": 64,
+    "m4_quarterly": 64,
+    "m4_yearly": 64,
+}
+
+_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)
+
+
+def main():
+  results_list = []
+  tfm = timesfm.TimesFm(
+      context_len=512,
+      horizon_len=_HORIZON.value,
+      input_patch_len=32,
+      output_patch_len=128,
+      num_layers=20,
+      model_dims=1280,
+      backend=_BACKEND.value,
+      per_core_batch_size=_BATCH_SIZE.value,
+      quantiles=QUANTILES,
+  )
+  tfm.load_from_checkpoint(
+      _MODEL_PATH.value,
+      checkpoint_type=checkpoints.CheckpointType.FLAX,
+  )
+  run_id = np.random.randint(100000)
+  model_name = "timesfm"
+  for dataset in dataset_names:
+    print(f"Evaluating model {model_name} on dataset {dataset}", flush=True)
+    exp = ExperimentHandler(dataset, quantiles=QUANTILES)
+
+    if dataset in context_dict:
+      context_len = context_dict[dataset]
+    else:
+      context_len = 512
+    train_df = exp.train_df
+    freq = exp.freq
+    init_time = time.time()
+    fcsts_df = tfm.forecast_on_df(
+        inputs=train_df,
+        freq=freq,
+        value_name="y",
+        model_name=model_name,
+        forecast_context_len=context_len,
+        num_jobs=_NUM_JOBS.value,
+    )
+    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
+    )
+    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()