diff --git a/notebooks/07.0-mc-optuna-find-hiddensize.ipynb b/notebooks/07.0-mc-optuna-find-hiddensize.ipynb new file mode 100644 index 0000000..d722151 --- /dev/null +++ b/notebooks/07.0-mc-optuna-find-hiddensize.ipynb @@ -0,0 +1,783 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "ExecuteTime": { + "end_time": "2020-10-10T01:25:12.788851Z", + "start_time": "2020-10-10T01:25:12.783398Z" + } + }, + "source": [ + "# Sequence to Sequence Models for Timeseries Regression\n", + "\n", + "\n", + "In this notebook we are going to find the optimal hidden_size for a model vs a dataset. We will use pytorch lightning and optuna." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "ExecuteTime": { + "end_time": "2020-11-08T02:52:04.993589Z", + "start_time": "2020-11-08T02:52:04.569061Z" + } + }, + "outputs": [], + "source": [ + "# OPTIONAL: Load the \"autoreload\" extension so that code can change. But blacklist large modules\n", + "%load_ext autoreload\n", + "%autoreload 2\n", + "%aimport -pandas\n", + "%aimport -torch\n", + "%aimport -numpy\n", + "%aimport -matplotlib\n", + "%aimport -dask\n", + "%aimport -tqdm\n", + "%matplotlib inline" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": { + "ExecuteTime": { + "end_time": "2020-11-08T02:52:06.671206Z", + "start_time": "2020-11-08T02:52:04.998087Z" + }, + "lines_to_next_cell": 0 + }, + "outputs": [], + "source": [ + "# Imports\n", + "import torch\n", + "from torch import nn, optim\n", + "from torch.nn import functional as F\n", + "from torch.autograd import Variable\n", + "import torch\n", + "import torch.utils.data\n", + "\n", + "import xarray as xr\n", + "import pandas as pd\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "\n", + "from pathlib import Path\n", + "from tqdm.auto import tqdm\n", + "\n", + "import pytorch_lightning as pl" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "ExecuteTime": { + "end_time": "2020-11-08T02:52:06.707927Z", + "start_time": "2020-11-08T02:52:06.674890Z" + } + }, + "outputs": [], + "source": [ + "from seq2seq_time.data.dataset import Seq2SeqDataSet, Seq2SeqDataSets\n", + "from seq2seq_time.predict import predict, predict_multi\n", + "from seq2seq_time.util import dset_to_nc" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": { + "ExecuteTime": { + "end_time": "2020-11-08T02:52:06.745323Z", + "start_time": "2020-11-08T02:52:06.711604Z" + } + }, + "outputs": [], + "source": [ + "import logging\n", + "import warnings\n", + "import seq2seq_time.silence \n", + "warnings.simplefilter('once')\n", + "warnings.simplefilter(action='ignore', category=FutureWarning)\n", + "warnings.simplefilter(action='ignore', category=DeprecationWarning)\n", + "warnings.filterwarnings('ignore', 'Consider increasing the value of the `num_workers` argument', UserWarning)\n", + "warnings.filterwarnings('ignore', 'Your val_dataloader has `shuffle=True`', UserWarning)\n", + "\n", + "from pytorch_lightning import _logger as log\n", + "log.setLevel(logging.WARN)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "ExecuteTime": { + "end_time": "2020-10-10T01:28:32.492160Z", + "start_time": "2020-10-10T01:28:32.488140Z" + } + }, + "source": [ + "## Parameters" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": { + "ExecuteTime": { + "end_time": "2020-11-08T02:52:06.843841Z", + "start_time": "2020-11-08T02:52:06.751591Z" + }, + "lines_to_next_cell": 0 + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "using cuda\n", + "20201108-095004\n" + ] + }, + { + "data": { + "text/plain": [ + "96" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "print(f'using {device}')\n", + "\n", + "timestamp = '20201108-095004'\n", + "print(timestamp)\n", + "window_past = 48*2\n", + "window_future = 48\n", + "batch_size = 64\n", + "num_workers = 5\n", + "datasets_root = Path('../data/processed/')\n", + "window_past" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "end_time": "2020-11-01T23:28:09.504323Z", + "start_time": "2020-11-01T23:28:09.453546Z" + }, + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Datasets\n", + "\n", + "From easy to hard, these dataset show different challenges, all of them with more than 20k datapoints and with a regression output. See the 00.01 notebook for more details, and the code for more information.\n", + "\n", + "Some such as MetroInterstateTraffic are easier, some are periodic such as BejingPM25, some are conditional on inputs such as GasSensor, and some are noisy and periodic like IMOSCurrentsVel" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": { + "ExecuteTime": { + "end_time": "2020-11-08T02:52:07.298057Z", + "start_time": "2020-11-08T02:52:06.850596Z" + }, + "lines_to_next_cell": 0 + }, + "outputs": [ + { + "data": { + "text/plain": [ + "[seq2seq_time.data.data.GasSensor,\n", + " seq2seq_time.data.data.IMOSCurrentsVel,\n", + " seq2seq_time.data.data.AppliancesEnergyPrediction,\n", + " seq2seq_time.data.data.BejingPM25,\n", + " seq2seq_time.data.data.MetroInterstateTraffic]" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from seq2seq_time.data.data import IMOSCurrentsVel, AppliancesEnergyPrediction, BejingPM25, GasSensor, MetroInterstateTraffic\n", + "datasets = [GasSensor, IMOSCurrentsVel, AppliancesEnergyPrediction, BejingPM25, MetroInterstateTraffic]\n", + "datasets" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Lightning\n", + "\n", + "We will use pytorch lightning to handle all the training scaffolding. We have a common pytorch lightning class that takes in the model and defines training steps and logging." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": { + "ExecuteTime": { + "end_time": "2020-11-08T02:52:07.347557Z", + "start_time": "2020-11-08T02:52:07.301918Z" + } + }, + "outputs": [], + "source": [ + "import pytorch_lightning as pl\n", + "\n", + "class PL_MODEL(pl.LightningModule):\n", + " def __init__(self, model, lr=3e-4, patience=None, weight_decay=0):\n", + " super().__init__()\n", + " self._model = model\n", + " self.lr = lr\n", + " self.patience = patience\n", + " self.weight_decay = weight_decay\n", + "\n", + " def forward(self, x_past, y_past, x_future, y_future=None):\n", + " \"\"\"Eval/Predict\"\"\"\n", + " y_dist, extra = self._model(x_past, y_past, x_future, y_future)\n", + " return y_dist, extra\n", + "\n", + " def training_step(self, batch, batch_idx, phase='train'):\n", + " x_past, y_past, x_future, y_future = batch\n", + " y_dist, extra = self.forward(*batch)\n", + " loss = -y_dist.log_prob(y_future).mean()\n", + " self.log_dict({f'loss/{phase}':loss})\n", + " if ('loss' in extra) and (phase=='train'):\n", + " # some models have a special loss\n", + " loss = extra['loss']\n", + " self.log_dict({f'model_loss/{phase}':loss})\n", + " return loss\n", + "\n", + " def validation_step(self, batch, batch_idx):\n", + " return self.training_step(batch, batch_idx, phase='val')\n", + " \n", + " def test_step(self, batch, batch_idx):\n", + " return self.training_step(batch, batch_idx, phase='test')\n", + " \n", + " def configure_optimizers(self):\n", + " optim = torch.optim.AdamW(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)\n", + " scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n", + " optim,\n", + " patience=self.patience,\n", + " verbose=False,\n", + " min_lr=1e-7,\n", + " ) if self.patience else None\n", + " return {'optimizer': optim, 'lr_scheduler': scheduler, 'monitor': 'loss/val'}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "start_time": "2020-11-08T02:52:04.592Z" + }, + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [ + "from torch.utils.data import DataLoader\n", + "from pytorch_lightning.loggers import CSVLogger, WandbLogger, TensorBoardLogger, TestTubeLogger\n", + "from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n", + "from pytorch_lightning.callbacks import LearningRateMonitor" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Models" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "start_time": "2020-11-08T02:52:04.595Z" + }, + "lines_to_end_of_cell_marker": 2, + "lines_to_next_cell": 0 + }, + "outputs": [], + "source": [ + "from seq2seq_time.models.baseline import BaselineLast, BaselineMean\n", + "from seq2seq_time.models.lstm_seq2seq import LSTMSeq2Seq\n", + "from seq2seq_time.models.lstm import LSTM\n", + "from seq2seq_time.models.transformer import Transformer\n", + "from seq2seq_time.models.transformer_seq2seq import TransformerSeq2Seq\n", + "from seq2seq_time.models.neural_process import RANP\n", + "from seq2seq_time.models.transformer_process import TransformerProcess\n", + "from seq2seq_time.models.tcn import TCNSeq\n", + "from seq2seq_time.models.inceptiontime import InceptionTimeSeq\n", + "from seq2seq_time.models.xattention import CrossAttention" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "start_time": "2020-11-08T02:52:04.599Z" + }, + "lines_to_next_cell": 0 + }, + "outputs": [], + "source": [ + "import gc\n", + "\n", + "def free_mem():\n", + " gc.collect()\n", + " torch.cuda.empty_cache()\n", + " gc.collect()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "end_time": "2020-11-02T06:10:41.904480Z", + "start_time": "2020-11-02T06:10:41.848613Z" + }, + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "start_time": "2020-11-08T02:52:04.605Z" + }, + "lines_to_next_cell": 0 + }, + "outputs": [], + "source": [ + "# PARAMS: model\n", + "dropout=0.0\n", + "layers=6\n", + "nhead=4\n", + "\n", + "models = [\n", + "# lambda xs, ys: BaselineLast(),\n", + "# lambda xs, ys, hidden_size: BaselineMean(),\n", + " lambda xs, ys, hidden_size, layers: Transformer(xs,\n", + " ys,\n", + " attention_dropout=dropout,\n", + " nhead=nhead,\n", + " nlayers=layers,\n", + " hidden_size=hidden_size),\n", + "\n", + " lambda xs, ys, hidden_size, layers:TransformerProcess(xs,\n", + " ys, hidden_size=hidden_size, nhead=nhead,\n", + " latent_dim=hidden_size//2, dropout=dropout,\n", + " nlayers=layers),\n", + " lambda xs, ys, hidden_size, layers:TCNSeq(xs, ys, hidden_size=hidden_size, nlayers=layers, dropout=dropout, kernel_size=2),\n", + " lambda xs, ys, hidden_size, layers: RANP(xs,\n", + " ys, hidden_dim=hidden_size, dropout=dropout, \n", + " latent_dim=hidden_size//2, n_decoder_layers=layers, n_latent_encoder_layers=layers, n_det_encoder_layers=layers),\n", + " lambda xs, ys, hidden_size, layers: TransformerSeq2Seq(xs,\n", + " ys,\n", + " hidden_size=hidden_size,\n", + " nhead=nhead,\n", + " nlayers=layers,\n", + " attention_dropout=dropout\n", + " ),\n", + " lambda xs, ys, hidden_size, layers: LSTM(xs,\n", + " ys,\n", + " hidden_size=hidden_size,\n", + " lstm_layers=layers//2,\n", + " lstm_dropout=dropout),\n", + " lambda xs, ys, hidden_size, layers: LSTMSeq2Seq(xs,\n", + " ys,\n", + " hidden_size=hidden_size,\n", + " lstm_layers=layers//2,\n", + " lstm_dropout=dropout),\n", + " lambda xs, ys, hidden_size, layers: CrossAttention(xs,\n", + " ys,\n", + " nlayers=layers,\n", + " hidden_size=hidden_size,),\n", + " lambda xs, ys, hidden_size, layers: InceptionTimeSeq(xs,\n", + " ys,\n", + " kernel_size=96,\n", + " layers=layers//2,\n", + " hidden_size=hidden_size,\n", + " bottleneck=hidden_size//4)\n", + "\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "start_time": "2020-11-08T02:52:04.608Z" + } + }, + "outputs": [], + "source": [ + "# DEBUG: sanity check\n", + "\n", + "for Dataset in datasets:\n", + " dataset_name = Dataset.__name__\n", + " dataset = Dataset(datasets_root)\n", + " ds_train, ds_val, ds_test = dataset.to_datasets(window_past=window_past,\n", + " window_future=window_future)\n", + "\n", + " # Init data\n", + " x_past, y_past, x_future, y_future = ds_train.get_rows(10)\n", + " xs = x_past.shape[-1]\n", + " ys = y_future.shape[-1]\n", + "\n", + " # Loaders\n", + " dl_train = DataLoader(ds_train,\n", + " batch_size=batch_size,\n", + " shuffle=True,\n", + " pin_memory=num_workers == 0,\n", + " num_workers=num_workers)\n", + " dl_val = DataLoader(ds_val,\n", + " shuffle=True,\n", + " batch_size=batch_size,\n", + " num_workers=num_workers)\n", + "\n", + " for m_fn in models:\n", + " free_mem()\n", + " pt_model = m_fn(xs, ys, 8, 4)\n", + " model_name = type(pt_model).__name__\n", + " print(timestamp, dataset_name, model_name)\n", + "\n", + " # Wrap in lightning\n", + " model = PL_MODEL(pt_model,\n", + " lr=3e-4\n", + " ).to(device)\n", + " trainer = pl.Trainer(\n", + " fast_dev_run=True,\n", + " # GPU\n", + " gpus=1,\n", + " amp_level='O1',\n", + " precision=16,\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "ExecuteTime": { + "end_time": "2020-11-01T07:30:40.569795Z", + "start_time": "2020-11-01T07:29:12.500374Z" + } + }, + "source": [ + "Lets summarize all models, and make sure they have a similar number of parameters" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "ExecuteTime": { + "end_time": "2020-10-23T23:36:11.052891Z", + "start_time": "2020-10-23T23:36:11.048874Z" + } + }, + "source": [ + "## Train" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "start_time": "2020-11-08T02:52:04.612Z" + }, + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [ + "from collections import defaultdict\n", + "from seq2seq_time.metrics import rmse, smape" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "start_time": "2020-11-08T02:52:04.617Z" + }, + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [ + "max_iters=20000" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "start_time": "2020-11-08T02:52:04.620Z" + } + }, + "outputs": [], + "source": [ + "tensorboard_dir = Path(f\"../outputs/{timestamp}\").resolve()\n", + "print(f'For tensorboard run:\\ntensorboard --logdir=\"{tensorboard_dir}\"')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "end_time": "2020-11-08T01:38:37.371764Z", + "start_time": "2020-11-08T01:38:37.315240Z" + } + }, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "start_time": "2020-11-08T02:52:04.626Z" + } + }, + "outputs": [], + "source": [ + "class MetricsCallback(pl.Callback):\n", + " \"\"\"PyTorch Lightning metric callback.\"\"\"\n", + "\n", + " def __init__(self):\n", + " super().__init__()\n", + " self.metrics = []\n", + "\n", + " def on_validation_end(self, trainer, pl_module):\n", + " self.metrics.append(trainer.callback_metrics)\n", + "\n", + "def objective(trial):\n", + " # sample\n", + " hidden_size_exp = trial.suggest_int(\"hidden_size_exp\", 2, 8)\n", + " hidden_size = 2**hidden_size_exp\n", + " \n", + " layers = trial.suggest_int(\"layers\", 2, 12)\n", + " \n", + " # Load model\n", + " pt_model = m_fn(xs, ys, hidden_size, layers)\n", + " model_name = type(pt_model).__name__\n", + " \n", + " # Wrap in lightning\n", + " patience = 2\n", + " model = PL_MODEL(pt_model,\n", + " lr=3e-4, patience=patience,\n", + " weight_decay=4e-5\n", + " ).to(device)\n", + "\n", + " \n", + " # The default logger in PyTorch Lightning writes to event files to be consumed by\n", + " # TensorBoard. We don't use any logger here as it requires us to implement several abstract\n", + " # methods. Instead we setup a simple callback, that saves metrics from each validation step.\n", + "# metrics_callback = MetricsCallback()\n", + " \n", + " save_dir = f\"../outputs/{timestamp}/{dataset_name}_{model_name}/{trial.number}\"\n", + " Path(save_dir).mkdir(exist_ok=True, parents=True)\n", + " trainer = pl.Trainer(\n", + " # Training length\n", + " min_epochs=2,\n", + " max_epochs=100,\n", + " limit_train_batches=max_iters//batch_size,\n", + " limit_val_batches=max_iters//batch_size//5,\n", + " # Misc\n", + " gradient_clip_val=20,\n", + " terminate_on_nan=True,\n", + " # GPU\n", + " gpus=1,\n", + " amp_level='O1',\n", + " precision=16,\n", + " # Callbacks\n", + " default_root_dir=save_dir,\n", + " logger=False,\n", + " callbacks=[\n", + "# metrics_callback, \n", + " EarlyStopping(monitor='loss/val', patience=patience * 2),\n", + " PyTorchLightningPruningCallback(trial, monitor=\"loss/val\")],\n", + " )\n", + " trainer.fit(model, dl_train, dl_val)\n", + " \n", + " # Run on all val data, using test mode\n", + " r = trainer.test(model, test_dataloader=dl_val, verbose=False)\n", + " return r[0]['loss/test']" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "end_time": "2020-11-08T02:45:44.106583Z", + "start_time": "2020-11-08T02:45:44.050637Z" + } + }, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "start_time": "2020-11-08T02:52:04.631Z" + } + }, + "outputs": [], + "source": [ + "import optuna\n", + "from optuna.integration import PyTorchLightningPruningCallback" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "start_time": "2020-11-08T02:52:04.634Z" + }, + "lines_to_next_cell": 0, + "scrolled": true + }, + "outputs": [], + "source": [ + "Path(f\"../outputs/{timestamp}\").mkdir(exist_ok=True)\n", + "results = defaultdict(dict)\n", + "for Dataset in tqdm(datasets, desc='datasets'):\n", + " dataset_name = Dataset.__name__\n", + " dataset = Dataset(datasets_root)\n", + " ds_train, ds_val, ds_test = dataset.to_datasets(window_past=window_past,\n", + " window_future=window_future)\n", + "\n", + " # Init data\n", + " x_past, y_past, x_future, y_future = ds_train.get_rows(10)\n", + " xs = x_past.shape[-1]\n", + " ys = y_future.shape[-1]\n", + "\n", + " # Loaders\n", + " dl_train = DataLoader(ds_train,\n", + " batch_size=batch_size,\n", + " shuffle=True,\n", + " pin_memory=num_workers == 0,\n", + " num_workers=num_workers)\n", + " dl_val = DataLoader(ds_val,\n", + " shuffle=False,\n", + " batch_size=batch_size,\n", + " num_workers=num_workers)\n", + "\n", + " for i, m_fn in enumerate(tqdm(models, desc=f'models ({dataset_name})')):\n", + " try:\n", + " model_name = type(m_fn(8, 8, 8, 2)).__name__\n", + " free_mem()\n", + " study_name = f'{timestamp}_{dataset_name}-{model_name}'\n", + " \n", + " storage = f\"sqlite:///../outputs/{timestamp}/optuna.db\"\n", + " pruner = optuna.pruners.MedianPruner()\n", + " study = optuna.create_study(storage=storage, \n", + " study_name=study_name, \n", + " pruner=pruner,\n", + " load_if_exists=True)\n", + " study.optimize(objective, n_trials=100, timeout=60*60)\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", + "\n", + " print(\" Params: \")\n", + " for key, value in trial.params.items():\n", + " print(\" {}: {}\".format(key, value))\n", + " \n", + " except Exception as e:\n", + " logging.exception('failed to run model')\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "@webio": { + "lastCommId": null, + "lastKernelId": null + }, + "jupytext": { + "encoding": "# -*- coding: utf-8 -*-", + "formats": "ipynb,py:light" + }, + "kernelspec": { + "display_name": "seq2seq-time", + "language": "python", + "name": "seq2seq-time" + }, + "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.7.8" + }, + "toc": { + "base_numbering": 1, + "nav_menu": {}, + "number_sections": true, + "sideBar": true, + "skip_h1_title": false, + "title_cell": "Table of Contents", + "title_sidebar": "Contents", + "toc_cell": false, + "toc_position": { + "height": "calc(100% - 180px)", + "left": "10px", + "top": "150px", + "width": "209.162px" + }, + "toc_section_display": true, + "toc_window_display": true + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/notebooks/07.0-mc-optuna-find-hiddensize.py b/notebooks/07.0-mc-optuna-find-hiddensize.py new file mode 100644 index 0000000..bfaed89 --- /dev/null +++ b/notebooks/07.0-mc-optuna-find-hiddensize.py @@ -0,0 +1,410 @@ +# -*- coding: utf-8 -*- +# --- +# jupyter: +# jupytext: +# formats: ipynb,py:light +# text_representation: +# extension: .py +# format_name: light +# format_version: '1.5' +# jupytext_version: 1.6.0 +# kernelspec: +# display_name: seq2seq-time +# language: python +# name: seq2seq-time +# --- + +# # Sequence to Sequence Models for Timeseries Regression +# +# +# In this notebook we are going to find the optimal hidden_size for a model vs a dataset. We will use pytorch lightning and optuna. + +# OPTIONAL: Load the "autoreload" extension so that code can change. But blacklist large modules +# %load_ext autoreload +# %autoreload 2 +# %aimport -pandas +# %aimport -torch +# %aimport -numpy +# %aimport -matplotlib +# %aimport -dask +# %aimport -tqdm +# %matplotlib inline + +# + +# Imports +import torch +from torch import nn, optim +from torch.nn import functional as F +from torch.autograd import Variable +import torch +import torch.utils.data + +import xarray as xr +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt + +from pathlib import Path +from tqdm.auto import tqdm + +import pytorch_lightning as pl +# - +from seq2seq_time.data.dataset import Seq2SeqDataSet, Seq2SeqDataSets +from seq2seq_time.predict import predict, predict_multi +from seq2seq_time.util import dset_to_nc + +# + +import logging +import warnings +import seq2seq_time.silence +warnings.simplefilter('once') +warnings.simplefilter(action='ignore', category=FutureWarning) +warnings.simplefilter(action='ignore', category=DeprecationWarning) +warnings.filterwarnings('ignore', 'Consider increasing the value of the `num_workers` argument', UserWarning) +warnings.filterwarnings('ignore', 'Your val_dataloader has `shuffle=True`', UserWarning) + +from pytorch_lightning import _logger as log +log.setLevel(logging.WARN) +# - + +# ## Parameters + +# + +device = "cuda" if torch.cuda.is_available() else "cpu" +print(f'using {device}') + +timestamp = '20201108-095004' +print(timestamp) +window_past = 48*2 +window_future = 48 +batch_size = 64 +num_workers = 5 +datasets_root = Path('../data/processed/') +window_past +# - + + + +# ## Datasets +# +# From easy to hard, these dataset show different challenges, all of them with more than 20k datapoints and with a regression output. See the 00.01 notebook for more details, and the code for more information. +# +# Some such as MetroInterstateTraffic are easier, some are periodic such as BejingPM25, some are conditional on inputs such as GasSensor, and some are noisy and periodic like IMOSCurrentsVel + +from seq2seq_time.data.data import IMOSCurrentsVel, AppliancesEnergyPrediction, BejingPM25, GasSensor, MetroInterstateTraffic +datasets = [GasSensor, IMOSCurrentsVel, AppliancesEnergyPrediction, BejingPM25, MetroInterstateTraffic] +datasets +# ## Lightning +# +# We will use pytorch lightning to handle all the training scaffolding. We have a common pytorch lightning class that takes in the model and defines training steps and logging. + +# + +import pytorch_lightning as pl + +class PL_MODEL(pl.LightningModule): + def __init__(self, model, lr=3e-4, patience=None, weight_decay=0): + super().__init__() + self._model = model + self.lr = lr + self.patience = patience + self.weight_decay = weight_decay + + def forward(self, x_past, y_past, x_future, y_future=None): + """Eval/Predict""" + y_dist, extra = self._model(x_past, y_past, x_future, y_future) + return y_dist, extra + + def training_step(self, batch, batch_idx, phase='train'): + x_past, y_past, x_future, y_future = batch + y_dist, extra = self.forward(*batch) + loss = -y_dist.log_prob(y_future).mean() + self.log_dict({f'loss/{phase}':loss}) + if ('loss' in extra) and (phase=='train'): + # some models have a special loss + loss = extra['loss'] + self.log_dict({f'model_loss/{phase}':loss}) + return loss + + def validation_step(self, batch, batch_idx): + return self.training_step(batch, batch_idx, phase='val') + + def test_step(self, batch, batch_idx): + return self.training_step(batch, batch_idx, phase='test') + + def configure_optimizers(self): + optim = torch.optim.AdamW(self.parameters(), lr=self.lr, weight_decay=self.weight_decay) + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optim, + patience=self.patience, + verbose=False, + min_lr=1e-7, + ) if self.patience else None + return {'optimizer': optim, 'lr_scheduler': scheduler, 'monitor': 'loss/val'} + + +# - + +from torch.utils.data import DataLoader +from pytorch_lightning.loggers import CSVLogger, WandbLogger, TensorBoardLogger, TestTubeLogger +from pytorch_lightning.callbacks.early_stopping import EarlyStopping +from pytorch_lightning.callbacks import LearningRateMonitor + + +# ## Models + +from seq2seq_time.models.baseline import BaselineLast, BaselineMean +from seq2seq_time.models.lstm_seq2seq import LSTMSeq2Seq +from seq2seq_time.models.lstm import LSTM +from seq2seq_time.models.transformer import Transformer +from seq2seq_time.models.transformer_seq2seq import TransformerSeq2Seq +from seq2seq_time.models.neural_process import RANP +from seq2seq_time.models.transformer_process import TransformerProcess +from seq2seq_time.models.tcn import TCNSeq +from seq2seq_time.models.inceptiontime import InceptionTimeSeq +from seq2seq_time.models.xattention import CrossAttention +# + +import gc + +def free_mem(): + gc.collect() + torch.cuda.empty_cache() + gc.collect() +# - + + + +# + +# PARAMS: model +dropout=0.0 +layers=6 +nhead=4 + +models = [ +# lambda xs, ys: BaselineLast(), +# lambda xs, ys, hidden_size: BaselineMean(), + lambda xs, ys, hidden_size, layers: Transformer(xs, + ys, + attention_dropout=dropout, + nhead=nhead, + nlayers=layers, + hidden_size=hidden_size), + + lambda xs, ys, hidden_size, layers:TransformerProcess(xs, + ys, hidden_size=hidden_size, nhead=nhead, + latent_dim=hidden_size//2, dropout=dropout, + nlayers=layers), + lambda xs, ys, hidden_size, layers:TCNSeq(xs, ys, hidden_size=hidden_size, nlayers=layers, dropout=dropout, kernel_size=2), + lambda xs, ys, hidden_size, layers: RANP(xs, + ys, hidden_dim=hidden_size, dropout=dropout, + latent_dim=hidden_size//2, n_decoder_layers=layers, n_latent_encoder_layers=layers, n_det_encoder_layers=layers), + lambda xs, ys, hidden_size, layers: TransformerSeq2Seq(xs, + ys, + hidden_size=hidden_size, + nhead=nhead, + nlayers=layers, + attention_dropout=dropout + ), + lambda xs, ys, hidden_size, layers: LSTM(xs, + ys, + hidden_size=hidden_size, + lstm_layers=layers//2, + lstm_dropout=dropout), + lambda xs, ys, hidden_size, layers: LSTMSeq2Seq(xs, + ys, + hidden_size=hidden_size, + lstm_layers=layers//2, + lstm_dropout=dropout), + lambda xs, ys, hidden_size, layers: CrossAttention(xs, + ys, + nlayers=layers, + hidden_size=hidden_size,), + lambda xs, ys, hidden_size, layers: InceptionTimeSeq(xs, + ys, + kernel_size=96, + layers=layers//2, + hidden_size=hidden_size, + bottleneck=hidden_size//4) + +] +# + +# DEBUG: sanity check + +for Dataset in datasets: + dataset_name = Dataset.__name__ + dataset = Dataset(datasets_root) + ds_train, ds_val, ds_test = dataset.to_datasets(window_past=window_past, + window_future=window_future) + + # Init data + x_past, y_past, x_future, y_future = ds_train.get_rows(10) + xs = x_past.shape[-1] + ys = y_future.shape[-1] + + # Loaders + dl_train = DataLoader(ds_train, + batch_size=batch_size, + shuffle=True, + pin_memory=num_workers == 0, + num_workers=num_workers) + dl_val = DataLoader(ds_val, + shuffle=True, + batch_size=batch_size, + num_workers=num_workers) + + for m_fn in models: + free_mem() + pt_model = m_fn(xs, ys, 8, 4) + model_name = type(pt_model).__name__ + print(timestamp, dataset_name, model_name) + + # Wrap in lightning + model = PL_MODEL(pt_model, + lr=3e-4 + ).to(device) + trainer = pl.Trainer( + fast_dev_run=True, + # GPU + gpus=1, + amp_level='O1', + precision=16, + ) +# - + +# Lets summarize all models, and make sure they have a similar number of parameters + +# ## Train + +from collections import defaultdict +from seq2seq_time.metrics import rmse, smape + + +max_iters=20000 + + +tensorboard_dir = Path(f"../outputs/{timestamp}").resolve() +print(f'For tensorboard run:\ntensorboard --logdir="{tensorboard_dir}"') + + + +# + +class MetricsCallback(pl.Callback): + """PyTorch Lightning metric callback.""" + + def __init__(self): + super().__init__() + self.metrics = [] + + def on_validation_end(self, trainer, pl_module): + self.metrics.append(trainer.callback_metrics) + +def objective(trial): + # sample + hidden_size_exp = trial.suggest_int("hidden_size_exp", 2, 8) + hidden_size = 2**hidden_size_exp + + layers = trial.suggest_int("layers", 2, 12) + + # Load model + pt_model = m_fn(xs, ys, hidden_size, layers) + model_name = type(pt_model).__name__ + + # Wrap in lightning + patience = 2 + model = PL_MODEL(pt_model, + lr=3e-4, patience=patience, + weight_decay=4e-5 + ).to(device) + + + # The default logger in PyTorch Lightning writes to event files to be consumed by + # TensorBoard. We don't use any logger here as it requires us to implement several abstract + # methods. Instead we setup a simple callback, that saves metrics from each validation step. +# metrics_callback = MetricsCallback() + + save_dir = f"../outputs/{timestamp}/{dataset_name}_{model_name}/{trial.number}" + Path(save_dir).mkdir(exist_ok=True, parents=True) + trainer = pl.Trainer( + # Training length + min_epochs=2, + max_epochs=100, + limit_train_batches=max_iters//batch_size, + limit_val_batches=max_iters//batch_size//5, + # Misc + gradient_clip_val=20, + terminate_on_nan=True, + # GPU + gpus=1, + amp_level='O1', + precision=16, + # Callbacks + default_root_dir=save_dir, + logger=False, + callbacks=[ +# metrics_callback, + EarlyStopping(monitor='loss/val', patience=patience * 2), + PyTorchLightningPruningCallback(trial, monitor="loss/val")], + ) + trainer.fit(model, dl_train, dl_val) + + # Run on all val data, using test mode + r = trainer.test(model, test_dataloader=dl_val, verbose=False) + return r[0]['loss/test'] +# - + + + +import optuna +from optuna.integration import PyTorchLightningPruningCallback + +Path(f"../outputs/{timestamp}").mkdir(exist_ok=True) +results = defaultdict(dict) +for Dataset in tqdm(datasets, desc='datasets'): + dataset_name = Dataset.__name__ + dataset = Dataset(datasets_root) + ds_train, ds_val, ds_test = dataset.to_datasets(window_past=window_past, + window_future=window_future) + + # Init data + x_past, y_past, x_future, y_future = ds_train.get_rows(10) + xs = x_past.shape[-1] + ys = y_future.shape[-1] + + # Loaders + dl_train = DataLoader(ds_train, + batch_size=batch_size, + shuffle=True, + pin_memory=num_workers == 0, + num_workers=num_workers) + dl_val = DataLoader(ds_val, + shuffle=False, + batch_size=batch_size, + num_workers=num_workers) + + for i, m_fn in enumerate(tqdm(models, desc=f'models ({dataset_name})')): + try: + model_name = type(m_fn(8, 8, 8, 2)).__name__ + free_mem() + study_name = f'{timestamp}_{dataset_name}-{model_name}' + + storage = f"sqlite:///../outputs/{timestamp}/optuna.db" + pruner = optuna.pruners.MedianPruner() + study = optuna.create_study(storage=storage, + study_name=study_name, + pruner=pruner, + load_if_exists=True) + study.optimize(objective, n_trials=100, timeout=60*60) + print("Number of finished trials: {}".format(len(study.trials))) + + print("Best trial:") + trial = study.best_trial + + print(" Value: {}".format(trial.value)) + + print(" Params: ") + for key, value in trial.params.items(): + print(" {}: {}".format(key, value)) + + except Exception as e: + logging.exception('failed to run model') + +