initial lstnet multivariate point forecasting model (#9)

* initial lstnet

* lstnet network

* fixed forward

* fix splitter

* fix prediction

* rename argument to what it is i.e. time_first

* fixed scaling and some default values

* scaler can now take time_first=False tensors

pts/model/lstnet/__init__.py

@@ -0,0 +1 @@
+from .lstnet_estimator import LSTNetEstimator

pts/model/lstnet/lstnet_estimator.py

@@ -0,0 +1,141 @@
+from typing import List, Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from pts import Trainer
+from pts.dataset import FieldName
+from pts.model import PTSEstimator, Predictor, PTSPredictor, copy_parameters
+from pts.transform import (
+    InstanceSplitter,
+    Transformation,
+    Chain,
+    RemoveFields,
+    ExpectedNumInstanceSampler,
+    AddObservedValuesIndicator,
+    AsNumpyArray,
+)
+from .lstnet_network import LSTNetTrain, LSTNetPredict
+
+
+class LSTNetEstimator(PTSEstimator):
+    def __init__(
+        self,
+        freq: str,
+        context_length: int,
+        num_series: int,
+        ar_window: int = 24,
+        skip_size: int = 24,
+        channels: int = 100,
+        kernel_size: int = 6,
+        prediction_length: Optional[int] = None,
+        horizon: Optional[int] = None,
+        trainer: Trainer = Trainer(),
+        dropout_rate: Optional[float] = 0.2,
+        output_activation: Optional[str] = None,
+        rnn_cell_type: str = "GRU",
+        rnn_num_cells: int = 100,
+        skip_rnn_cell_type: str = "GRU",
+        skip_rnn_num_cells: int = 5,
+        scaling: bool = True,
+        dtype: np.dtype = np.float32,
+    ):
+        super().__init__(trainer, dtype=dtype)
+
+        self.freq = freq
+        self.num_series = num_series
+        self.skip_size = skip_size
+        self.ar_window = ar_window
+        self.horizon = horizon
+        self.prediction_length = prediction_length
+
+        self.future_length = horizon if horizon is not None else prediction_length
+        self.context_length = context_length
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.dropout_rate = dropout_rate
+        self.output_activation = output_activation
+        self.rnn_cell_type = rnn_cell_type
+        self.rnn_num_cells = rnn_num_cells
+        self.skip_rnn_cell_type = skip_rnn_cell_type
+        self.skip_rnn_num_cells = skip_rnn_num_cells
+        self.scaling = scaling
+        self.dtype = dtype
+
+    def create_transformation(self) -> Transformation:
+        return Chain(
+            trans=[
+                AsNumpyArray(field=FieldName.TARGET, expected_ndim=2, dtype=self.dtype),
+                AddObservedValuesIndicator(
+                    target_field=FieldName.TARGET,
+                    output_field=FieldName.OBSERVED_VALUES,
+                    dtype=self.dtype,
+                ),
+                InstanceSplitter(
+                    target_field=FieldName.TARGET,
+                    is_pad_field=FieldName.IS_PAD,
+                    start_field=FieldName.START,
+                    forecast_start_field=FieldName.FORECAST_START,
+                    train_sampler=ExpectedNumInstanceSampler(num_instances=1),
+                    time_series_fields=[FieldName.OBSERVED_VALUES],
+                    past_length=self.context_length,
+                    future_length=self.future_length,
+                    time_first=False,
+                ),
+            ]
+        )
+
+    def create_training_network(self, device: torch.device) -> LSTNetTrain:
+        return LSTNetTrain(
+            num_series=self.num_series,
+            channels=self.channels,
+            kernel_size=self.kernel_size,
+            rnn_cell_type=self.rnn_cell_type,
+            rnn_num_cells=self.rnn_num_cells,
+            skip_rnn_cell_type=self.skip_rnn_cell_type,
+            skip_rnn_num_cells=self.skip_rnn_num_cells,
+            skip_size=self.skip_size,
+            ar_window=self.ar_window,
+            context_length=self.context_length,
+            horizon=self.horizon,
+            prediction_length=self.prediction_length,
+            dropout_rate=self.dropout_rate,
+            output_activation=self.output_activation,
+            scaling=self.scaling,
+        ).to(device)
+
+    def create_predictor(
+        self,
+        transformation: Transformation,
+        trained_network: LSTNetTrain,
+        device: torch.device,
+    ) -> PTSPredictor:
+        prediction_network = LSTNetPredict(
+            num_series=self.num_series,
+            channels=self.channels,
+            kernel_size=self.kernel_size,
+            rnn_cell_type=self.rnn_cell_type,
+            rnn_num_cells=self.rnn_num_cells,
+            skip_rnn_cell_type=self.skip_rnn_cell_type,
+            skip_rnn_num_cells=self.skip_rnn_num_cells,
+            skip_size=self.skip_size,
+            ar_window=self.ar_window,
+            context_length=self.context_length,
+            horizon=self.horizon,
+            prediction_length=self.prediction_length,
+            dropout_rate=self.dropout_rate,
+            output_activation=self.output_activation,
+            scaling=self.scaling,
+        ).to(device)
+
+        copy_parameters(trained_network, prediction_network)
+
+        return PTSPredictor(
+            input_transform=transformation,
+            prediction_net=prediction_network,
+            batch_size=self.trainer.batch_size,
+            freq=self.freq,
+            prediction_length=self.prediction_length,
+            device=device,
+        )

pts/model/lstnet/lstnet_network.py

@@ -0,0 +1,189 @@
+from typing import List, Tuple, Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from pts.modules import MeanScaler, NOPScaler
+
+
+class LSTNetBase(nn.Module):
+    def __init__(
+        self,
+        num_series: int,
+        channels: int,
+        kernel_size: int,
+        rnn_cell_type: str,
+        rnn_num_cells: int,
+        skip_rnn_cell_type: str,
+        skip_rnn_num_cells: int,
+        skip_size: int,
+        ar_window: int,
+        context_length: int,
+        horizon: Optional[int],
+        prediction_length: Optional[int],
+        dropout_rate: float,
+        output_activation: Optional[str],
+        scaling: bool,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+
+        self.num_series = num_series
+        self.channels = channels
+        assert (
+            channels % skip_size == 0
+        ), "number of conv1d `channels` must be divisible by the `skip_size`"
+        self.skip_size = skip_size
+        assert ar_window > 0, "auto-regressive window must be a positive integer"
+        self.ar_window = ar_window
+        assert not ((horizon is None)) == (
+            prediction_length is None
+        ), "Exactly one of `horizon` and `prediction_length` must be set at a time"
+        assert horizon is None or horizon > 0, "`horizon` must be greater than zero"
+        assert (
+            prediction_length is None or prediction_length > 0
+        ), "`prediction_length` must be greater than zero"
+        self.prediction_length = prediction_length
+        self.horizon = horizon
+        assert context_length > 0, "`context_length` must be greater than zero"
+        self.context_length = context_length
+        if output_activation is not None:
+            assert output_activation in [
+                "sigmoid",
+                "tanh",
+            ], "`output_activation` must be either 'sigmiod' or 'tanh' "
+        self.output_activation = output_activation
+        assert rnn_cell_type in [
+            "GRU",
+            "LSTM",
+        ], "`rnn_cell_type` must be either 'GRU' or 'LSTM' "
+        assert skip_rnn_cell_type in [
+            "GRU",
+            "LSTM",
+        ], "`skip_rnn_cell_type` must be either 'GRU' or 'LSTM' "
+
+        self.conv_out = context_length - kernel_size
+        self.conv_skip = self.conv_out // skip_size
+        assert self.conv_skip > 0, (
+            "conv1d output size must be greater than or equal to `skip_size`\n"
+            "Choose a smaller `kernel_size` or bigger `context_length`"
+        )
+        self.skip_rnn_c_dim = channels * skip_size
+
+        self.cnn = nn.Conv2d(
+            in_channels=1, out_channels=channels, kernel_size=(num_series, kernel_size)
+        )
+
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+        rnn = {"LSTM": nn.LSTM, "GRU": nn.GRU}[rnn_cell_type]
+        self.rnn = rnn(
+            input_size=channels,
+            hidden_size=rnn_num_cells,
+            # dropout=dropout_rate,
+        )
+
+        skip_rnn = {"LSTM": nn.LSTM, "GRU": nn.GRU}[skip_rnn_cell_type]
+        self.skip_rnn_num_cells = skip_rnn_num_cells
+        self.skip_rnn = skip_rnn(
+            input_size=channels,
+            hidden_size=skip_rnn_num_cells,
+            # dropout=dropout_rate,
+        )
+
+        self.fc = nn.Linear(rnn_num_cells + skip_size * skip_rnn_num_cells, num_series)
+
+        if self.horizon:
+            self.ar_fc = nn.Linear(ar_window, 1)
+        else:
+            self.ar_fc = nn.Linear(ar_window, prediction_length)
+
+        if scaling:
+            self.scaler = MeanScaler(keepdim=True, time_first=False)
+        else:
+            self.scaler = NOPScaler(keepdim=True, time_first=False)
+
+    def forward(
+        self, past_target: torch.Tensor, past_observed_values: torch.Tensor
+    ) -> torch.Tensor:
+        scaled_past_target, scale = self.scaler(
+            past_target[..., -self.context_length :],  # [B, C, T]
+            past_observed_values[..., -self.context_length :]  # [B, C, T]
+        )
+
+        # CNN
+        c = F.relu(self.cnn(scaled_past_target.unsqueeze(1)))
+        c = self.dropout(c)
+        c = c.squeeze()  # [B, C, T]
+
+        # RNN
+        r = c.permute(2, 0, 1)  # [F (T), B, C]
+        _, r = self.rnn(r)  # [1, B, H]
+        r = self.dropout(r.squeeze()) # [B, H]
+
+        # Skip-RNN
+        skip_c = c[..., -self.conv_skip * self.skip_size :]
+        skip_c = skip_c.reshape(-1, self.channels, self.conv_skip, self.skip_size)
+        skip_c = skip_c.permute(2, 0, 3, 1)
+        skip_c = skip_c.reshape((self.conv_skip, -1, self.channels))
+        _, skip_c = self.skip_rnn(skip_c)
+        skip_c = skip_c.reshape((-1, self.skip_size * self.skip_rnn_num_cells))
+        skip_c = self.dropout(skip_c)
+
+        res = self.fc(torch.cat((r, skip_c), 1)).unsqueeze(-1)
+
+        # Highway
+        ar_x = scaled_past_target[..., -self.ar_window :]
+        ar_x = ar_x.reshape(-1, self.ar_window)
+
+        ar_x = self.ar_fc(ar_x)
+        if self.horizon:
+            ar_x = ar_x.reshape(-1, self.num_series, 1)
+        else:
+            ar_x = ar_x.reshape(-1, self.num_series, self.prediction_length)
+        out = res + ar_x
+
+        if self.output_activation is None:
+            return out, scale
+
+        return (
+            (
+                torch.sigmoid(out)
+                if self.output_activation == "sigmoid"
+                else torch.tanh(out)
+            ),
+            scale,
+        )
+
+
+class LSTNetTrain(LSTNetBase):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.loss_fn = nn.L1Loss()
+
+    def forward(
+        self,
+        past_target: torch.Tensor,
+        past_observed_values: torch.Tensor,
+        future_target: torch.Tensor,
+    ) -> torch.Tensor:
+        ret, scale = super().forward(past_target, past_observed_values)
+
+        if self.horizon:
+            future_target = future_target[..., -1:]
+
+        loss = self.loss_fn(ret*scale, future_target)
+        return loss
+
+
+class LSTNetPredict(LSTNetBase):
+    def forward(
+        self, past_target: torch.Tensor, past_observed_values: torch.Tensor
+    ) -> torch.Tensor:
+        ret, scale = super().forward(past_target, past_observed_values)
+        ret = (ret*scale).permute(0, 2, 1)
+
+        return ret.unsqueeze(1)

pts/modules/scaler.py

@@ -6,9 +6,10 @@ import torch.nn as nn
 
 
 class Scaler(ABC, nn.Module):
-    def __init__(self, keepdim: bool = False):
+    def __init__(self, keepdim: bool = False, time_first: bool = True):
         super().__init__()
         self.keepdim = keepdim
+        self.time_first = time_first
 
     @abstractmethod
     def compute_scale(
@@ -23,7 +24,8 @@ class Scaler(ABC, nn.Module):
         Parameters
         ----------
         data
-            tensor of shape (N, T, C) containing the data to be scaled
+            tensor of shape (N, T, C) if ``time_first == True`` or (N, C, T)
+            if ``time_first == False`` containing the data to be scaled
 
         observed_indicator
             observed_indicator: binary tensor with the same shape as
@@ -33,19 +35,23 @@ class Scaler(ABC, nn.Module):
         Returns
         -------
         Tensor
-            Tensor containing the "scaled" data, shape: (N, T, C).
+            Tensor containing the "scaled" data, shape: (N, T, C) or (N, C, T).
         Tensor
-            Tensor containing the scale, of shape (N, C) if ``keepdim == False``, and shape
-            (N, 1, C) if ``keepdim == True``.
+            Tensor containing the scale, of shape (N, C) if ``keepdim == False``, 
+            and shape (N, 1, C) or (N, C, 1) if ``keepdim == True``.
         """
 
         scale = self.compute_scale(data, observed_indicator)
 
+        if self.time_first:
+            dim = 1
+        else:
+            dim = 2
         if self.keepdim:
-            scale = scale.unsqueeze(1)
+            scale = scale.unsqueeze(dim=dim)
             return data / scale, scale
         else:
-            return data / scale.unsqueeze(1), scale
+            return data / scale.unsqueeze(dim=dim), scale
 
 
 class MeanScaler(Scaler):
@@ -69,9 +75,15 @@ class MeanScaler(Scaler):
     def compute_scale(
         self, data: torch.Tensor, observed_indicator: torch.Tensor
     ) -> torch.Tensor:
+
+        if self.time_first:
+            dim = 1
+        else:
+            dim = 2
+
         # these will have shape (N, C)
-        num_observed = observed_indicator.sum(dim=1)
-        sum_observed = (data.abs() * observed_indicator).sum(dim=1)
+        num_observed = observed_indicator.sum(dim=dim)
+        sum_observed = (data.abs() * observed_indicator).sum(dim=dim)
 
         # first compute a global scale per-dimension
         total_observed = num_observed.sum(dim=0)
@@ -105,4 +117,8 @@ class NOPScaler(Scaler):
     def compute_scale(
         self, data: torch.Tensor, observed_indicator: torch.Tensor
     ) -> torch.Tensor:
-        return torch.ones_like(data).mean(dim=1)
+        if self.time_first:
+            dim = 1
+        else:
+            dim = 2
+        return torch.ones_like(data).mean(dim=dim)

pts/transform/split.py

@@ -92,7 +92,7 @@ class InstanceSplitter(FlatMapTransformation):
         length of the target seen before making prediction
     future_length
         length of the target that must be predicted
-    batch_first
+    time_first
         whether to have time series output in (time, dimension) or in
         (dimension, time) layout
     time_series_fields
@@ -116,7 +116,7 @@ class InstanceSplitter(FlatMapTransformation):
         train_sampler: InstanceSampler,
         past_length: int,
         future_length: int,
-        batch_first: bool = True,
+        time_first: bool = True,
         time_series_fields: Optional[List[str]] = None,
         pick_incomplete: bool = True,
     ) -> None:
@@ -126,7 +126,7 @@ class InstanceSplitter(FlatMapTransformation):
         self.train_sampler = train_sampler
         self.past_length = past_length
         self.future_length = future_length
-        self.batch_first = batch_first
+        self.time_first = time_first
         self.ts_fields = time_series_fields if time_series_fields is not None else []
         self.target_field = target_field
         self.is_pad_field = is_pad_field
@@ -197,7 +197,7 @@ class InstanceSplitter(FlatMapTransformation):
             if pad_length > 0:
                 pad_indicator[:pad_length] = 1
 
-            if self.batch_first:
+            if self.time_first:
                 for ts_field in slice_cols:
                     d[self._past(ts_field)] = d[self._past(ts_field)].transpose()
                     d[self._future(ts_field)] = d[self._future(ts_field)].transpose()
@@ -245,7 +245,7 @@ class CanonicalInstanceSplitter(FlatMapTransformation):
         instance sampler that provides sampling indices given a time-series
     instance_length
         length of the target seen before making prediction
-    batch_first
+    time_first
         whether to have time series output in (time, dimension) or in
         (dimension, time) layout
     time_series_fields
@@ -270,7 +270,7 @@ class CanonicalInstanceSplitter(FlatMapTransformation):
         forecast_start_field: str,
         instance_sampler: InstanceSampler,
         instance_length: int,
-        batch_first: bool = True,
+        time_first: bool = True,
         time_series_fields: List[str] = [],
         allow_target_padding: bool = False,
         pad_value: float = 0.0,
@@ -279,7 +279,7 @@ class CanonicalInstanceSplitter(FlatMapTransformation):
     ) -> None:
         self.instance_sampler = instance_sampler
         self.instance_length = instance_length
-        self.batch_first = batch_first
+        self.time_first = time_first
         self.dynamic_feature_fields = time_series_fields
         self.target_field = target_field
         self.allow_target_padding = allow_target_padding
@@ -349,14 +349,14 @@ class CanonicalInstanceSplitter(FlatMapTransformation):
                 else:
                     past_ts = full_ts[..., (i - self.instance_length) : i]
 
-                past_ts = past_ts.transpose() if self.batch_first else past_ts
+                past_ts = past_ts.transpose() if self.time_first else past_ts
                 d[self._past(ts_field)] = past_ts
 
                 if self.use_prediction_features and not is_train:
                     if not ts_field == self.target_field:
                         future_ts = full_ts[..., i : i + self.prediction_length]
                         future_ts = (
-                            future_ts.transpose() if self.batch_first else future_ts
+                            future_ts.transpose() if self.time_first else future_ts
                         )
                         d[self._future(ts_field)] = future_ts
 

test/test_transform.py

@@ -363,7 +363,7 @@ def test_multi_dim_transformation(is_train):
                 past_length=train_length,
                 future_length=pred_length,
                 time_series_fields=["dynamic_feat", "observed_values"],
-                batch_first=False,
+                time_first=False,
             ),
         ]
     )