tutorials/Model Selection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "60611bed",
   "metadata": {},
   "source": [
    "# Model Selection in TSGM\n",
    "This is a simple example of model selection through hyperparameter optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6220fcdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import numpy as np\n",
    "import functools\n",
    "import sklearn\n",
    "import copy\n",
    "import sklearn.model_selection\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "\n",
    "import tsgm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0290b339",
   "metadata": {},
   "source": [
    "## 0. Install optuna\n",
    "\n",
    "Let's first install Optuna."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cf29099",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "!{sys.executable} -m pip install optuna\n",
    "import optuna"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a375b04e",
   "metadata": {},
   "source": [
    "## 1. Load data\n",
    "We are using a small dataset generated by the `tsgm.utils.gen_sine_dataset` function. We scale the features using `tsgm.utils.TSFeatureWiseScaler` to ensure each feature falls within the range of $[0, 1]$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3da6d8eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data():\n",
    "    data = tsgm.utils.gen_sine_dataset(10000, 24, 5)\n",
    "    scaler = tsgm.utils.TSFeatureWiseScaler()        \n",
    "    scaled_data = scaler.fit_transform(data)\n",
    "    return scaled_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b76dd508",
   "metadata": {},
   "source": [
    "## 2. Define the optimization problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ecc6c87",
   "metadata": {},
   "outputs": [],
   "source": [
    "# we aim at minimizing the discrepancy metric defined in next cell\n",
    "study = optuna.create_study(direction=\"minimize\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0099c91e",
   "metadata": {},
   "outputs": [],
   "source": [
    "metric_to_optimize = tsgm.metrics.metrics.DistanceMetric(\n",
    "            statistics=[\n",
    "                functools.partial(tsgm.metrics.statistics.axis_max_s, axis=None),\n",
    "                functools.partial(tsgm.metrics.statistics.axis_min_s, axis=None),\n",
    "                functools.partial(tsgm.metrics.statistics.axis_max_s, axis=1),\n",
    "                functools.partial(tsgm.metrics.statistics.axis_min_s, axis=1),\n",
    "            ],\n",
    "            discrepancy=lambda x, y: np.linalg.norm(x - y),\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bbab28c",
   "metadata": {},
   "source": [
    "## 3. Define the search space for the optimizer\n",
    "We can optimize the choice of the optimizer and hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4d398a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# optimizers and the search space for the hyperparameters\n",
    "def _create_optimizer(trial):\n",
    "    # optimize the choice of optimizers as well as their parameters\n",
    "    kwargs = {}\n",
    "    optimizer_options = [\"RMSprop\", \"Adam\", \"SGD\"]\n",
    "    optimizer_selected = trial.suggest_categorical(\"optimizer\", optimizer_options)\n",
    "    if optimizer_selected == \"RMSprop\":\n",
    "        kwargs[\"learning_rate\"] = trial.suggest_float(\n",
    "            \"rmsprop_learning_rate\", 1e-5, 1e-1, log=True\n",
    "        )\n",
    "        kwargs[\"momentum\"] = trial.suggest_float(\n",
    "            \"rmsprop_momentum\", 1e-5, 1e-1, log=True\n",
    "        )\n",
    "    elif optimizer_selected == \"Adam\":\n",
    "        kwargs[\"learning_rate\"] = trial.suggest_float(\n",
    "            \"adam_learning_rate\", 1e-5, 1e-1, log=True\n",
    "        )\n",
    "    elif optimizer_selected == \"SGD\":\n",
    "        kwargs[\"learning_rate\"] = trial.suggest_float(\n",
    "            \"sgd_opt_learning_rate\", 1e-5, 1e-1, log=True\n",
    "        )\n",
    "        kwargs[\"momentum\"] = trial.suggest_float(\n",
    "            \"sgd_opt_momentum\", 1e-5, 1e-1, log=True\n",
    "        )\n",
    "\n",
    "    optimizer = getattr(tf.optimizers, optimizer_selected)(**kwargs)\n",
    "    return optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70786d6c",
   "metadata": {},
   "source": [
    "## 4. Define the objective function\n",
    "In the objective function, we load the data and use them to train a TimeGAN model (`tsgm.models.timeGAN.TimeGAN`) while changing its parameters. We use the fitted TimeGAN model to generate synthetic samples, and finally use them to compute the value of the metric we want to optimize. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32bb7160",
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective(trial):\n",
    "    # Get data\n",
    "    train_data = get_data()\n",
    "\n",
    "    # Define the search space\n",
    "    n_layers = trial.suggest_int(name=\"n_layers\", low=1, high=10)\n",
    "    num_hidden = trial.suggest_int(name=\"num_hidden\", low=4, high=128, log=True)\n",
    "    \n",
    "    # Build TimeGAN model\n",
    "    model = tsgm.models.timeGAN.TimeGAN(\n",
    "        seq_len=24,\n",
    "        module=\"gru\",\n",
    "        hidden_dim=num_hidden,\n",
    "        n_features=5,\n",
    "        n_layers=n_layers,\n",
    "        batch_size=256,\n",
    "        gamma=1.0,\n",
    "    )\n",
    "    # get optimizer\n",
    "    optimizer = _create_optimizer(trial)\n",
    "    \n",
    "    # compile model\n",
    "    model.compile(optimizer)\n",
    "\n",
    "    # Training and validating\n",
    "    EPOCHS = 100\n",
    "    model.fit(data=train_data, epochs=EPOCHS)\n",
    "    \n",
    "    # Generate 10 samples of synthetic data\n",
    "    _y = model.generate(n_samples=10)\n",
    "    \n",
    "    # Evaluate them vs the first 10 samples of training data\n",
    "    objective_to_optimize = metric_to_optimize(_y, np.array(train_data[:10]))\n",
    "    \n",
    "    # Return last validation score\n",
    "    return objective_to_optimize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae0de98f",
   "metadata": {},
   "outputs": [],
   "source": [
    "study.optimize(objective, n_trials=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76ceca1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters corresponding to best trial\n",
    "study.best_trial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14ccbf23",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}