LagTST-hyperparameter_tuning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "86ff733d-8ea7-4d77-8d2a-cd329ab8f385",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import torch\n",
    "\n",
    "from gluonts.dataset.repository.datasets import get_dataset\n",
    "from gluonts.dataset.common import ListDataset\n",
    "from gluonts.evaluation import make_evaluation_predictions, Evaluator\n",
    "\n",
    "# from gluonts.torch.model.lag_tst.estimator import LagTSTEstimator\n",
    "from gluonts.torch.distributions import NegativeBinomialOutput\n",
    "from gluonts.torch.modules.loss import NegativeLogLikelihood\n",
    "\n",
    "from LagTST import LagTSTEstimator\n",
    "\n",
    "# from pts.modules import StudentTOutput\n",
    "\n",
    "import random\n",
    "import numpy as np\n",
    "import time\n",
    "import optuna\n",
    "from optuna.samplers import TPESampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "944676b0-2a9f-4301-92e6-f382f5693639",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class LagTSTObjective:\n",
    "    def __init__(self, dataset, metric_type=\"mean_wQuantileLoss\"):\n",
    "        self.dataset = dataset\n",
    "        self.metric_type = metric_type\n",
    "\n",
    "    def get_params(self, trial) -> dict:\n",
    "        return {\n",
    "            \"context_length\": trial.suggest_int(\n",
    "                \"context_length\",\n",
    "                dataset.metadata.prediction_length,\n",
    "                dataset.metadata.prediction_length * 5,\n",
    "                1,\n",
    "            ),\n",
    "            \"batch_size\": trial.suggest_int(\"batch_size\", 32, 256, 32),\n",
    "            \"d_model\": trial.suggest_int(\"d_model\", 16, 64, 16),\n",
    "            \"dim\": trial.suggest_int(\"dim\", 16, 64, 16),\n",
    "            \"patch_size\": trial.suggest_int(\"patch_size\", 2, 16, 4),\n",
    "            \"kernel_size\": trial.suggest_int(\"kernel_size\", 9, 18, 3),\n",
    "        }\n",
    "\n",
    "    def __call__(self, trial):\n",
    "        params = self.get_params(trial)\n",
    "\n",
    "        estimator = LagTSTEstimator(\n",
    "            prediction_length=self.dataset.metadata.prediction_length,\n",
    "            context_length=params[\"context_length\"],\n",
    "            freq=dataset.metadata.freq,\n",
    "            scaling=\"std\",\n",
    "            # distr_output=NegativeBinomialOutput(),\n",
    "            # loss=NegativeLogLikelihood(beta=0.2),\n",
    "            d_model=params[\"d_model\"],\n",
    "            dim_feedforward=params[\"dim\"],\n",
    "            batch_size=params[\"batch_size\"],\n",
    "            patch_reverse_mapping_layer=\"mlp\",\n",
    "            num_batches_per_epoch=100,\n",
    "            trainer_kwargs=dict(accelerator=\"gpu\", max_epochs=30),\n",
    "        )\n",
    "        predictor = estimator.train(\n",
    "            training_data=self.dataset.train,\n",
    "            cache_data=True,\n",
    "            shuffle_buffer_length=1024,\n",
    "            validation_data=self.dataset.test,\n",
    "        )\n",
    "\n",
    "        forecast_it, ts_it = make_evaluation_predictions(\n",
    "            dataset=dataset.test,\n",
    "            predictor=predictor,\n",
    "        )\n",
    "        forecasts = list(forecast_it)\n",
    "        # if layer == layers[0]:\n",
    "        tss = list(ts_it)\n",
    "        evaluator = Evaluator()\n",
    "        agg_metrics, _ = evaluator(iter(tss), iter(forecasts))\n",
    "        return agg_metrics[self.metric_type]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d6c1da22-755e-44d5-86b6-10465d8c25e8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataset = get_dataset(\n",
    "    \"solar-energy\", regenerate=False\n",
    ")  # dataset = get_dataset(\"electricity\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ded5c21-ea78-4f99-98e3-1573f5abfbf2",
   "metadata": {
    "tags": [],
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "seed = 42\n",
    "random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "start_time = time.time()\n",
    "sampler = TPESampler(seed=seed)\n",
    "study = optuna.create_study(sampler=sampler, direction=\"minimize\")\n",
    "study.optimize(LagTSTObjective(dataset), n_trials=10)\n",
    "\n",
    "print(\"Number of finished trials: {}\".format(len(study.trials)))\n",
    "\n",
    "print(\"Best trial:\")\n",
    "trial = study.best_trial\n",
    "\n",
    "print(\"  Value: {}\".format(trial.value))\n",
    "print(\"  Params: \")\n",
    "for key, value in trial.params.items():\n",
    "    print(\"    {}: {}\".format(key, value))\n",
    "print(time.time() - start_time)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py38",
   "language": "python",
   "name": "py38"
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}