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

Adding v2.0 support

Browse Source
This commit is contained in:
Rajat Sen
2024-12-30 23:55:36 +00:00
parent 5a69171296
commit 73704e5936
12 changed files with 1467 additions and 1410 deletions
Download Patch File Download Diff File Expand all files Collapse all files
notebooks/covariates.ipynb
+33 -53
View File
@@ -9,7 +9,8 @@
"This toturial notebook demonstrates how to utilize exogenous covariates with TimesFM when making forecasts. Before running this notebook, make sure:\n",
"\n",
"- You've read through the README of TimesFM.\n",
"- A local kernel with Python 3.10 is up and running."
"- A local kernel with Python 3.10 is up and running, for the jax version.\n",
"- Install the JAX version following the installation instructions."
]
},
{
@@ -19,36 +20,15 @@
"## Setup the environment and install TimesFM."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'\n",
"os.environ['JAX_PMAP_USE_TENSORSTORE'] = 'false'"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!pip install timesfm\n",
"import timesfm"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load the checkpoint\n",
"\n",
"**Notice:** Please set up the backend as per your machine (\"cpu\", \"gpu\" or \"tpu\"). This notebook will run by default on CPU.\n",
"**Notice:** Please set up the backend as per your machine (\"cpu\", \"gpu\" or \"tpu\"). This notebook will run by default on GPU.\n",
"\n",
"We load the 1.0-200m model checkpoint from HuggingFace."
"We load the 2.0-500m model checkpoint from HuggingFace."
]
},
{
@@ -57,23 +37,21 @@
"metadata": {},
"outputs": [],
"source": [
"timesfm_backend = \"cpu\" # @param\n",
"\n",
"from jax._src import config\n",
"config.update(\n",
" \"jax_platforms\", {\"cpu\": \"cpu\", \"gpu\": \"cuda\", \"tpu\": \"\"}[timesfm_backend]\n",
")\n",
"import timesfm\n",
"timesfm_backend = \"gpu\" # @param\n",
"\n",
"model = timesfm.TimesFm(\n",
" context_len=512,\n",
" horizon_len=128,\n",
" input_patch_len=32,\n",
" output_patch_len=128,\n",
" num_layers=20,\n",
" model_dims=1280,\n",
" backend=timesfm_backend,\n",
")\n",
"model.load_from_checkpoint(repo_id=\"google/timesfm-1.0-200m\")"
" hparams=timesfm.TimesFmHparams(\n",
" backend=timesfm_backend,\n",
" per_core_batch_size=32,\n",
" horizon_len=128,\n",
" num_layers=50,\n",
" use_positional_embedding=False,\n",
" context_len=2048,\n",
" ),\n",
" checkpoint=timesfm.TimesFmCheckpoint(\n",
" huggingface_repo_id=\"google/timesfm-2.0-500m-jax\"),\n",
" )"
]
},
{
@@ -140,7 +118,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
@@ -176,7 +154,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
@@ -208,7 +186,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
@@ -239,13 +217,15 @@
"metadata": {},
"outputs": [],
"source": [
"import time\n",
"\n",
"# Benchmark\n",
"batch_size = 128\n",
"context_len = 120\n",
"horizon_len = 24\n",
"input_data = get_batched_data_fn(batch_size = 128)\n",
"metrics = defaultdict(list)\n",
"import time\n",
"\n",
"\n",
"for i, example in enumerate(input_data()):\n",
" raw_forecast, _ = model.forecast(\n",
@@ -307,15 +287,15 @@
"source": [
"You should see results close to \n",
"```\n",
"eval_mae_timesfm: 6.762283045916956\n",
"eval_mae_xreg_timesfm: 5.39219617611074\n",
"eval_mae_xreg: 37.15275842572484\n",
"eval_mse_timesfm: 166.7771466306823\n",
"eval_mse_xreg_timesfm: 120.64757721021306\n",
"eval_mse_xreg: 1672.2116821201796\n",
"eval_mae_timesfm: 6.729583250571446\n",
"eval_mae_xreg_timesfm: 5.3375301110158\n",
"eval_mae_xreg: 37.152760709266\n",
"eval_mse_timesfm: 162.3132151851567\n",
"eval_mse_xreg_timesfm: 120.9900627409689\n",
"eval_mse_xreg: 1672.208769045399\n",
"```\n",
"\n",
"With the covariates, the TimesFM forecast Mean Absolute Error improves from 6.76 to 5.39, and Mean Squred Error from 166.78 to 120.65. The results of purely fitting the linear model are also provided for reference."
"With the covariates, the TimesFM forecast Mean Absolute Error improves from 6.73 to 5.34, and Mean Squred Error from 162.31 to 120.99. The results of purely fitting the linear model are also provided for reference."
]
},
{
@@ -381,9 +361,9 @@
],
"metadata": {
"kernelspec": {
"display_name": "cuda-gpt",
"display_name": "chronos-v2",
"language": "python",
"name": "cuda"
"name": "python3"
},
"language_info": {
"codemirror_mode": {
@@ -395,7 +375,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.14"
"version": "3.10.15"
}
},
"nbformat": 4,
notebooks/finetuning.ipynb
+28 -20
View File
@@ -9,7 +9,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
@@ -34,7 +34,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
@@ -61,14 +61,22 @@
"metadata": {},
"outputs": [],
"source": [
"timesfm_backend = \"gpu\" # @param\n",
"\n",
"tfm = timesfm.TimesFm(\n",
" hparams=timesfm.TimesFmHparams(\n",
" backend=\"gpu\",\n",
" backend=timesfm_backend,\n",
" per_core_batch_size=32,\n",
" horizon_len=128,\n",
" num_layers=50,\n",
" # Se this to True for v1.0 checkpoints\n",
" use_positional_embedding=False,\n",
" # Note that we could set this to as high as 2048 but keeping it 512 here so that\n",
" # both v1.0 and 2.0 checkpoints work\n",
" context_len=512,\n",
" ),\n",
" checkpoint=timesfm.TimesFmCheckpoint(\n",
" huggingface_repo_id=\"google/timesfm-1.0-200m\"),\n",
" huggingface_repo_id=\"google/timesfm-2.0-500m-jax\"),\n",
" )"
]
},
@@ -81,7 +89,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
@@ -126,7 +134,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
@@ -147,7 +155,7 @@
"pred_len = 96\n",
"\n",
"num_ts = len(ts_cols)\n",
"batch_size = 16\n",
"batch_size = 8\n",
"\n",
"dtl = data_loader.TimeSeriesdata(\n",
" data_path=data_path,\n",
@@ -187,7 +195,7 @@
"outputs": [],
"source": [
"for tbatch in tqdm(train_batches.as_numpy_iterator()):\n",
" pass\n",
" break\n",
"print(tbatch[0].shape)"
]
},
@@ -224,7 +232,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 9,
"metadata": {},
"outputs": [],
"source": [
@@ -248,7 +256,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
@@ -276,7 +284,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 11,
"metadata": {},
"outputs": [],
"source": [
@@ -296,7 +304,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
@@ -325,7 +333,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
@@ -416,7 +424,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 19,
"metadata": {},
"outputs": [],
"source": [
@@ -445,7 +453,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 20,
"metadata": {},
"outputs": [],
"source": [
@@ -455,7 +463,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 21,
"metadata": {},
"outputs": [],
"source": [
@@ -470,7 +478,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 22,
"metadata": {},
"outputs": [],
"source": [
@@ -583,13 +591,13 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"## There is around a __9%__ reduction in MAE from finetuning."
"## There is around a __7%__ reduction in MAE from finetuning."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "tfm_env_v3",
"display_name": "chronos-v2",
"language": "python",
"name": "python3"
},
@@ -603,7 +611,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.14"
"version": "3.10.15"
}
},
"nbformat": 4,