Merge pull request #243 from misha-chertushkin/notebook-import-fix

Fix imports in Notebook examples

README.md

@@ -21,6 +21,10 @@ We recommend at least 32GB RAM to load TimesFM dependencies.
 - Launched [finetuning support](https://github.com/google-research/timesfm/blob/master/notebooks/finetuning.ipynb) that lets you finetune the weights of the pretrained TimesFM model on your own data.
 - Launched [~zero-shot covariate support](https://github.com/google-research/timesfm/blob/master/notebooks/covariates.ipynb) with external regressors. More details [here](https://github.com/google-research/timesfm?tab=readme-ov-file#covariates-support).
 
+## Update - Feb. 17, 2024
+- We are providing the option for [finetuning using Pytorch](https://github.com/google-research/timesfm/blob/master/notebooks/finetuning_torch.ipynb), which mimics the previously added functionality from [finetuning support](https://github.com/google-research/timesfm/blob/master/notebooks/finetuning.ipynb).
+- We are also providing the Multi-GPU finetuining with Pytorch. We currently support DDP multi-gpu finetuning, other variants of multi-gpu training (pipeline parallelism/model parallelism) might be added later. In order to use it, follow the steps in [finetuning example](https://github.com/google-research/timesfm/blob/master/finetuning/finetuning_example.py) .
+
 ## Checkpoint timesfm-1.0-200m (-pytorch)
 
 timesfm-1.0-200m is our first open model checkpoint:

notebooks/finetuning_torch.ipynb

@@ -20,9 +20,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
-   "outputs": [],
+   "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",
@@ -32,7 +42,7 @@
     "import torch\n",
     "import torch.multiprocessing as mp\n",
     "import yfinance as yf\n",
-    "from timesfm.finetuning_torch import FinetuningConfig, TimesFMFinetuner\n",
+    "from finetuning.finetuning_torch import FinetuningConfig, TimesFMFinetuner\n",
     "from huggingface_hub import snapshot_download\n",
     "from torch.utils.data import Dataset\n",
     "\n",
@@ -139,7 +149,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -168,7 +178,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -252,10 +262,31 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "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",
@@ -290,17 +321,202 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "metadata": {},
-   "outputs": [],
+   "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-311",
+   "display_name": "timesfm-DnAbSweh-py3.11",
    "language": "python",
    "name": "python3"
   },
@@ -314,7 +530,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.11"
+   "version": "3.11.10"
   }
  },
  "nbformat": 4,