Simplified encoder decoder model and moved curriculum learning outside

model/pytorch/dcrnn_model.py

@@ -1,12 +1,10 @@
-import numpy as np
 import torch
 import torch.nn as nn
-from abc import ABC, abstractmethod
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
-class DCRNNModel(metaclass=ABC):
+class DCRNNModel:
     def __init__(self, is_training, scale_factor, adj_mx, **model_kwargs):
         super().__init__()
         self.adj_mx = adj_mx
@@ -22,56 +20,6 @@ class DCRNNModel(metaclass=ABC):
         self.input_dim = int(model_kwargs.get('input_dim', 1))
         self.hidden_state_size = self.num_nodes * self.rnn_units
 
-    @abstractmethod
-    @property
-    def dcgru_layers(self):
-        pass
-
-    def _forward_cell(self, cell_input, prev_hidden_states):
-        """
-        Runs for 1 time step through all layers.
-        :param cell_input: shape (batch_size, input feature size)
-        :param prev_hidden_states: (num_layers, batch_size, hidden size)
-        :return: output of cell: shape(batch_size, hidden size)
-                all hidden states from layer: shape(num_layers, batch_size, hidden size)
-        """
-        hidden_states = []
-        output = cell_input
-        for layer_num, dcgru_layer in enumerate(self.dcgru_layers):
-            hidden_state = dcgru_layer(output, prev_hidden_states[layer_num])
-            hidden_states.append(hidden_state)
-            output = hidden_state
-
-        return output, torch.cat(hidden_states, dim=1)  # runs in O(num_layers) so not too slow #todo: check dim
-
-    def _forward_impl(self, inputs, hidden_state):
-        """
-        forward pass.
-
-        :param inputs: shape (batch_size, timesteps, input_size)
-        :param hidden_state: (num_layers, batch_size, self.hidden_state_size) -> optional, zeros if not provided
-        :return: output: # shape (timesteps, batch_size, self.hidden_state_size)
-                 hidden_state # shape (num_layers, batch_size, self.hidden_state_size) (lower indices mean lower layers)
-        """
-        batch_size, timesteps, _ = inputs.size()
-        if hidden_state is None:
-            hidden_state = torch.zeros((self.num_rnn_layers, batch_size, self.hidden_state_size),
-                                       device=device)
-        output = torch.empty((timesteps, batch_size, self.hidden_state_size))
-        for t in range(timesteps):
-            hidden_state = self.t_step_forward_pass(hidden_state, inputs, output, t)
-
-        output = output.permute(1, 0, 2)
-        return output, hidden_state
-
-    @abstractmethod
-    def t_step_forward_pass(self, hidden_state, inputs, output, t):
-        """
-        Implements the forward pass for timestep t.
-        """
-        # this is to accommodate curriculum learning
-        pass
-
 
 class EncoderModel(nn.Module, DCRNNModel):
     def __init__(self, is_training, scaler, adj_mx, **model_kwargs):
@@ -79,14 +27,9 @@ class EncoderModel(nn.Module, DCRNNModel):
         # https://pytorch.org/docs/stable/nn.html#gru
         self.seq_len = int(model_kwargs.get('seq_len'))  # for the encoder
 
-    def t_step_forward_pass(self, hidden_state, inputs, output, t):
+    @property
-        cell_input = inputs[:, t, :]  # (batch_size, input_size)
-        cell_output, hidden_state = self._forward_cell(cell_input, hidden_state)
-        output[t] = cell_output
-        return hidden_state
-
     def dcgru_layers(self):
-        # input shape is supposed to be Input (batch_size, timesteps, num_sensor*input_dim)
+        # input shape is supposed to be Input (batch_size, num_sensor*input_dim)
         # first layer takes input shape and subsequent layer take input from the first layer
         return [nn.GRUCell(input_size=self.num_nodes * self.input_dim,
                            hidden_size=self.hidden_state_size,
@@ -99,12 +42,23 @@ class EncoderModel(nn.Module, DCRNNModel):
         """
         Encoder forward pass.
 
-        :param inputs: shape (batch_size, timesteps, num_nodes * input_dim)
+        :param inputs: shape (batch_size, self.num_nodes * self.input_dim)
         :param hidden_state: (num_layers, batch_size, self.hidden_state_size) -> optional, zeros if not provided
-        :return: output: # shape (timesteps, batch_size, self.hidden_state_size)
+        :return: output: # shape (batch_size, self.hidden_state_size)
                  hidden_state # shape (num_layers, batch_size, self.hidden_state_size) (lower indices mean lower layers)
         """
-        return self._forward_impl(inputs, hidden_state)
+        batch_size, _ = inputs.size()
+        if hidden_state is None:
+            hidden_state = torch.zeros((self.num_rnn_layers, batch_size, self.hidden_state_size),
+                                       device=device)
+        hidden_states = []
+        output = inputs
+        for layer_num, dcgru_layer in enumerate(self.dcgru_layers):
+            hidden_state = dcgru_layer(output, hidden_state)
+            hidden_states.append(hidden_state)
+            output = hidden_state
+
+        return output, torch.cat(hidden_states, dim=1)  # runs in O(num_layers) so not too slow # todo: check dim
 
 
 class DecoderModel(nn.Module, DCRNNModel):
@@ -115,6 +69,7 @@ class DecoderModel(nn.Module, DCRNNModel):
         self.horizon = int(model_kwargs.get('horizon', 1))  # for the decoder
         self.projection_layer = nn.Linear(self.hidden_state_size, self.num_nodes * self.output_dim)
 
+    @property
     def dcgru_layers(self):
         return [nn.GRUCell(input_size=self.num_nodes * self.output_dim,
                            hidden_size=self.hidden_state_size,
@@ -123,28 +78,25 @@ class DecoderModel(nn.Module, DCRNNModel):
                                                      bias=True) for _ in
                                           range(self.num_rnn_layers - 1)]
 
-    def t_step_forward_pass(self, hidden_state, inputs, output, t):
-        cell_input = inputs[:, t, :]  # (batch_size, input_size)
-
-        if self.is_training:
-            if t > 0 and self.use_curriculum_learning:
-                c = np.random.uniform(0, 1)
-                if c >= self._compute_sampling_threshold(): #todo
-                    cell_input = output[
-                        t - 1]  # todo: this won't work because the linear layer is applied after forward_impl
-
-        cell_output, hidden_state = self._forward_cell(cell_input, hidden_state)
-        output[t] = cell_output
-        return hidden_state
-
     def forward(self, inputs, hidden_state=None):
         """
         Decoder forward pass.
 
-        :param inputs: shape (batch_size, timesteps, num_nodes * input_dim)
+        :param inputs: shape (batch_size, self.num_nodes * self.output_dim)
         :param hidden_state: (num_layers, batch_size, self.hidden_state_size) -> optional, zeros if not provided
-        :return: output: # shape (timesteps, batch_size, self.num_nodes * self.output_dim)
+        :return: output: # shape (batch_size, self.num_nodes * self.output_dim)
                  hidden_state # shape (num_layers, batch_size, self.hidden_state_size) (lower indices mean lower layers)
         """
-        output, hidden = self._forward_impl(inputs, hidden_state)
+        batch_size, _ = inputs.size()
-        return self.projection_layer(output), hidden_state
+        if hidden_state is None:
+            hidden_state = torch.zeros((self.num_rnn_layers, batch_size, self.hidden_state_size),
+                                       device=device)
+        hidden_states = []
+        output = inputs
+        for layer_num, dcgru_layer in enumerate(self.dcgru_layers):
+            hidden_state = dcgru_layer(output, hidden_state)
+            hidden_states.append(hidden_state)
+            output = hidden_state
+
+        return self.projection_layer(output), torch.cat(hidden_states,
+                                                        dim=1)  # runs in O(num_layers) so not too slow #todo: check dim