Full pytorch support

This commit is contained in:
Rajat Sen
2024-09-12 23:30:46 +00:00
parent 61fa1b2ef2
commit 1b95563eea
19 changed files with 2402 additions and 1935 deletions
-28
View File
@@ -1,28 +0,0 @@
name: tfm_env
channels:
- conda-forge
- defaults
- anaconda
dependencies:
- jupyterlab
- pip
- python=3.10
- pip:
- datasetsforecast
- fire
- git+https://github.com/awslabs/gluon-ts.git
- huggingface_hub[cli]
- neuralforecast
- orjson
- statsforecast
- utilsforecast
- git+https://github.com/amazon-science/chronos-forecasting.git
- praxis
- paxml
- jax[cuda12]==0.4.26
- einshape
- python-dotenv
- nixtla>=0.5.1
- rich
- scikit-learn
-28
View File
@@ -1,28 +0,0 @@
name: tfm_env
channels:
- conda-forge
- defaults
- anaconda
dependencies:
- jupyterlab
- pip
- python=3.10
- pip:
- datasetsforecast
- fire
- git+https://github.com/awslabs/gluon-ts.git
- huggingface_hub[cli]
- neuralforecast
- orjson
- statsforecast
- utilsforecast
- git+https://github.com/amazon-science/chronos-forecasting.git
- praxis
- paxml
- jax[cpu]==0.4.26
- einshape
- python-dotenv
- nixtla>=0.5.1
- rich
- scikit-learn
+11 -5
View File
@@ -5,14 +5,20 @@ The benchmark setting has been borrowed from Nixtla's original [benchmarking](ht
## Running TimesFM on the benchmark
Install the environment and the package as detailed in the main README and then follow the steps from the base directory.
We need to add the following packages for running these benchmarks. Follow the installation instructions till before `poetry lock`. Then,
```
conda activate tfm_env
TF_CPP_MIN_LOG_LEVEL=2 XLA_PYTHON_CLIENT_PREALLOCATE=false python3 -m experiments.extended_benchmarks.run_timesfm --model_path=<model_path> --backend="gpu"
poetry add git+https://github.com/awslabs/gluon-ts.git
poetry lock
poetry install --only <pax or pytorch>
```
To run the timesfm on the benchmark do:
```
poetry run python3 -m experiments.extended_benchmarks.run_timesfm --model_path=google/timesfm-1.0-200m(-pytorch) --backend="gpu"
```
In the above, `<model_path>` should point to the checkpoint directory that can be downloaded from HuggingFace.
Note: In the current version of TimesFM we focus on point forecasts and therefore the mase, smape have been calculated using the quantile head corresponding to the median i.e 0.5 quantile. We do offer 10 quantile heads but they have not been calibrated after pretraining. We recommend using them with caution or calibrate/conformalize them on a hold out for your applications. More to follow on later versions.
@@ -22,7 +28,7 @@ Note: In the current version of TimesFM we focus on point forecasts and therefor
__Update:__ We have added TimeGPT-1 to the benchmark results. We had to remove the Dominick dataset as we were not able to run TimeGPT-1 on this benchmark. Note that the previous results including Dominick remain available at `./tfm_results.png`. In order to reproduce the results for TimeGPT-1, please run `run_timegpt.py`.
_Remark:_ All baselines except the ones involving TimeGPT were run performed on a [g2-standard-32](https://cloud.google.com/compute/docs/gpus). Since TimeGPT-1 can only be accessed by an API, the time column might not reflect the true speed of the model as it also includes the communication cost. Moreover, we are not sure about the exact backend hardware for TimeGPT.
_Remark:_ All baselines except the ones involving TimeGPT were run performed on a [g2-standard-32](https://cloud.google.com/compute/docs/gpus). Since TimeGPT-1 can only be accessed by an API, the time column might not reflect the true speed of the model as it also includes the communication cost. Moreover, we are not sure about the exact backend hardware for TimeGPT. The TimesFM latency numbers are from the PAX version.
We can see that TimesFM performs the best in terms of both mase and smape. More importantly it is much faster than the other methods, in particular it is more than 600x faster than StatisticalEnsemble and 80x faster than Chronos (Large).
+12 -23
View File
@@ -11,7 +11,6 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Evaluation script for timesfm."""
import os
@@ -21,12 +20,10 @@ 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",
@@ -74,35 +71,27 @@ 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", "google/timesfm-1.0-200m",
"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,
hparams=timesfm.TimesFmHparams(
backend=_BACKEND.value,
per_core_batch_size=_BATCH_SIZE.value,
horizon_len=_HORIZON.value,
),
checkpoint=timesfm.TimesFmCheckpoint(
huggingface_repo_id=_MODEL_PATH.value),
)
run_id = np.random.randint(100000)
model_name = "timesfm"
@@ -127,9 +116,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)
+14 -5
View File
@@ -6,12 +6,21 @@ All experiments were performed on a [g2-standard-32](https://cloud.google.com/co
## Running TimesFM on the benchmark
Install the environment and the package as detailed in the main README and then follow the steps from the base directory.
We need to add the following packages for running these benchmarks. Follow the installation instructions till before `poetry lock`. Then,
```
conda activate tfm_env
TF_CPP_MIN_LOG_LEVEL=2 XLA_PYTHON_CLIENT_PREALLOCATE=false python3 -m experiments.long_horizon_benchmarks.run_eval \
--model_path=<model_path> --backend="gpu" \
poetry add git+https://github.com/awslabs/gluon-ts.git
poetry add git+https://github.com/amazon-science/chronos-forecasting.git
poetry lock
poetry install --only pax
```
Note that for now only the pax version runs on this benchmark, because we had to remove the old tf dependency from the pytorch version. We will fix this issue soon.
To run the timesfm on the benchmark do:
```
poetry run python3 -m experiments.long_horizon_benchmarks.run_eval \
--model_path=google/timesfm-1.0-200m --backend="gpu" \
--pred_len=96 --context_len=512 --dataset=etth1
```
@@ -20,7 +29,7 @@ In the above, `<model_path>` should point to the checkpoint directory that can b
For running chronos on the same benchmark you can run the command,
```
TF_CPP_MIN_LOG_LEVEL=2 XLA_PYTHON_CLIENT_PREALLOCATE=false python3 -m experiments.long_horizon_benchmarks.run_eval \
poetry run python3 -m experiments.long_horizon_benchmarks.run_eval \
--model_path=amazon/chronos-t5-mini --backend="gpu" \
--pred_len=96 --context_len=512 --dataset=etth1
```
+27 -48
View File
@@ -11,7 +11,6 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Eval pipeline."""
import json
@@ -22,44 +21,33 @@ from absl import flags
import chronos
import numpy as np
import pandas as pd
from paxml import checkpoints
import timesfm
from timesfm import data_loader
import torch
import tqdm
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 +164,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(
@@ -188,19 +175,12 @@ def eval():
)
else:
model = timesfm.TimesFm(
context_len=_CONTEXT_LEN.value,
horizon_len=_PRED_LEN.value,
input_patch_len=32,
output_patch_len=128,
num_layers=20,
model_dims=1280,
backend=_BACKEND.value,
per_core_batch_size=batch_size,
quantiles=QUANTILES,
)
model.load_from_checkpoint(
model_path,
checkpoint_type=checkpoints.CheckpointType.FLAX,
hparams=timesfm.TimesFmHparams(
backend=_BACKEND.value,
per_core_batch_size=_BATCH_SIZE.value,
horizon_len=_PRED_LEN.value,
),
checkpoint=timesfm.TimesFmCheckpoint(huggingface_repo_id=model_path),
)
smape_run_losses = []
mse_run_losses = []
@@ -213,10 +193,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())