From bdce241a8f4b1e6478a8080f21021bb116d856a7 Mon Sep 17 00:00:00 2001
From: Chintan Shah <chintan.shah@apprenticehealth.com>
Date: Mon, 30 Sep 2019 20:32:31 -0400
Subject: [PATCH] Implemented abstract method and changed scheme to do all
 layers first for each timestep

---
 model/pytorch/dcrnn_model.py | 57 +++++++++++++++++++++++-------------
 1 file changed, 37 insertions(+), 20 deletions(-)

diff --git a/model/pytorch/dcrnn_model.py b/model/pytorch/dcrnn_model.py
index 10e5783..6b8b489 100644
--- a/model/pytorch/dcrnn_model.py
+++ b/model/pytorch/dcrnn_model.py
@@ -22,43 +22,54 @@ class DCRNNModel(metaclass=ABC):
         self.hidden_state_size = self.num_nodes * self.rnn_units
 
     @abstractmethod
+    @property
     def dcgru_layers(self):
         pass
 
-    @staticmethod
-    def _forward_layer(inputs, dcgru_layer, hidden_state):
-        # inputs shape = (timesteps, batch_size, input_size)
-        outputs = []
-        for cell_input in inputs[:, ]:
-            hidden_state = dcgru_layer(cell_input, hidden_state)
-            outputs.append(hidden_state)
+    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 torch.cat(outputs, dim=1)  # runs in O(timesteps) not too slow
+        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, num_nodes * input_dim)
+        :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)
         """
-        layer_input = inputs.permute(1, 0, 2)  # first axis is now timesteps
+        batch_size, timesteps, _ = inputs.size()
         if hidden_state is None:
-            batch_size = inputs.size()[0]
             hidden_state = torch.zeros((self.num_rnn_layers, batch_size, self.hidden_state_size),
                                        device=device)
-        hidden = torch.empty_like(hidden_state)
-        # noinspection PyTypeChecker
-        for layer_num, dcgru_layer in enumerate(self.dcgru_layers):
-            layer_states = self._forward_layer(layer_input, dcgru_layer, hidden_state[layer_num])
-            # append last time step's hidden state
-            hidden[layer_num] = layer_states[-1]
-            layer_input = layer_states
+        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 = layer_input  # last layer's output
-        return output, hidden
+        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):
@@ -67,6 +78,12 @@ 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):
+        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)
         # first layer takes input shape and subsequent layer take input from the first layer