notebooks/hfdemo/patch_tsmixer_getting_started.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7478e0e2-b7af-4fd4-b44e-ca58e0c31b71",
   "metadata": {},
   "source": [
    "# Getting started with `PatchTSMixer`\n",
    "## Direct forecasting example\n",
    "\n",
    "This notebooke demonstrates the usage of a `PatchTSMixer` model for a multivariate time series forecasting task. This notebook has a dependecy on HuggingFace [transformers](https://github.com/huggingface/transformers) repo. For details related to model architecture, refer to the [TSMixer paper](https://arxiv.org/abs/2306.09364)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f63ae353-96df-4380-89f6-1e6cebf684fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/dccstor/dnn_forecasting/conda_envs/envs/hf/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "2023-12-11 04:17:23.961875: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2023-12-11 04:17:23.962007: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2023-12-11 04:17:24.020776: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2023-12-11 04:17:30.476420: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "# Standard\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "\n",
    "# Third Party\n",
    "from transformers import (\n",
    "    EarlyStoppingCallback,\n",
    "    PatchTSMixerConfig,\n",
    "    PatchTSMixerForPrediction,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    ")\n",
    "\n",
    "# First Party\n",
    "from tsfm_public.toolkit.dataset import ForecastDFDataset\n",
    "from tsfm_public.toolkit.time_series_preprocessor import TimeSeriesPreprocessor\n",
    "from tsfm_public.toolkit.util import select_by_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a826c4f3-1c6c-4088-b6af-f430f45fd380",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set seed for reproducibility\n",
    "SEED = 42\n",
    "torch.manual_seed(SEED)\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e4eb9be-c19f-448f-a4bd-c600e068633f",
   "metadata": {},
   "source": [
    "## Load and prepare datasets\n",
    "\n",
    "In the next cell, please adjust the following parameters to suit your application:\n",
    "- `dataset_path`: path to local .csv file, or web address to a csv file for the data of interest. Data is loaded with pandas, so anything supported by\n",
    "`pd.read_csv` is supported: (https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_csv.html).\n",
    "- `timestamp_column`: column name containing timestamp information, use None if there is no such column\n",
    "- `id_columns`: List of column names specifying the IDs of different time series. If no ID column exists, use []\n",
    "- `forecast_columns`: List of columns to be modeled\n",
    "- `context_length`: The amount of historical data used as input to the model. Windows of the input time series data with length equal to\n",
    "context_length will be extracted from the input dataframe. In the case of a multi-time series dataset, the context windows will be created\n",
    "so that they are contained within a single time series (i.e., a single ID).\n",
    "- `forecast_horizon`: Number of time stamps to forecast in future.\n",
    "- `train_start_index`, `train_end_index`: the start and end indices in the loaded data which delineate the training data.\n",
    "- `valid_start_index`, `valid_end_index`: the start and end indices in the loaded data which delineate the validation data.\n",
    "- `test_start_index`, `test_end_index`: the start and end indices in the loaded data which delineate the test data.\n",
    "- `patch_length`: The patch length for the `PatchTSMixer` model. Recommended to have a value so that `context_length` is divisible by it.\n",
    "- `num_workers`: Number of dataloder workers in pytorch dataloader.\n",
    "- `batch_size`: Batch size. \n",
    "The data is first loaded into a Pandas dataframe and split into training, validation, and test parts. Then the pandas dataframes are converted\n",
    "to the appropriate torch dataset needed for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d4c1e812-f2d6-4ccb-a79c-47879b562d85",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = \"ETTh1\"\n",
    "num_workers = 8  # Reduce this if you have low number of CPU cores\n",
    "batch_size = 32  # Reduce if not enough GPU memory available\n",
    "context_length = 512\n",
    "forecast_horizon = 96\n",
    "patch_length = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "19ca5a76-64d4-4f8d-92f5-67f76c1685fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading target dataset: ETTh1\n"
     ]
    }
   ],
   "source": [
    "print(f\"Loading target dataset: {dataset}\")\n",
    "dataset_path = f\"https://raw.githubusercontent.com/zhouhaoyi/ETDataset/main/ETT-small/{dataset}.csv\"\n",
    "timestamp_column = \"date\"\n",
    "id_columns = []\n",
    "forecast_columns = [\"HUFL\", \"HULL\", \"MUFL\", \"MULL\", \"LUFL\", \"LULL\", \"OT\"]\n",
    "train_start_index = None  # None indicates beginning of dataset\n",
    "train_end_index = 12 * 30 * 24\n",
    "\n",
    "# we shift the start of the validation/test period back by context length so that\n",
    "# the first validation/test timestamp is immediately following the training data\n",
    "valid_start_index = 12 * 30 * 24 - context_length\n",
    "valid_end_index = 12 * 30 * 24 + 4 * 30 * 24\n",
    "\n",
    "test_start_index = 12 * 30 * 24 + 4 * 30 * 24 - context_length\n",
    "test_end_index = 12 * 30 * 24 + 8 * 30 * 24"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0e856ce6-11e9-4c8c-9aeb-3e4a8b05f3eb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TimeSeriesPreprocessor {\n",
       "  \"context_length\": 64,\n",
       "  \"feature_extractor_type\": \"TimeSeriesPreprocessor\",\n",
       "  \"id_columns\": [],\n",
       "  \"input_columns\": [\n",
       "    \"HUFL\",\n",
       "    \"HULL\",\n",
       "    \"MUFL\",\n",
       "    \"MULL\",\n",
       "    \"LUFL\",\n",
       "    \"LULL\",\n",
       "    \"OT\"\n",
       "  ],\n",
       "  \"output_columns\": [\n",
       "    \"HUFL\",\n",
       "    \"HULL\",\n",
       "    \"MUFL\",\n",
       "    \"MULL\",\n",
       "    \"LUFL\",\n",
       "    \"LULL\",\n",
       "    \"OT\"\n",
       "  ],\n",
       "  \"prediction_length\": null,\n",
       "  \"processor_class\": \"TimeSeriesPreprocessor\",\n",
       "  \"scaler_dict\": {\n",
       "    \"0\": {\n",
       "      \"copy\": true,\n",
       "      \"feature_names_in_\": [\n",
       "        \"HUFL\",\n",
       "        \"HULL\",\n",
       "        \"MUFL\",\n",
       "        \"MULL\",\n",
       "        \"LUFL\",\n",
       "        \"LULL\",\n",
       "        \"OT\"\n",
       "      ],\n",
       "      \"mean_\": [\n",
       "        7.937742245659508,\n",
       "        2.0210386567335163,\n",
       "        5.079770601157927,\n",
       "        0.7461858799957015,\n",
       "        2.781762386375555,\n",
       "        0.7884531235540096,\n",
       "        17.1282616982271\n",
       "      ],\n",
       "      \"n_features_in_\": 7,\n",
       "      \"n_samples_seen_\": 8640,\n",
       "      \"scale_\": [\n",
       "        5.812749409143771,\n",
       "        2.0901046504076,\n",
       "        5.518793579036245,\n",
       "        1.9263792741329822,\n",
       "        1.0235226594952194,\n",
       "        0.6302366362251923,\n",
       "        9.176491024944335\n",
       "      ],\n",
       "      \"var_\": [\n",
       "        33.78805569350125,\n",
       "        4.368537449655475,\n",
       "        30.457082568011693,\n",
       "        3.710937107809115,\n",
       "        1.0475986345001667,\n",
       "        0.39719821764044544,\n",
       "        84.20798753088393\n",
       "      ],\n",
       "      \"with_mean\": true,\n",
       "      \"with_std\": true\n",
       "    }\n",
       "  },\n",
       "  \"scaling\": true,\n",
       "  \"time_series_task\": \"forecasting\",\n",
       "  \"timestamp_column\": \"date\"\n",
       "}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = pd.read_csv(\n",
    "    dataset_path,\n",
    "    parse_dates=[timestamp_column],\n",
    ")\n",
    "\n",
    "train_data = select_by_index(\n",
    "    data,\n",
    "    id_columns=id_columns,\n",
    "    start_index=train_start_index,\n",
    "    end_index=train_end_index,\n",
    ")\n",
    "valid_data = select_by_index(\n",
    "    data,\n",
    "    id_columns=id_columns,\n",
    "    start_index=valid_start_index,\n",
    "    end_index=valid_end_index,\n",
    ")\n",
    "test_data = select_by_index(\n",
    "    data,\n",
    "    id_columns=id_columns,\n",
    "    start_index=test_start_index,\n",
    "    end_index=test_end_index,\n",
    ")\n",
    "\n",
    "tsp = TimeSeriesPreprocessor(\n",
    "    timestamp_column=timestamp_column,\n",
    "    id_columns=id_columns,\n",
    "    target_columns=forecast_columns,\n",
    "    scaling=True,\n",
    ")\n",
    "tsp.train(train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "678d849d-41fc-450d-a855-1dde27179b31",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = ForecastDFDataset(\n",
    "    tsp.preprocess(train_data),\n",
    "    id_columns=id_columns,\n",
    "    target_columns=forecast_columns,\n",
    "    context_length=context_length,\n",
    "    prediction_length=forecast_horizon,\n",
    ")\n",
    "valid_dataset = ForecastDFDataset(\n",
    "    tsp.preprocess(valid_data),\n",
    "    id_columns=id_columns,\n",
    "    target_columns=forecast_columns,\n",
    "    context_length=context_length,\n",
    "    prediction_length=forecast_horizon,\n",
    ")\n",
    "test_dataset = ForecastDFDataset(\n",
    "    tsp.preprocess(test_data),\n",
    "    id_columns=id_columns,\n",
    "    target_columns=forecast_columns,\n",
    "    context_length=context_length,\n",
    "    prediction_length=forecast_horizon,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae939491-8813-44c9-bc3d-d1c6a5764cd4",
   "metadata": {},
   "source": [
    "## Testing with a `PatchTSMixer` model that was trained on the training part of the `ETTh1` data\n",
    "\n",
    "A pre-trained model (on `ETTh1` data) is available at [ibm-granite/granite-timeseries-patchtsmixer](https://huggingface.co/ibm-granite/granite-timeseries-patchtsmixer)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b839a219-14b9-4445-aa37-f0e466a35e62",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading pretrained model\n",
      "Done\n"
     ]
    }
   ],
   "source": [
    "print(\"Loading pretrained model\")\n",
    "inference_forecast_model = PatchTSMixerForPrediction.from_pretrained(\"ibm-granite/granite-timeseries-patchtsmixer\")\n",
    "print(\"Done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9b17c69f-f44d-4ccc-b8d4-5dd34cd1b43e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Doing testing on Etth1/test data\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='349' max='349' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [349/349 00:00]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.3698882460594177, 'eval_runtime': 2.6656, 'eval_samples_per_second': 1044.798, 'eval_steps_per_second': 130.928}\n"
     ]
    }
   ],
   "source": [
    "inference_forecast_trainer = Trainer(\n",
    "    model=inference_forecast_model,\n",
    ")\n",
    "\n",
    "print(\"\\n\\nDoing testing on Etth1/test data\")\n",
    "result = inference_forecast_trainer.evaluate(test_dataset)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19456329-1293-45bf-99c7-e5ccf0534846",
   "metadata": {},
   "source": [
    "## If we want to train from scratch\n",
    "\n",
    "Adjust the following model parameters according to need.\n",
    "- `d_model` (`int`, *optional*, defaults to 8):\n",
    "    Hidden dimension of the model. Recommended to set it as a multiple of patch_length (i.e. 2-8X of\n",
    "    patch_len). Larger value indicates more complex model.\n",
    "- `expansion_factor` (`int`, *optional*, defaults to 2):\n",
    "    Expansion factor to use inside MLP. Recommended range is 2-5. Larger value indicates more complex model.\n",
    "- `num_layers` (`int`, *optional*, defaults to 3):\n",
    "    Number of layers to use. Recommended range is 3-15. Larger value indicates more complex model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "226b904e-1ab2-478b-98b4-ce99bc23f1c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = PatchTSMixerConfig(\n",
    "    context_length=context_length,\n",
    "    prediction_length=forecast_horizon,\n",
    "    patch_length=patch_length,\n",
    "    num_input_channels=len(forecast_columns),\n",
    "    patch_stride=patch_length,\n",
    "    d_model=48,\n",
    "    num_layers=3,\n",
    "    expansion_factor=3,\n",
    "    dropout=0.5,\n",
    "    head_dropout=0.7,\n",
    "    mode=\"common_channel\",  # change it `mix_channel` if we need to explicitly model channel correlations\n",
    "    scaling=\"std\",\n",
    ")\n",
    "model = PatchTSMixerForPrediction(config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "27812e8c-c0f6-45e3-a075-310929329460",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Doing forecasting training on Etth1/train\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='2016' max='25200' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [ 2016/25200 00:26 < 05:03, 76.43 it/s, Epoch 8/100]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.499100</td>\n",
       "      <td>0.705115</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.406400</td>\n",
       "      <td>0.689810</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.382100</td>\n",
       "      <td>0.681682</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.366600</td>\n",
       "      <td>0.681303</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>0.358800</td>\n",
       "      <td>0.684668</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>0.351900</td>\n",
       "      <td>0.690913</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>0.347900</td>\n",
       "      <td>0.690822</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>0.347600</td>\n",
       "      <td>0.701287</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=2016, training_loss=0.3825622891622876, metrics={'train_runtime': 28.0282, 'train_samples_per_second': 28660.44, 'train_steps_per_second': 899.095, 'total_flos': 597127745765376.0, 'train_loss': 0.3825622891622876, 'epoch': 8.0})"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_args = TrainingArguments(\n",
    "    output_dir=\"./checkpoint/patchtsmixer/direct/train/output/\",\n",
    "    overwrite_output_dir=True,\n",
    "    learning_rate=0.0001,\n",
    "    num_train_epochs=100,\n",
    "    do_eval=True,\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    dataloader_num_workers=num_workers,\n",
    "    report_to=\"tensorboard\",\n",
    "    save_strategy=\"epoch\",\n",
    "    logging_strategy=\"epoch\",\n",
    "    save_total_limit=3,\n",
    "    logging_dir=\"./checkpoint/patchtsmixer/direct/train/logs/\",  # Make sure to specify a logging directory\n",
    "    load_best_model_at_end=True,  # Load the best model when training ends\n",
    "    metric_for_best_model=\"eval_loss\",  # Metric to monitor for early stopping\n",
    "    greater_is_better=False,  # For loss\n",
    "    label_names=[\"future_values\"],\n",
    ")\n",
    "\n",
    "# Create a new early stopping callback with faster convergence properties\n",
    "early_stopping_callback = EarlyStoppingCallback(\n",
    "    early_stopping_patience=5,  # Number of epochs with no improvement after which to stop\n",
    "    early_stopping_threshold=0.001,  # Minimum improvement required to consider as improvement\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=train_args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=valid_dataset,\n",
    "    callbacks=[early_stopping_callback],\n",
    ")\n",
    "\n",
    "print(\"\\n\\nDoing forecasting training on Etth1/train\")\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e06f3931-6b5f-450e-b17d-360ae3984e67",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='88' max='88' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [88/88 00:00]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'eval_loss': 0.367517352104187,\n",
       " 'eval_runtime': 0.6978,\n",
       " 'eval_samples_per_second': 3991.024,\n",
       " 'eval_steps_per_second': 126.108,\n",
       " 'epoch': 8.0}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.evaluate(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e206fac-39ef-4101-8d56-d8fd2c76b72f",
   "metadata": {},
   "source": [
    "## If we want to train from scratch for a few specific forecast channels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1e1a42c9-a3dd-4a54-ab20-bdd10a1903ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "forecast_channel_indices = [\n",
    "    -4,\n",
    "    -1,\n",
    "]  # add the channel indices (i.e., the column number) for which the model should forecast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d7c7dbbc-2952-4a8b-ac4d-a05257c7afa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = PatchTSMixerConfig(\n",
    "    context_length=context_length,\n",
    "    prediction_length=forecast_horizon,\n",
    "    patch_length=patch_length,\n",
    "    num_input_channels=len(forecast_columns),\n",
    "    patch_stride=patch_length,\n",
    "    d_model=48,\n",
    "    num_layers=3,\n",
    "    expansion_factor=3,\n",
    "    dropout=0.5,\n",
    "    head_dropout=0.7,\n",
    "    mode=\"common_channel\",\n",
    "    scaling=\"std\",\n",
    "    prediction_channel_indices=forecast_channel_indices,\n",
    ")\n",
    "model = PatchTSMixerForPrediction(config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "63aa0a38-c56f-4abd-92a6-9e9e5daadf51",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Doing forecasting training on Etth1/train\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='4284' max='25200' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [ 4284/25200 00:59 < 04:51, 71.82 it/s, Epoch 17/100]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.275300</td>\n",
       "      <td>0.496316</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.231200</td>\n",
       "      <td>0.485542</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.218200</td>\n",
       "      <td>0.478069</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.209900</td>\n",
       "      <td>0.470516</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>0.206400</td>\n",
       "      <td>0.477010</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>0.202600</td>\n",
       "      <td>0.474555</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>0.200600</td>\n",
       "      <td>0.474283</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>0.198300</td>\n",
       "      <td>0.472296</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>0.196000</td>\n",
       "      <td>0.464579</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>0.194800</td>\n",
       "      <td>0.467563</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>11</td>\n",
       "      <td>0.193500</td>\n",
       "      <td>0.467156</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>12</td>\n",
       "      <td>0.192500</td>\n",
       "      <td>0.456910</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>13</td>\n",
       "      <td>0.190300</td>\n",
       "      <td>0.464192</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>14</td>\n",
       "      <td>0.189500</td>\n",
       "      <td>0.465797</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>15</td>\n",
       "      <td>0.188200</td>\n",
       "      <td>0.468837</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>16</td>\n",
       "      <td>0.187800</td>\n",
       "      <td>0.471794</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>17</td>\n",
       "      <td>0.186300</td>\n",
       "      <td>0.473285</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=4284, training_loss=0.20359806520264356, metrics={'train_runtime': 60.2954, 'train_samples_per_second': 13322.735, 'train_steps_per_second': 417.942, 'total_flos': 1268896459751424.0, 'train_loss': 0.20359806520264356, 'epoch': 17.0})"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=train_args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=valid_dataset,\n",
    "    callbacks=[early_stopping_callback],\n",
    ")\n",
    "\n",
    "print(\"\\n\\nDoing forecasting training on Etth1/train\")\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ec5de472-7e2e-4bb2-8cdb-6371d0a33ce1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'eval_loss': 0.1160622164607048,\n",
       " 'eval_runtime': 0.7379,\n",
       " 'eval_samples_per_second': 3774.245,\n",
       " 'eval_steps_per_second': 119.258,\n",
       " 'epoch': 17.0}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.evaluate(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b2e8f0a-8367-4c10-84bd-a72e8f21ccc4",
   "metadata": {},
   "source": [
    "#### Sanity check: Compute number of forecasting channels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a1cdf078-10e8-4788-b7e4-3d1640ab8f8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = trainer.predict(test_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "4db1b83f-7381-4fe6-97e4-42ee210ab71d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2785, 96, 2)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output.predictions[0].shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}