Files
pytorch-transformer-ts/template/ts-template.ipynb
T

605 lines
24 KiB
Plaintext

{
"cells": [
{
"cell_type": "markdown",
"id": "329d23e6",
"metadata": {},
"source": [
"# PyTorch Transformer for Time Series Implementation\n",
"\n",
"The estimator consits of the:\n",
"\n",
"* Estimator\n",
"* Lightning module\n",
"* Model\n",
"\n",
"Next we will describe all these components and how they are connected with each other.\n",
"\n",
"\n",
"## 1. Estimator Class\n",
"\n",
"The Estimator is the main class which when instantiated allow you to train a particular model given the data and a trainer. In our case we will subclass from a `PyTorchLightningEstimator`. \n",
"\n",
"1. The Estimator's initializer needs to take as input:\n",
" * all the arguments of the model\n",
" * arguments for the pytorch trainer as a dictionary\n",
" * arguments to the dataloader e.g. batch size and number of batches per epoch etc.\n",
" * feature sizes and properties of the dataset (e.g. context length, prediction length, output distribution etc.)\n",
"1. The Estimator needs to define the chain of transformations to the dataset via the `create_transformation` function:\n",
" * for example we can add time feature to the dynamic real features of the dataset etc.\n",
" * we can concatenate the age feature to the dynamic real feature array etc.\n",
"1. The Estimator needs a way to sample the time series during training/validation and testing via the `_create_instance_splitter` method\n",
"1. The Estimator needs methods to create train/val dataloader\n",
"1. The Estimator needs a method that returns the predictor model and corresponding prediction transformations\n",
"1. The Estimator needs a method that return the corresponding lightning module\n",
"\n",
"\n",
"## 2. Lightning Module\n",
"\n",
"The lightning module is like any standard lightning module class and defines:\n",
"\n",
"1. `forward` method which will return the loss for training\n",
"1. `training_step` which will call the `forward` and log the loss\n",
"1. `validation_step` for potentially logging any validation metrics\n",
"1. `configure_optimizers` for configuring the optimizer and potential schedulars etc.\n",
"\n",
"\n",
"## 3. The Model\n",
"\n",
"The core of the logic is in the model class whose intializers take in the arguments you sent it via the Estimator and intialize the `nn.Module` layers like for embedding the categorical features as well as the main network and the probablistic head which will map outputs from the network to parameters of the choosen distribution if needed. \n",
"\n",
"1. The `forward` of this model will be used for prediction,\n",
"1. Therefore the model would need a helper which returns the parameters of the distrubtion or the actual distrubution which can then be used by the lightning module to calcuate the loss\n",
"1. The model would need some helper methods to calculate feature sizes in order to create the appropriate `nn.Modules`\n",
"\n",
"\n",
"## Homework\n",
"\n",
"Go through the code below and add as comments the approiate section number e.g. `2.3` to `validation_step` to identify all the parts"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "8dc0844c",
"metadata": {},
"outputs": [],
"source": [
"from typing import List, Optional, Iterable, Dict, Any\n",
"\n",
"import numpy as np\n",
"\n",
"import torch\n",
"import torch.nn as nn\n",
"from torch.utils.data import DataLoader\n",
"\n",
"import pytorch_lightning as pl\n",
"\n",
"from gluonts.core.component import validated\n",
"from gluonts.dataset.common import Dataset\n",
"from gluonts.dataset.field_names import FieldName\n",
"from gluonts.itertools import Cyclic, PseudoShuffled, IterableSlice\n",
"from gluonts.time_feature import (\n",
" TimeFeature,\n",
" time_features_from_frequency_str,\n",
")\n",
"from gluonts.torch.modules.loss import DistributionLoss, NegativeLogLikelihood\n",
"from gluonts.transform import (\n",
" Transformation,\n",
" Chain,\n",
" RemoveFields,\n",
" SetField,\n",
" AsNumpyArray,\n",
" AddObservedValuesIndicator,\n",
" AddTimeFeatures,\n",
" AddAgeFeature,\n",
" VstackFeatures,\n",
" InstanceSplitter,\n",
" ValidationSplitSampler,\n",
" TestSplitSampler,\n",
" ExpectedNumInstanceSampler,\n",
" SelectFields,\n",
" InstanceSampler,\n",
")\n",
"from gluonts.torch.util import (\n",
" IterableDataset,\n",
")\n",
"from gluonts.torch.model.estimator import PyTorchLightningEstimator\n",
"from gluonts.torch.model.predictor import PyTorchPredictor\n",
"from gluonts.torch.distributions import (\n",
" DistributionOutput,\n",
" StudentTOutput,\n",
")\n",
"from gluonts.torch.modules.scaler import MeanScaler, NOPScaler\n",
"from gluonts.torch.modules.feature import FeatureEmbedder\n",
"from gluonts.time_feature import get_lags_for_frequency"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b8af6d90",
"metadata": {},
"outputs": [],
"source": [
"class TransformerModel(nn.Module):\n",
" @validated()\n",
" def __init__(\n",
" self,\n",
" freq: str,\n",
" context_length: int,\n",
" prediction_length: int,\n",
" num_feat_dynamic_real: int,\n",
" num_feat_static_real: int,\n",
" num_feat_static_cat: int,\n",
" cardinality: List[int],\n",
" embedding_dimension: Optional[List[int]] = None,\n",
" # TODO transformer arguments\n",
" \n",
" dropout_rate: float = 0.1,\n",
" distr_output: DistributionOutput = StudentTOutput(),\n",
" lags_seq: Optional[List[int]] = None,\n",
" scaling: bool = True,\n",
" num_parallel_samples: int = 100,\n",
" ) -> None:\n",
" super().__init__()\n",
" \n",
" self.context_length = context_length\n",
" self.prediction_length = prediction_length\n",
" self.distr_output = distr_output\n",
" self.param_proj = distr_output.get_args_proj(hidden_size)\n",
" self.target_shape = distr_output.event_shape\n",
" self.num_feat_dynamic_real = num_feat_dynamic_real\n",
" self.num_feat_static_cat = num_feat_static_cat\n",
" self.num_feat_static_real = num_feat_static_real\n",
" self.embedding_dimension = (\n",
" embedding_dimension\n",
" if embedding_dimension is not None or cardinality is None\n",
" else [min(50, (cat + 1) // 2) for cat in cardinality]\n",
" )\n",
" self.lags_seq = lags_seq or get_lags_for_frequency(freq_str=freq)\n",
" self.num_parallel_samples = num_parallel_samples\n",
" self.history_length = self.context_length + max(self.lags_seq)\n",
" self.embedder = FeatureEmbedder(\n",
" cardinalities=cardinality,\n",
" embedding_dims=self.embedding_dimension,\n",
" )\n",
" if scaling:\n",
" self.scaler = MeanScaler(dim=1, keepdim=True)\n",
" else:\n",
" self.scaler = NOPScaler(dim=1, keepdim=True)\n",
" \n",
" # TODO transformer enc-decoder and mask initializer\n",
" \n",
" \n",
" # TODO\n",
" # add method that does the forward for training\n",
" \n",
" @property\n",
" def _number_of_features(self) -> int:\n",
" return (\n",
" sum(self.embedding_dimension)\n",
" + self.num_feat_dynamic_real\n",
" + self.num_feat_static_real\n",
" + 1 # the log(scale)\n",
" )\n",
"\n",
" @property\n",
" def _past_length(self) -> int:\n",
" return self.context_length + max(self.lags_seq)\n",
" \n",
" # for prediction\n",
" def forward(\n",
" self,\n",
" feat_static_cat: torch.Tensor,\n",
" feat_static_real: torch.Tensor,\n",
" past_time_feat: torch.Tensor,\n",
" past_target: torch.Tensor,\n",
" past_observed_values: torch.Tensor,\n",
" future_time_feat: torch.Tensor,\n",
" num_parallel_samples: Optional[int] = None,\n",
" ) -> torch.Tensor:\n",
" if num_parallel_samples is None:\n",
" num_parallel_samples = self.num_parallel_samples\n",
" \n",
" # TODO"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9b1529a8",
"metadata": {},
"outputs": [],
"source": [
"class TransformerLightningModule(pl.LightningModule):\n",
" def __init__(\n",
" self,\n",
" model: TransformerModel,\n",
" loss: DistributionLoss = NegativeLogLikelihood(),\n",
" lr: float = 1e-3,\n",
" weight_decay: float = 1e-8,\n",
" ) -> None:\n",
" super().__init__()\n",
" self.save_hyperparameters()\n",
" self.model = model\n",
" self.loss = loss\n",
" self.lr = lr\n",
" self.weight_decay = weight_decay\n",
" \n",
" def training_step(self, batch, batch_idx: int):\n",
" \"\"\"Execute training step\"\"\"\n",
" train_loss = self(batch)\n",
" self.log(\n",
" \"train_loss\",\n",
" train_loss,\n",
" on_epoch=True,\n",
" on_step=False,\n",
" prog_bar=True,\n",
" )\n",
" return train_loss\n",
"\n",
" def validation_step(self, batch, batch_idx: int):\n",
" \"\"\"Execute validation step\"\"\"\n",
" val_loss = self(batch)\n",
" self.log(\n",
" \"val_loss\", val_loss, on_epoch=True, on_step=False, prog_bar=True\n",
" )\n",
" return val_loss\n",
"\n",
" def configure_optimizers(self):\n",
" \"\"\"Returns the optimizer to use\"\"\"\n",
" return torch.optim.Adam(\n",
" self.model.parameters(),\n",
" lr=self.lr,\n",
" weight_decay=self.weight_decay,\n",
" )\n",
"\n",
" def forward(self, batch):\n",
" feat_static_cat = batch[\"feat_static_cat\"]\n",
" feat_static_real = batch[\"feat_static_real\"]\n",
" past_time_feat = batch[\"past_time_feat\"]\n",
" past_target = batch[\"past_target\"]\n",
" future_time_feat = batch[\"future_time_feat\"]\n",
" future_target = batch[\"future_target\"]\n",
" past_observed_values = batch[\"past_observed_values\"]\n",
" future_observed_values = batch[\"future_observed_values\"]\n",
" \n",
" # TODO\n",
" # params, scale, _, _ = self.model.compute_params(...) #TODO\n",
" distr = self.model.output_distribution(params, scale)\n",
"\n",
" context_target = past_target[:, -self.model.context_length + 1 :]\n",
" target = torch.cat(\n",
" (context_target, future_target),\n",
" dim=1,\n",
" )\n",
" loss_values = self.loss(distr, target)\n",
"\n",
" context_observed = past_observed_values[\n",
" :, -self.model.context_length + 1 :\n",
" ]\n",
" observed_values = torch.cat(\n",
" (context_observed, future_observed_values), dim=1\n",
" )\n",
"\n",
" if len(self.model.target_shape) == 0:\n",
" loss_weights = observed_values\n",
" else:\n",
" loss_weights = observed_values.min(dim=-1, keepdim=False)\n",
"\n",
" return weighted_average(loss_values, weights=loss_weights)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f7b4c72f",
"metadata": {},
"outputs": [],
"source": [
"PREDICTION_INPUT_NAMES = [\n",
" \"feat_static_cat\",\n",
" \"feat_static_real\",\n",
" \"past_time_feat\",\n",
" \"past_target\",\n",
" \"past_observed_values\",\n",
" \"future_time_feat\",\n",
"]\n",
"\n",
"TRAINING_INPUT_NAMES = PREDICTION_INPUT_NAMES + [\n",
" \"future_target\",\n",
" \"future_observed_values\",\n",
"]"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1937891a",
"metadata": {},
"outputs": [],
"source": [
"class TransformerEstimator(PyTorchLightningEstimator):\n",
" @validated()\n",
" def __init__(\n",
" self,\n",
" freq: str,\n",
" prediction_length: int,\n",
" context_length: Optional[int] = None,\n",
" \n",
" # Transformer arguments\n",
" # TODO\n",
" \n",
" dropout_rate: float = 0.1,\n",
" num_feat_dynamic_real: int = 0,\n",
" num_feat_static_cat: int = 0,\n",
" num_feat_static_real: int = 0,\n",
" cardinality: Optional[List[int]] = None,\n",
" embedding_dimension: Optional[List[int]] = None,\n",
" distr_output: DistributionOutput = StudentTOutput(),\n",
" loss: DistributionLoss = NegativeLogLikelihood(),\n",
" scaling: bool = True,\n",
" lags_seq: Optional[List[int]] = None,\n",
" time_features: Optional[List[TimeFeature]] = None,\n",
" num_parallel_samples: int = 100,\n",
" batch_size: int = 32,\n",
" num_batches_per_epoch: int = 50,\n",
" trainer_kwargs: Optional[Dict[str, Any]] = dict(),\n",
" train_sampler: Optional[InstanceSampler] = None,\n",
" validation_sampler: Optional[InstanceSampler] = None,\n",
" ) -> None:\n",
" trainer_kwargs = {\n",
" \"max_epochs\": 100,\n",
" \"gradient_clip_val\": 10.0,\n",
" **trainer_kwargs,\n",
" }\n",
" super().__init__(trainer_kwargs=trainer_kwargs)\n",
" \n",
" self.freq = freq\n",
" self.context_length = (\n",
" context_length if context_length is not None else prediction_length\n",
" )\n",
" self.prediction_length = prediction_length\n",
" self.distr_output = distr_output\n",
" self.loss = loss\n",
" \n",
" self.dropout_rate = dropout_rate\n",
" self.num_feat_dynamic_real = num_feat_dynamic_real\n",
" self.num_feat_static_cat = num_feat_static_cat\n",
" self.num_feat_static_real = num_feat_static_real\n",
" self.cardinality = (\n",
" cardinality if cardinality and num_feat_static_cat > 0 else [1]\n",
" )\n",
" self.embedding_dimension = embedding_dimension\n",
" self.scaling = scaling\n",
" self.lags_seq = lags_seq\n",
" self.time_features = (\n",
" time_features\n",
" if time_features is not None\n",
" else time_features_from_frequency_str(self.freq)\n",
" )\n",
"\n",
" self.num_parallel_samples = num_parallel_samples\n",
" self.batch_size = batch_size\n",
" self.num_batches_per_epoch = num_batches_per_epoch\n",
"\n",
" self.train_sampler = train_sampler or ExpectedNumInstanceSampler(\n",
" num_instances=1.0, min_future=prediction_length\n",
" )\n",
" self.validation_sampler = validation_sampler or ValidationSplitSampler(\n",
" min_future=prediction_length\n",
" )\n",
" \n",
" def create_transformation(self) -> Transformation:\n",
" remove_field_names = []\n",
" if self.num_feat_static_real == 0:\n",
" remove_field_names.append(FieldName.FEAT_STATIC_REAL)\n",
" if self.num_feat_dynamic_real == 0:\n",
" remove_field_names.append(FieldName.FEAT_DYNAMIC_REAL)\n",
"\n",
" return Chain(\n",
" [RemoveFields(field_names=remove_field_names)]\n",
" + (\n",
" [SetField(output_field=FieldName.FEAT_STATIC_CAT, value=[0])]\n",
" if not self.num_feat_static_cat > 0\n",
" else []\n",
" )\n",
" + (\n",
" [\n",
" SetField(\n",
" output_field=FieldName.FEAT_STATIC_REAL, value=[0.0]\n",
" )\n",
" ]\n",
" if not self.num_feat_static_real > 0\n",
" else []\n",
" )\n",
" + [\n",
" AsNumpyArray(\n",
" field=FieldName.FEAT_STATIC_CAT,\n",
" expected_ndim=1,\n",
" dtype=np.long,\n",
" ),\n",
" AsNumpyArray(\n",
" field=FieldName.FEAT_STATIC_REAL,\n",
" expected_ndim=1,\n",
" ),\n",
" AsNumpyArray(\n",
" field=FieldName.TARGET,\n",
" # in the following line, we add 1 for the time dimension\n",
" expected_ndim=1 + len(self.distr_output.event_shape),\n",
" ),\n",
" AddObservedValuesIndicator(\n",
" target_field=FieldName.TARGET,\n",
" output_field=FieldName.OBSERVED_VALUES,\n",
" ),\n",
" AddTimeFeatures(\n",
" start_field=FieldName.START,\n",
" target_field=FieldName.TARGET,\n",
" output_field=FieldName.FEAT_TIME,\n",
" time_features=self.time_features,\n",
" pred_length=self.prediction_length,\n",
" ),\n",
" AddAgeFeature(\n",
" target_field=FieldName.TARGET,\n",
" output_field=FieldName.FEAT_AGE,\n",
" pred_length=self.prediction_length,\n",
" log_scale=True,\n",
" ),\n",
" VstackFeatures(\n",
" output_field=FieldName.FEAT_TIME,\n",
" input_fields=[FieldName.FEAT_TIME, FieldName.FEAT_AGE]\n",
" + (\n",
" [FieldName.FEAT_DYNAMIC_REAL]\n",
" if self.num_feat_dynamic_real > 0\n",
" else []\n",
" ),\n",
" ),\n",
" ]\n",
" )\n",
"\n",
" def _create_instance_splitter(\n",
" self, module: TransformerLightningModule, mode: str\n",
" ):\n",
" assert mode in [\"training\", \"validation\", \"test\"]\n",
"\n",
" instance_sampler = {\n",
" \"training\": self.train_sampler,\n",
" \"validation\": self.validation_sampler,\n",
" \"test\": TestSplitSampler(),\n",
" }[mode]\n",
"\n",
" return InstanceSplitter(\n",
" target_field=FieldName.TARGET,\n",
" is_pad_field=FieldName.IS_PAD,\n",
" start_field=FieldName.START,\n",
" forecast_start_field=FieldName.FORECAST_START,\n",
" instance_sampler=instance_sampler,\n",
" past_length=module.model._past_length,\n",
" future_length=self.prediction_length,\n",
" time_series_fields=[\n",
" FieldName.FEAT_TIME,\n",
" FieldName.OBSERVED_VALUES,\n",
" ],\n",
" dummy_value=self.distr_output.value_in_support,\n",
" )\n",
"\n",
" def create_training_data_loader(\n",
" self,\n",
" data: Dataset,\n",
" module: TransformerLightningModule,\n",
" shuffle_buffer_length: Optional[int] = None,\n",
" **kwargs,\n",
" ) -> Iterable:\n",
" transformation = self._create_instance_splitter(\n",
" module, \"training\"\n",
" ) + SelectFields(TRAINING_INPUT_NAMES)\n",
"\n",
" training_instances = transformation.apply(\n",
" Cyclic(data)\n",
" if shuffle_buffer_length is None\n",
" else PseudoShuffled(\n",
" Cyclic(data), shuffle_buffer_length=shuffle_buffer_length\n",
" )\n",
" )\n",
"\n",
" return IterableSlice(\n",
" iter(\n",
" DataLoader(\n",
" IterableDataset(training_instances),\n",
" batch_size=self.batch_size,\n",
" **kwargs,\n",
" )\n",
" ),\n",
" self.num_batches_per_epoch,\n",
" )\n",
"\n",
" def create_validation_data_loader(\n",
" self,\n",
" data: Dataset,\n",
" module: TransformerLightningModule,\n",
" **kwargs,\n",
" ) -> Iterable:\n",
" transformation = self._create_instance_splitter(\n",
" module, \"validation\"\n",
" ) + SelectFields(TRAINING_INPUT_NAMES)\n",
"\n",
" validation_instances = transformation.apply(data)\n",
"\n",
" return DataLoader(\n",
" IterableDataset(validation_instances),\n",
" batch_size=self.batch_size,\n",
" **kwargs,\n",
" )\n",
" \n",
" def create_predictor(\n",
" self,\n",
" transformation: Transformation,\n",
" module: TransformerLightningModule,\n",
" ) -> PyTorchPredictor:\n",
" prediction_splitter = self._create_instance_splitter(module, \"test\")\n",
"\n",
" return PyTorchPredictor(\n",
" input_transform=transformation + prediction_splitter,\n",
" input_names=PREDICTION_INPUT_NAMES,\n",
" prediction_net=module.model,\n",
" batch_size=self.batch_size,\n",
" prediction_length=self.prediction_length,\n",
" )\n",
"\n",
" def create_lightning_module(self) -> TransformerLightningModule:\n",
" model = TransformerModel(\n",
" freq=self.freq,\n",
" context_length=self.context_length,\n",
" prediction_length=self.prediction_length,\n",
" num_feat_dynamic_real=1 + self.num_feat_dynamic_real + len(self.time_features),\n",
" num_feat_static_real=max(1, self.num_feat_static_real),\n",
" num_feat_static_cat=max(1, self.num_feat_static_cat),\n",
" cardinality=self.cardinality,\n",
" embedding_dimension=self.embedding_dimension,\n",
"\n",
" # TODO Transformer arguments\n",
" \n",
" # univariate input\n",
" input_size=self.input_size,\n",
" distr_output=self.distr_output,\n",
" lags_seq=self.lags_seq,\n",
" scaling=self.scaling,\n",
" num_parallel_samples=self.num_parallel_samples,\n",
" )\n",
" \n",
" return TransformerLightningModule(model=model, loss=self.loss)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c75639e7",
"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.8.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
}