MlpTSMixer-hyperparameter_tuning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "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.multivariate_grouper import MultivariateGrouper\n",
    "from gluonts.dataset.repository.datasets import dataset_recipes, get_dataset\n",
    "from gluonts.evaluation.backtest import make_evaluation_predictions\n",
    "from gluonts.evaluation import MultivariateEvaluator\n",
    "\n",
    "from MlpTSMixer import MlpTSMixerEstimator\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": null,
   "id": "944676b0-2a9f-4301-92e6-f382f5693639",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class MlpTSMixerObjective:\n",
    "    def __init__(\n",
    "        self,\n",
    "        dataset,\n",
    "        train_grouper,\n",
    "        test_grouper,\n",
    "        metric_type=\"m_sum_mean_wQuantileLoss\",\n",
    "    ):\n",
    "        self.metric_type = metric_type\n",
    "        self.dataset = dataset\n",
    "        self.dataset_train = train_grouper(self.dataset.train)\n",
    "        self.dataset_test = test_grouper(self.dataset.test)\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",
    "            \"depth\": trial.suggest_int(\"depth\", 2, 16, 4),\n",
    "            \"dim\": trial.suggest_int(\"dim\", 16, 64, 16),\n",
    "            \"patch_size\": trial.suggest_int(\"patch_size\", 2, 16, 4),\n",
    "            \"expansion_factor\": trial.suggest_int(\"expansion_factor\", 2, 8, 2),\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",
    "        estimator = estimator = MlpTSMixerEstimator(\n",
    "            # distr_output=StudentTOutput(dim=int(dataset.metadata.feat_static_cat[0].cardinality)),\n",
    "            input_size=int(self.dataset.metadata.feat_static_cat[0].cardinality),\n",
    "            prediction_length=self.dataset.metadata.prediction_length,\n",
    "            context_length=params[\"context_length\"],\n",
    "            freq=self.dataset.metadata.freq,\n",
    "            scaling=\"std\",\n",
    "            depth=params[\"depth\"],\n",
    "            patch_size=(params[\"patch_size\"], params[\"patch_size\"]),\n",
    "            dim=params[\"dim\"],\n",
    "            batch_size=params[\"batch_size\"],\n",
    "            num_batches_per_epoch=100,\n",
    "            patch_reverse_mapping_layer=\"mlp\",\n",
    "            trainer_kwargs=dict(accelerator=\"cuda\", max_epochs=30),\n",
    "        )\n",
    "        predictor = estimator.train(\n",
    "            training_data=self.dataset_train, num_workers=8, shuffle_buffer_length=1024\n",
    "        )\n",
    "\n",
    "        forecast_it, ts_it = make_evaluation_predictions(\n",
    "            dataset=self.dataset_test, predictor=predictor, num_samples=100\n",
    "        )\n",
    "        forecasts = list(forecast_it)\n",
    "        tss = list(ts_it)\n",
    "        evaluator = MultivariateEvaluator(\n",
    "            quantiles=(np.arange(20) / 20.0)[1:], target_agg_funcs={\"sum\": np.sum}\n",
    "        )\n",
    "        agg_metrics, _ = evaluator(iter(tss), iter(forecasts))\n",
    "        return agg_metrics[self.metric_type]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6c1da22-755e-44d5-86b6-10465d8c25e8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataset = get_dataset(\"solar_nips\", regenerate=False)\n",
    "train_grouper = MultivariateGrouper(\n",
    "    max_target_dim=int(dataset.metadata.feat_static_cat[0].cardinality)\n",
    ")\n",
    "\n",
    "test_grouper = MultivariateGrouper(\n",
    "    num_test_dates=int(len(dataset.test) / len(dataset.train)),\n",
    "    max_target_dim=int(dataset.metadata.feat_static_cat[0].cardinality),\n",
    ")\n",
    "dataset_train = train_grouper(dataset.train)\n",
    "dataset_test = test_grouper(dataset.test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ded5c21-ea78-4f99-98e3-1573f5abfbf2",
   "metadata": {
    "tags": []
   },
   "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(MlpTSMixerObjective(dataset, train_grouper, test_grouper), 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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9819ac62-f056-4dcb-814f-075f8ce15bac",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}