2.0.0 initial

This commit is contained in:
siriuz42
2025-09-12 00:18:08 +00:00
parent d70708d42a
commit 7d8f3d971d
52 changed files with 1882 additions and 394 deletions
+383
View File
@@ -0,0 +1,383 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# TimesFM with Covariates\n",
"\n",
"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, for the jax version.\n",
"- Install the JAX version following the installation instructions."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setup the environment and install 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 GPU.\n",
"\n",
"We load the 2.0-500m model checkpoint from HuggingFace."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import timesfm\n",
"timesfm_backend = \"gpu\" # @param\n",
"\n",
"model = timesfm.TimesFm(\n",
" 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",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Covariates\n",
"\n",
"Let's take a toy example of forecasting sales for a grocery store: \n",
"\n",
"**Task:** Given the observed the daily sales of this week (7 days), forecast the daily sales of next week (7 days).\n",
"\n",
"```\n",
"Product: ice cream\n",
"Daily_sales: [30, 30, 4, 5, 7, 8, 10]\n",
"Category: food\n",
"Base_price: 1.99\n",
"Weekday: [0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6]\n",
"Has_promotion: [Yes, Yes, No, No, No, Yes, Yes, No, No, No, No, No, No, No]\n",
"Daily_temperature: [31.0, 24.3, 19.4, 26.2, 24.6, 30.0, 31.1, 32.4, 30.9, 26.0, 25.0, 27.8, 29.5, 31.2]\n",
"```\n",
"\n",
"```\n",
"Product: sunscreen\n",
"Daily_sales: [5, 7, 12, 13, 5, 6, 10]\n",
"Category: skin product\n",
"Base_price: 29.99\n",
"Weekday: [0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6]\n",
"Has_promotion: [No, No, Yes, Yes, No, No, No, Yes, Yes, Yes, Yes, Yes, Yes, Yes]\n",
"Daily_temperature: [31.0, 24.3, 19.4, 26.2, 24.6, 30.0, 31.1, 32.4, 30.9, 26.0, 25.0, 27.8, 29.5, 31.2]\n",
"```\n",
"\n",
"In this example, besides the `Daily_sales`, we also have covariates `Category`, `Base_price`, `Weekday`, `Has_promotion`, `Daily_temperature`. Let's introduce some concepts:\n",
"\n",
"**Static covariates** are covariates for each time series. \n",
"- In our example, `Category` is a **static categorical covariate**, \n",
"- `Base_price` is a **static numerical covariates**.\n",
"\n",
"**Dynamic covariates** are covaraites for each time stamps.\n",
"- Date / time related features can be usually treated as dynamic covariates.\n",
"- In our example, `Weekday` and `Has_promotion` are **dynamic categorical covariates**.\n",
"- `Daily_temperate` is a **dynamic numerical covariate**.\n",
"\n",
"**Notice:** Here we make it mandatory that the dynamic covariates need to cover both the forecasting context and horizon. For example, all dynamic covariates in the example have 14 values: the first 7 correspond to the observed 7 days, and the last 7 correspond to the next 7 days."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# TimesFM with Covariates\n",
"\n",
"\n",
"The strategy we take here is to treat covariates as batched in-context exogenous regressors (XReg) and fit linear models on them outside of TimesFM. The final forecast will be the sum of the TimesFM forecast and the linear model forecast.\n",
"\n",
" In simple words, we consider these two options.\n",
"\n",
"**Option 1:** Get the TimesFM forecast, and fit the linear model regressing the residuals on the covariates (\"timesfm + xreg\").\n",
"\n",
"**Option 2:** Fit the linear model of the time series itself on the covariates, then forecast the residuals using TimesFM (\"xreg + timesfm\").\n",
"\n",
"Let's take a code at the example of Electricity Price Forecasting (EPF). \n"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"import pandas as pd\n",
"import numpy as np\n",
"from collections import defaultdict"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"df = pd.read_csv('https://datasets-nixtla.s3.amazonaws.com/EPF_FR_BE.csv')\n",
"df['ds'] = pd.to_datetime(df['ds'])\n",
"df"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This dataset has a few covariates beside the hourly target `y`:\n",
"\n",
"- `unique_id`: a static categorical covariate indicating the country.\n",
"- `gen_forecast`: a dynamic numerical covariate indicating the estimated electricity to be generated.\n",
"- `system_load`: the observed system load. Notice that this **CANNOT** be considered as a dynamic numerical covariate because we cannot know its values over the forecasting horizon in advance.\n",
"- `weekday`: a dynamic categorical covariate.\\\n",
"\n",
"Let's now make some forecasting tasks for TimesFM based on this dataset. For simplicity we create forecast contexts of 120 time points (hours) and forecast horizons of 24 time points."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"# Data pipelining\n",
"def get_batched_data_fn(\n",
" batch_size: int = 128, \n",
" context_len: int = 120, \n",
" horizon_len: int = 24,\n",
"):\n",
" examples = defaultdict(list)\n",
"\n",
" num_examples = 0\n",
" for country in (\"FR\", \"BE\"):\n",
" sub_df = df[df[\"unique_id\"] == country]\n",
" for start in range(0, len(sub_df) - (context_len + horizon_len), horizon_len):\n",
" num_examples += 1\n",
" examples[\"country\"].append(country)\n",
" examples[\"inputs\"].append(sub_df[\"y\"][start:(context_end := start + context_len)].tolist())\n",
" examples[\"gen_forecast\"].append(sub_df[\"gen_forecast\"][start:context_end + horizon_len].tolist())\n",
" examples[\"week_day\"].append(sub_df[\"week_day\"][start:context_end + horizon_len].tolist())\n",
" examples[\"outputs\"].append(sub_df[\"y\"][context_end:(context_end + horizon_len)].tolist())\n",
" \n",
" def data_fn():\n",
" for i in range(1 + (num_examples - 1) // batch_size):\n",
" yield {k: v[(i * batch_size) : ((i + 1) * batch_size)] for k, v in examples.items()}\n",
" \n",
" return data_fn\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"# Define metrics\n",
"def mse(y_pred, y_true):\n",
" y_pred = np.array(y_pred)\n",
" y_true = np.array(y_true)\n",
" return np.mean(np.square(y_pred - y_true), axis=1, keepdims=True)\n",
"\n",
"def mae(y_pred, y_true):\n",
" y_pred = np.array(y_pred)\n",
" y_true = np.array(y_true)\n",
" return np.mean(np.abs(y_pred - y_true), axis=1, keepdims=True)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now let's try `model.forecast_with_covariates`. \n",
"\n",
"In particular, the output is a tuple whose first element is the new forecast."
]
},
{
"cell_type": "code",
"execution_count": null,
"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",
"\n",
"\n",
"for i, example in enumerate(input_data()):\n",
" raw_forecast, _ = model.forecast(\n",
" inputs=example[\"inputs\"], freq=[0] * len(example[\"inputs\"])\n",
" )\n",
" start_time = time.time()\n",
" # Forecast with covariates\n",
" # Output: new forecast, forecast by the xreg\n",
" cov_forecast, ols_forecast = model.forecast_with_covariates( \n",
" inputs=example[\"inputs\"],\n",
" dynamic_numerical_covariates={\n",
" \"gen_forecast\": example[\"gen_forecast\"],\n",
" },\n",
" dynamic_categorical_covariates={\n",
" \"week_day\": example[\"week_day\"],\n",
" },\n",
" static_numerical_covariates={},\n",
" static_categorical_covariates={\n",
" \"country\": example[\"country\"]\n",
" },\n",
" freq=[0] * len(example[\"inputs\"]),\n",
" xreg_mode=\"xreg + timesfm\", # default\n",
" ridge=0.0,\n",
" force_on_cpu=False,\n",
" normalize_xreg_target_per_input=True, # default\n",
" )\n",
" print(\n",
" f\"\\rFinished batch {i} linear in {time.time() - start_time} seconds\",\n",
" end=\"\",\n",
" )\n",
" metrics[\"eval_mae_timesfm\"].extend(\n",
" mae(raw_forecast[:, :horizon_len], example[\"outputs\"])\n",
" )\n",
" metrics[\"eval_mae_xreg_timesfm\"].extend(mae(cov_forecast, example[\"outputs\"]))\n",
" metrics[\"eval_mae_xreg\"].extend(mae(ols_forecast, example[\"outputs\"]))\n",
" metrics[\"eval_mse_timesfm\"].extend(\n",
" mse(raw_forecast[:, :horizon_len], example[\"outputs\"])\n",
" )\n",
" metrics[\"eval_mse_xreg_timesfm\"].extend(mse(cov_forecast, example[\"outputs\"]))\n",
" metrics[\"eval_mse_xreg\"].extend(mse(ols_forecast, example[\"outputs\"]))\n",
"\n",
"print()\n",
"\n",
"for k, v in metrics.items():\n",
" print(f\"{k}: {np.mean(v)}\")\n",
"\n",
"# My output:\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"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You should see results close to \n",
"```\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.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."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Formatting Your Request\n",
"\n",
"It is quite crucial to get the covariates properly formatted so that we can call this `model.forecast_with_covariates`. Please see its docstring for details. Here let's also grab a batch from a toy data input pipeline for quick explanations."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"toy_input_pipeline = get_batched_data_fn(batch_size=2, context_len=5, horizon_len=2)\n",
"print(next(toy_input_pipeline()))\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You should see something similar to this\n",
"```\n",
"{\n",
" 'country': ['FR', 'FR'], \n",
" 'inputs': [[53.48, 51.93, 48.76, 42.27, 38.41], [48.76, 42.27, 38.41, 35.72, 32.66]], \n",
" 'gen_forecast': [[76905.0, 75492.0, 74394.0, 72639.0, 69347.0, 67960.0, 67564.0], [74394.0, 72639.0, 69347.0, 67960.0, 67564.0, 67277.0, 67019.0]], \n",
" 'week_day': [[3, 3, 3, 3, 3, 3, 3], [3, 3, 3, 3, 3, 3, 3]], \n",
" 'outputs': [[35.72, 32.66], [32.83, 30.06]],\n",
"}\n",
"```\n",
"\n",
"Notice:\n",
"- We have two examples in this batch.\n",
"- For each example we support different context lengths and horizon lengths just as `model.forecast`. Although it is not demonstrated in this dataset.\n",
"- If dynamic covariates are present, the horizon lengths will be inferred from them, e.g. how many values are provided in additional to the ones corresponding to the inputs. Make sure all your dynamic covariates have the same length per example.\n",
"- The static covariates are one per example.\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## More Applications\n",
"\n",
"### Past Dynamic Covariates\n",
"\n",
"Past dynamic covariates are covariates that are only available for the context. For instance in our example `system_load` is a past dynamic covariate. Time series models generally can handle this, however it is something the batched in context regression cannot address, because these regressors are not available in the future. If you do have those covariates and consider them very meaningful, there are two hacky options to try immediately:\n",
"\n",
"1. Shift and repeat these past dynamic covariates to use their delayed version. For example, if you think the `system_load` for this week is meaningful for forecasting next week, you can create a `delay_7_system_load` by shifting 7 timestamps and use this as one dynamic numerical covariate for TimesFM.\n",
"2. Bootstrap, that is to run TimesFM once to forecast these past dynamic covariates into the horizon, then call TimesFM again using these forecasts as the future part for these dynamic covariates.\n",
"\n",
"### Multivariate Time Series\n",
"\n",
"For multivariate time series, if we need univariate forecast, we can try treating the main time series as the target and use the rest as the dynamic covariates."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "chronos-v2",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.15"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
+619
View File
@@ -0,0 +1,619 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Importing relevant packages for finetuning"
]
},
{
"cell_type": "code",
"execution_count": 1,
"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": [
"import timesfm\n",
"import gc\n",
"import numpy as np\n",
"import pandas as pd\n",
"from timesfm import patched_decoder\n",
"from timesfm import data_loader"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"from tqdm import tqdm\n",
"import dataclasses\n",
"import IPython\n",
"import IPython.display\n",
"import matplotlib as mpl\n",
"import matplotlib.pyplot as plt\n",
"mpl.rcParams['figure.figsize'] = (8, 6)\n",
"mpl.rcParams['axes.grid'] = False"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Loading TimesFM pretrained checkpoint"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"timesfm_backend = \"gpu\" # @param\n",
"\n",
"tfm = timesfm.TimesFm(\n",
" hparams=timesfm.TimesFmHparams(\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-2.0-500m-jax\"),\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Evaluating pretrained checkpoint on ETT datasets"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"DATA_DICT = {\n",
" \"ettm2\": {\n",
" \"boundaries\": [34560, 46080, 57600],\n",
" \"data_path\": \"../datasets/ETT-small/ETTm2.csv\",\n",
" \"freq\": \"15min\",\n",
" },\n",
" \"ettm1\": {\n",
" \"boundaries\": [34560, 46080, 57600],\n",
" \"data_path\": \"../datasets/ETT-small/ETTm1.csv\",\n",
" \"freq\": \"15min\",\n",
" },\n",
" \"etth2\": {\n",
" \"boundaries\": [8640, 11520, 14400],\n",
" \"data_path\": \"../datasets/ETT-small/ETTh2.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"etth1\": {\n",
" \"boundaries\": [8640, 11520, 14400],\n",
" \"data_path\": \"../datasets/ETT-small/ETTh1.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"elec\": {\n",
" \"boundaries\": [18413, 21044, 26304],\n",
" \"data_path\": \"../datasets/electricity/electricity.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"traffic\": {\n",
" \"boundaries\": [12280, 14036, 17544],\n",
" \"data_path\": \"../datasets/traffic/traffic.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"weather\": {\n",
" \"boundaries\": [36887, 42157, 52696],\n",
" \"data_path\": \"../datasets/weather/weather.csv\",\n",
" \"freq\": \"10min\",\n",
" },\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"dataset = \"ettm1\"\n",
"data_path = DATA_DICT[dataset][\"data_path\"]\n",
"freq = DATA_DICT[dataset][\"freq\"]\n",
"int_freq = timesfm.freq_map(freq)\n",
"boundaries = DATA_DICT[dataset][\"boundaries\"]\n",
"\n",
"data_df = pd.read_csv(open(data_path, \"r\"))\n",
"\n",
"\n",
"ts_cols = [col for col in data_df.columns if col != \"date\"]\n",
"num_cov_cols = None\n",
"cat_cov_cols = None\n",
"\n",
"context_len = 512\n",
"pred_len = 96\n",
"\n",
"num_ts = len(ts_cols)\n",
"batch_size = 8\n",
"\n",
"dtl = data_loader.TimeSeriesdata(\n",
" data_path=data_path,\n",
" datetime_col=\"date\",\n",
" num_cov_cols=num_cov_cols,\n",
" cat_cov_cols=cat_cov_cols,\n",
" ts_cols=np.array(ts_cols),\n",
" train_range=[0, boundaries[0]],\n",
" val_range=[boundaries[0], boundaries[1]],\n",
" test_range=[boundaries[1], boundaries[2]],\n",
" hist_len=context_len,\n",
" pred_len=pred_len,\n",
" batch_size=num_ts,\n",
" freq=freq,\n",
" normalize=True,\n",
" epoch_len=None,\n",
" holiday=False,\n",
" permute=True,\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_batches = dtl.tf_dataset(mode=\"train\", shift=1).batch(batch_size)\n",
"val_batches = dtl.tf_dataset(mode=\"val\", shift=pred_len)\n",
"test_batches = dtl.tf_dataset(mode=\"test\", shift=pred_len)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"for tbatch in tqdm(train_batches.as_numpy_iterator()):\n",
" break\n",
"print(tbatch[0].shape)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### MAE on the test split for the pretrained TimesFM model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"mae_losses = []\n",
"for batch in tqdm(test_batches.as_numpy_iterator()):\n",
" past = batch[0]\n",
" actuals = batch[3]\n",
" forecasts, _ = tfm.forecast(list(past), [0] * past.shape[0], normalize=True)\n",
" forecasts = forecasts[:, 0 : actuals.shape[1]]\n",
" mae_losses.append(np.abs(forecasts - actuals).mean())\n",
"\n",
"print(f\"MAE: {np.mean(mae_losses)}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Finetuning the model on the ETT dataset"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [],
"source": [
"import jax\n",
"from jax import numpy as jnp\n",
"from praxis import pax_fiddle\n",
"from praxis import py_utils\n",
"from praxis import pytypes\n",
"from praxis import base_model\n",
"from praxis import optimizers\n",
"from praxis import schedules\n",
"from praxis import base_hyperparams\n",
"from praxis import base_layer\n",
"from paxml import tasks_lib\n",
"from paxml import trainer_lib\n",
"from paxml import checkpoints\n",
"from paxml import learners\n",
"from paxml import partitioning\n",
"from paxml import checkpoint_types"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"# PAX shortcuts\n",
"NestedMap = py_utils.NestedMap\n",
"WeightInit = base_layer.WeightInit\n",
"WeightHParams = base_layer.WeightHParams\n",
"InstantiableParams = py_utils.InstantiableParams\n",
"JTensor = pytypes.JTensor\n",
"NpTensor = pytypes.NpTensor\n",
"WeightedScalars = pytypes.WeightedScalars\n",
"instantiate = base_hyperparams.instantiate\n",
"LayerTpl = pax_fiddle.Config[base_layer.BaseLayer]\n",
"AuxLossStruct = base_layer.AuxLossStruct\n",
"\n",
"AUX_LOSS = base_layer.AUX_LOSS\n",
"template_field = base_layer.template_field\n",
"\n",
"# Standard prng key names\n",
"PARAMS = base_layer.PARAMS\n",
"RANDOM = base_layer.RANDOM\n",
"\n",
"key = jax.random.PRNGKey(seed=1234)"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [],
"source": [
"model = pax_fiddle.Config(\n",
" patched_decoder.PatchedDecoderFinetuneModel,\n",
" name='patched_decoder_finetune',\n",
" core_layer_tpl=tfm.model_p,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### We will hold the transformer layers fixed while finetuning, while training all other components."
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"@pax_fiddle.auto_config\n",
"def build_learner() -> learners.Learner:\n",
" return pax_fiddle.Config(\n",
" learners.Learner,\n",
" name='learner',\n",
" loss_name='avg_qloss',\n",
" optimizer=optimizers.Adam(\n",
" epsilon=1e-7,\n",
" clip_threshold=1e2,\n",
" learning_rate=1e-2,\n",
" lr_schedule=pax_fiddle.Config(\n",
" schedules.Cosine,\n",
" initial_value=1e-3,\n",
" final_value=1e-4,\n",
" total_steps=40000,\n",
" ),\n",
" ema_decay=0.9999,\n",
" ),\n",
" # Linear probing i.e we hold the transformer layers fixed.\n",
" bprop_variable_exclusion=['.*/stacked_transformer_layer/.*'],\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"task_p = tasks_lib.SingleTask(\n",
" name='ts-learn',\n",
" model=model,\n",
" train=tasks_lib.SingleTask.Train(\n",
" learner=build_learner(),\n",
" ),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"task_p.model.ici_mesh_shape = [1, 1, 1]\n",
"task_p.model.mesh_axis_names = ['replica', 'data', 'mdl']\n",
"\n",
"DEVICES = np.array(jax.devices()).reshape([1, 1, 1])\n",
"MESH = jax.sharding.Mesh(DEVICES, ['replica', 'data', 'mdl'])\n",
"\n",
"num_devices = jax.local_device_count()\n",
"print(f'num_devices: {num_devices}')\n",
"print(f'device kind: {jax.local_devices()[0].device_kind}')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"jax_task = task_p\n",
"key, init_key = jax.random.split(key)\n",
"\n",
"# To correctly prepare a batch of data for model initialization (now that shape\n",
"# inference is merged), we take one devices*batch_size tensor tuple of data,\n",
"# slice out just one batch, then run the prepare_input_batch function over it.\n",
"\n",
"\n",
"def process_train_batch(batch):\n",
" past_ts = batch[0].reshape(batch_size * num_ts, -1)\n",
" actual_ts = batch[3].reshape(batch_size * num_ts, -1)\n",
" return NestedMap(input_ts=past_ts, actual_ts=actual_ts)\n",
"\n",
"\n",
"def process_eval_batch(batch):\n",
" past_ts = batch[0]\n",
" actual_ts = batch[3]\n",
" return NestedMap(input_ts=past_ts, actual_ts=actual_ts)\n",
"\n",
"\n",
"jax_model_states, _ = trainer_lib.initialize_model_state(\n",
" jax_task,\n",
" init_key,\n",
" process_train_batch(tbatch),\n",
" checkpoint_type=checkpoint_types.CheckpointType.GDA,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Setting the initial model weights to the pretrained TimesFM parameters."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"jax_model_states.mdl_vars['params']['core_layer'] = tfm._train_state.mdl_vars['params']\n",
"jax_vars = jax_model_states.mdl_vars\n",
"gc.collect()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Training loop"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {},
"outputs": [],
"source": [
"jax_task = task_p\n",
"\n",
"\n",
"def train_step(states, prng_key, inputs):\n",
" return trainer_lib.train_step_single_learner(\n",
" jax_task, states, prng_key, inputs\n",
" )\n",
"\n",
"\n",
"def eval_step(states, prng_key, inputs):\n",
" states = states.to_eval_state()\n",
" return trainer_lib.eval_step_single_learner(\n",
" jax_task, states, prng_key, inputs\n",
" )\n",
"\n",
"key, train_key, eval_key = jax.random.split(key, 3)\n",
"train_prng_seed = jax.random.split(train_key, num=jax.local_device_count())\n",
"eval_prng_seed = jax.random.split(eval_key, num=jax.local_device_count())\n",
"\n",
"p_train_step = jax.pmap(train_step, axis_name='batch')\n",
"p_eval_step = jax.pmap(eval_step, axis_name='batch')"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [],
"source": [
"replicated_jax_states = trainer_lib.replicate_model_state(jax_model_states)\n",
"replicated_jax_vars = replicated_jax_states.mdl_vars"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {},
"outputs": [],
"source": [
"best_eval_loss = 1e7\n",
"step_count = 0\n",
"patience = 0\n",
"NUM_EPOCHS = 100\n",
"PATIENCE = 5\n",
"TRAIN_STEPS_PER_EVAL = 1000\n",
"CHECKPOINT_DIR='/home/senrajat_google_com/ettm1_finetune'"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {},
"outputs": [],
"source": [
"def reshape_batch_for_pmap(batch, num_devices):\n",
" def _reshape(input_tensor):\n",
" bsize = input_tensor.shape[0]\n",
" residual_shape = list(input_tensor.shape[1:])\n",
" nbsize = bsize // num_devices\n",
" return jnp.reshape(input_tensor, [num_devices, nbsize] + residual_shape)\n",
"\n",
" return jax.tree.map(_reshape, batch)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"for epoch in range(NUM_EPOCHS):\n",
" print(f\"__________________Epoch: {epoch}__________________\", flush=True)\n",
" train_its = train_batches.as_numpy_iterator()\n",
" if patience >= PATIENCE:\n",
" print(\"Early stopping.\", flush=True)\n",
" break\n",
" for batch in tqdm(train_its):\n",
" train_losses = []\n",
" if patience >= PATIENCE:\n",
" print(\"Early stopping.\", flush=True)\n",
" break\n",
" tbatch = process_train_batch(batch)\n",
" tbatch = reshape_batch_for_pmap(tbatch, num_devices)\n",
" replicated_jax_states, step_fun_out = p_train_step(\n",
" replicated_jax_states, train_prng_seed, tbatch\n",
" )\n",
" train_losses.append(step_fun_out.loss[0])\n",
" if step_count % TRAIN_STEPS_PER_EVAL == 0:\n",
" print(\n",
" f\"Train loss at step {step_count}: {np.mean(train_losses)}\",\n",
" flush=True,\n",
" )\n",
" train_losses = []\n",
" print(\"Starting eval.\", flush=True)\n",
" val_its = val_batches.as_numpy_iterator()\n",
" eval_losses = []\n",
" for ev_batch in tqdm(val_its):\n",
" ebatch = process_eval_batch(ev_batch)\n",
" ebatch = reshape_batch_for_pmap(ebatch, num_devices)\n",
" _, step_fun_out = p_eval_step(\n",
" replicated_jax_states, eval_prng_seed, ebatch\n",
" )\n",
" eval_losses.append(step_fun_out.loss[0])\n",
" mean_loss = np.mean(eval_losses)\n",
" print(f\"Eval loss at step {step_count}: {mean_loss}\", flush=True)\n",
" if mean_loss < best_eval_loss or np.isnan(mean_loss):\n",
" best_eval_loss = mean_loss\n",
" print(\"Saving checkpoint.\")\n",
" jax_state_for_saving = py_utils.maybe_unreplicate_for_fully_replicated(\n",
" replicated_jax_states\n",
" )\n",
" checkpoints.save_checkpoint(\n",
" jax_state_for_saving, CHECKPOINT_DIR, overwrite=True\n",
" )\n",
" patience = 0\n",
" del jax_state_for_saving\n",
" gc.collect()\n",
" else:\n",
" patience += 1\n",
" print(f\"patience: {patience}\")\n",
" step_count += 1"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Loading and evaluating the best (according to validation loss) finetuned checkpoint"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_state = checkpoints.restore_checkpoint(jax_model_states, CHECKPOINT_DIR)\n",
"print(train_state.step)\n",
"tfm._train_state.mdl_vars['params'] = train_state.mdl_vars['params']['core_layer']\n",
"tfm.jit_decode()\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"mae_losses = []\n",
"for batch in tqdm(test_batches.as_numpy_iterator()):\n",
" past = batch[0]\n",
" actuals = batch[3]\n",
" _, forecasts = tfm.forecast(list(past), [0] * past.shape[0])\n",
" forecasts = forecasts[:, 0 : actuals.shape[1], 5]\n",
" mae_losses.append(np.abs(forecasts - actuals).mean())\n",
"\n",
"print(f\"MAE: {np.mean(mae_losses)}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## There is around a __7%__ reduction in MAE from finetuning."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "chronos-v2",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.15"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
+538
View File
@@ -0,0 +1,538 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Introduction\n",
"This notebook shows how to use TimesFM with finetuning. \n",
"\n",
"In order to perform finetuning, you need to create the Pytorch Dataset in a proper format. The example of the Dataset is provided below.\n",
"The finetuning code can be found in timesfm.finetuning_torch.py. This notebook just imports the methods from finetuning"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Dataset Creation"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"TimesFM v1.2.0. See https://github.com/google-research/timesfm/blob/master/README.md for updated APIs.\n",
"Loaded Jax TimesFM.\n",
"Loaded PyTorch TimesFM.\n"
]
}
],
"source": [
"from os import path\n",
"from typing import Optional, Tuple\n",
"\n",
"import numpy as np\n",
"import pandas as pd\n",
"import torch\n",
"import torch.multiprocessing as mp\n",
"import yfinance as yf\n",
"from finetuning.finetuning_torch import FinetuningConfig, TimesFMFinetuner\n",
"from huggingface_hub import snapshot_download\n",
"from torch.utils.data import Dataset\n",
"\n",
"from timesfm import TimesFm, TimesFmCheckpoint, TimesFmHparams\n",
"from timesfm.pytorch_patched_decoder import PatchedTimeSeriesDecoder\n",
"import os\n",
"\n",
"\n",
"class TimeSeriesDataset(Dataset):\n",
" \"\"\"Dataset for time series data compatible with TimesFM.\"\"\"\n",
"\n",
" def __init__(self,\n",
" series: np.ndarray,\n",
" context_length: int,\n",
" horizon_length: int,\n",
" freq_type: int = 0):\n",
" \"\"\"\n",
" Initialize dataset.\n",
"\n",
" Args:\n",
" series: Time series data\n",
" context_length: Number of past timesteps to use as input\n",
" horizon_length: Number of future timesteps to predict\n",
" freq_type: Frequency type (0, 1, or 2)\n",
" \"\"\"\n",
" if freq_type not in [0, 1, 2]:\n",
" raise ValueError(\"freq_type must be 0, 1, or 2\")\n",
"\n",
" self.series = series\n",
" self.context_length = context_length\n",
" self.horizon_length = horizon_length\n",
" self.freq_type = freq_type\n",
" self._prepare_samples()\n",
"\n",
" def _prepare_samples(self) -> None:\n",
" \"\"\"Prepare sliding window samples from the time series.\"\"\"\n",
" self.samples = []\n",
" total_length = self.context_length + self.horizon_length\n",
"\n",
" for start_idx in range(0, len(self.series) - total_length + 1):\n",
" end_idx = start_idx + self.context_length\n",
" x_context = self.series[start_idx:end_idx]\n",
" x_future = self.series[end_idx:end_idx + self.horizon_length]\n",
" self.samples.append((x_context, x_future))\n",
"\n",
" def __len__(self) -> int:\n",
" return len(self.samples)\n",
"\n",
" def __getitem__(\n",
" self, index: int\n",
" ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n",
" x_context, x_future = self.samples[index]\n",
"\n",
" x_context = torch.tensor(x_context, dtype=torch.float32)\n",
" x_future = torch.tensor(x_future, dtype=torch.float32)\n",
"\n",
" input_padding = torch.zeros_like(x_context)\n",
" freq = torch.tensor([self.freq_type], dtype=torch.long)\n",
"\n",
" return x_context, input_padding, freq, x_future\n",
"\n",
"def prepare_datasets(series: np.ndarray,\n",
" context_length: int,\n",
" horizon_length: int,\n",
" freq_type: int = 0,\n",
" train_split: float = 0.8) -> Tuple[Dataset, Dataset]:\n",
" \"\"\"\n",
" Prepare training and validation datasets from time series data.\n",
"\n",
" Args:\n",
" series: Input time series data\n",
" context_length: Number of past timesteps to use\n",
" horizon_length: Number of future timesteps to predict\n",
" freq_type: Frequency type (0, 1, or 2)\n",
" train_split: Fraction of data to use for training\n",
"\n",
" Returns:\n",
" Tuple of (train_dataset, val_dataset)\n",
" \"\"\"\n",
" train_size = int(len(series) * train_split)\n",
" train_data = series[:train_size]\n",
" val_data = series[train_size:]\n",
"\n",
" # Create datasets with specified frequency type\n",
" train_dataset = TimeSeriesDataset(train_data,\n",
" context_length=context_length,\n",
" horizon_length=horizon_length,\n",
" freq_type=freq_type)\n",
"\n",
" val_dataset = TimeSeriesDataset(val_data,\n",
" context_length=context_length,\n",
" horizon_length=horizon_length,\n",
" freq_type=freq_type)\n",
"\n",
" return train_dataset, val_dataset\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Model Creation"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"def get_model(load_weights: bool = False):\n",
" device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
" repo_id = \"google/timesfm-2.0-500m-pytorch\"\n",
" hparams = TimesFmHparams(\n",
" backend=device,\n",
" per_core_batch_size=32,\n",
" horizon_len=128,\n",
" num_layers=50,\n",
" use_positional_embedding=False,\n",
" context_len=\n",
" 192, # Context length can be anything up to 2048 in multiples of 32\n",
" )\n",
" tfm = TimesFm(hparams=hparams,\n",
" checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))\n",
"\n",
" model = PatchedTimeSeriesDecoder(tfm._model_config)\n",
" if load_weights:\n",
" checkpoint_path = path.join(snapshot_download(repo_id), \"torch_model.ckpt\")\n",
" loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)\n",
" model.load_state_dict(loaded_checkpoint)\n",
" return model, hparams, tfm._model_config\n"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"def plot_predictions(\n",
" model: TimesFm,\n",
" val_dataset: Dataset,\n",
" save_path: Optional[str] = \"predictions.png\",\n",
") -> None:\n",
" \"\"\"\n",
" Plot model predictions against ground truth for a batch of validation data.\n",
"\n",
" Args:\n",
" model: Trained TimesFM model\n",
" val_dataset: Validation dataset\n",
" save_path: Path to save the plot\n",
" \"\"\"\n",
" import matplotlib.pyplot as plt\n",
"\n",
" model.eval()\n",
"\n",
" x_context, x_padding, freq, x_future = val_dataset[0]\n",
" x_context = x_context.unsqueeze(0) # Add batch dimension\n",
" x_padding = x_padding.unsqueeze(0)\n",
" freq = freq.unsqueeze(0)\n",
" x_future = x_future.unsqueeze(0)\n",
"\n",
" device = next(model.parameters()).device\n",
" x_context = x_context.to(device)\n",
" x_padding = x_padding.to(device)\n",
" freq = freq.to(device)\n",
" x_future = x_future.to(device)\n",
"\n",
" with torch.no_grad():\n",
" predictions = model(x_context, x_padding.float(), freq)\n",
" predictions_mean = predictions[..., 0] # [B, N, horizon_len]\n",
" last_patch_pred = predictions_mean[:, -1, :] # [B, horizon_len]\n",
"\n",
" context_vals = x_context[0].cpu().numpy()\n",
" future_vals = x_future[0].cpu().numpy()\n",
" pred_vals = last_patch_pred[0].cpu().numpy()\n",
"\n",
" context_len = len(context_vals)\n",
" horizon_len = len(future_vals)\n",
"\n",
" plt.figure(figsize=(12, 6))\n",
"\n",
" plt.plot(range(context_len),\n",
" context_vals,\n",
" label=\"Historical Data\",\n",
" color=\"blue\",\n",
" linewidth=2)\n",
"\n",
" plt.plot(\n",
" range(context_len, context_len + horizon_len),\n",
" future_vals,\n",
" label=\"Ground Truth\",\n",
" color=\"green\",\n",
" linestyle=\"--\",\n",
" linewidth=2,\n",
" )\n",
"\n",
" plt.plot(range(context_len, context_len + horizon_len),\n",
" pred_vals,\n",
" label=\"Prediction\",\n",
" color=\"red\",\n",
" linewidth=2)\n",
"\n",
" plt.xlabel(\"Time Step\")\n",
" plt.ylabel(\"Value\")\n",
" plt.title(\"TimesFM Predictions vs Ground Truth\")\n",
" plt.legend()\n",
" plt.grid(True)\n",
"\n",
" if save_path:\n",
" plt.savefig(save_path)\n",
" print(f\"Plot saved to {save_path}\")\n",
"\n",
" plt.close()\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"def get_data(context_len: int,\n",
" horizon_len: int,\n",
" freq_type: int = 0) -> Tuple[Dataset, Dataset]:\n",
" df = yf.download(\"AAPL\", start=\"2010-01-01\", end=\"2019-01-01\")\n",
" time_series = df[\"Close\"].values\n",
"\n",
" train_dataset, val_dataset = prepare_datasets(\n",
" series=time_series,\n",
" context_length=context_len,\n",
" horizon_length=horizon_len,\n",
" freq_type=freq_type,\n",
" train_split=0.8,\n",
" )\n",
"\n",
" print(f\"Created datasets:\")\n",
" print(f\"- Training samples: {len(train_dataset)}\")\n",
" print(f\"- Validation samples: {len(val_dataset)}\")\n",
" print(f\"- Using frequency type: {freq_type}\")\n",
" return train_dataset, val_dataset\n",
"\n",
"\n",
"\n",
"def single_gpu_example():\n",
" \"\"\"Basic example of finetuning TimesFM on stock data.\"\"\"\n",
" model, hparams, tfm_config = get_model(load_weights=True)\n",
" config = FinetuningConfig(batch_size=256,\n",
" num_epochs=5,\n",
" learning_rate=1e-4,\n",
" use_wandb=True,\n",
" freq_type=1,\n",
" log_every_n_steps=10,\n",
" val_check_interval=0.5,\n",
" use_quantile_loss=True)\n",
"\n",
" train_dataset, val_dataset = get_data(128,\n",
" tfm_config.horizon_len,\n",
" freq_type=config.freq_type)\n",
" finetuner = TimesFMFinetuner(model, config)\n",
"\n",
" print(\"\\nStarting finetuning...\")\n",
" results = finetuner.finetune(train_dataset=train_dataset,\n",
" val_dataset=val_dataset)\n",
"\n",
" print(\"\\nFinetuning completed!\")\n",
" print(f\"Training history: {len(results['history']['train_loss'])} epochs\")\n",
"\n",
" plot_predictions(\n",
" model=model,\n",
" val_dataset=val_dataset,\n",
" save_path=\"timesfm_predictions.png\",\n",
" )\n"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "ac84aeda3a1749ae8f30b06859067bb1",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Fetching 3 files: 0%| | 0/3 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "6d9d8081fc514c6d8601a2e0e63954a2",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Fetching 3 files: 0%| | 0/3 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"[*********************100%***********************] 1 of 1 completed\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Created datasets:\n",
"- Training samples: 1556\n",
"- Validation samples: 198\n",
"- Using frequency type: 1\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[34m\u001b[1mwandb\u001b[0m: Using wandb-core as the SDK backend. Please refer to https://wandb.me/wandb-core for more information.\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mmishacamry\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
]
},
{
"data": {
"text/html": [
"Tracking run with wandb version 0.19.1"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"Run data is saved locally in <code>/home/chertushkin/forks/timesfm/notebooks/wandb/run-20250217_114343-tjs63ml2</code>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"Syncing run <strong><a href='https://wandb.ai/mishacamry/timesfm-finetuning/runs/tjs63ml2' target=\"_blank\">chocolate-eon-50</a></strong> to <a href='https://wandb.ai/mishacamry/timesfm-finetuning' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target=\"_blank\">docs</a>)<br>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
" View project at <a href='https://wandb.ai/mishacamry/timesfm-finetuning' target=\"_blank\">https://wandb.ai/mishacamry/timesfm-finetuning</a>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
" View run at <a href='https://wandb.ai/mishacamry/timesfm-finetuning/runs/tjs63ml2' target=\"_blank\">https://wandb.ai/mishacamry/timesfm-finetuning/runs/tjs63ml2</a>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"Starting finetuning...\n"
]
},
{
"data": {
"text/html": [],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<br> <style><br> .wandb-row {<br> display: flex;<br> flex-direction: row;<br> flex-wrap: wrap;<br> justify-content: flex-start;<br> width: 100%;<br> }<br> .wandb-col {<br> display: flex;<br> flex-direction: column;<br> flex-basis: 100%;<br> flex: 1;<br> padding: 10px;<br> }<br> </style><br><div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>epoch</td><td>▁▃▅▆█</td></tr><tr><td>learning_rate</td><td>▁▁▁▁▁</td></tr><tr><td>train_loss</td><td>█▃▂▁▁</td></tr><tr><td>val_loss</td><td>█▁▄▁▂</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>epoch</td><td>5</td></tr><tr><td>learning_rate</td><td>0.0001</td></tr><tr><td>train_loss</td><td>2.85423</td></tr><tr><td>val_loss</td><td>26.7628</td></tr></table><br/></div></div>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
" View run <strong style=\"color:#cdcd00\">chocolate-eon-50</strong> at: <a href='https://wandb.ai/mishacamry/timesfm-finetuning/runs/tjs63ml2' target=\"_blank\">https://wandb.ai/mishacamry/timesfm-finetuning/runs/tjs63ml2</a><br> View project at: <a href='https://wandb.ai/mishacamry/timesfm-finetuning' target=\"_blank\">https://wandb.ai/mishacamry/timesfm-finetuning</a><br>Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"Find logs at: <code>./wandb/run-20250217_114343-tjs63ml2/logs</code>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"Finetuning completed!\n",
"Training history: 5 epochs\n",
"Plot saved to timesfm_predictions.png\n"
]
}
],
"source": [
"single_gpu_example()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "timesfm-DnAbSweh-py3.11",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.10"
}
},
"nbformat": 4,
"nbformat_minor": 2
}