{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "## Importing relevant packages for finetuning" ] }, { "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": [ "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": null, "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": [ "tfm = timesfm.TimesFm(\n", " hparams=timesfm.TimesFmHparams(\n", " backend=\"gpu\",\n", " per_core_batch_size=32,\n", " horizon_len=128,\n", " ),\n", " checkpoint=timesfm.TimesFmCheckpoint(\n", " huggingface_repo_id=\"google/timesfm-1.0-200m\"),\n", " )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Evaluating pretrained checkpoint on ETT datasets" ] }, { "cell_type": "code", "execution_count": null, "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": null, "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 = 16\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", " pass\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": null, "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": null, "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": null, "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": null, "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": null, "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": null, "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": null, "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": null, "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": null, "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 __9%__ reduction in MAE from finetuning." ] } ], "metadata": { "kernelspec": { "display_name": "tfm_env_v3", "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.14" } }, "nbformat": 4, "nbformat_minor": 2 }