Added dcrnn_cell

model/pytorch/dcrnn_cell.py

@@ -1,17 +1,14 @@
-from typing import Optional
-
 import torch
-from torch import Tensor
 
 from lib import utils
 
 
 class LayerParams:
-    def __init__(self, rnn_network: torch.nn.RNN, type: str):
+    def __init__(self, rnn_network: torch.nn.Module, layer_type: str):
         self._rnn_network = rnn_network
         self._params_dict = {}
         self._biases_dict = {}
-        self._type = type
+        self._type = layer_type
 
     def get_weights(self, shape):
         if shape not in self._params_dict:
@@ -31,32 +28,24 @@ class LayerParams:
         return self._biases_dict[length]
 
 
-class DCGRUCell(torch.nn.RNN):
-    def __init__(self, num_units, adj_mx, max_diffusion_step, num_nodes, input_size: int,
-                 hidden_size: int,
-                 num_layers: int = 1,
-                 num_proj=None,
-                 nonlinearity='tanh', filter_type="laplacian", use_gc_for_ru=True):
+class DCGRUCell(torch.nn.Module):
+    def __init__(self, num_units, adj_mx, max_diffusion_step, num_nodes, nonlinearity='tanh',
+                 filter_type="laplacian", use_gc_for_ru=True):
         """
 
         :param num_units:
         :param adj_mx:
         :param max_diffusion_step:
         :param num_nodes:
-        :param input_size:
-        :param num_proj:
         :param nonlinearity:
         :param filter_type: "laplacian", "random_walk", "dual_random_walk".
         :param use_gc_for_ru: whether to use Graph convolution to calculate the reset and update gates.
         """
-        super(DCGRUCell, self).__init__(input_size, hidden_size, bias=True,
-                                        # bias param does not exist in tf code?
-                                        num_layers=num_layers,
-                                        nonlinearity=nonlinearity)
+
+        super().__init__()
         self._activation = torch.tanh if nonlinearity == 'tanh' else torch.relu
         # support other nonlinearities up here?
         self._num_nodes = num_nodes
-        self._num_proj = num_proj
         self._num_units = num_units
         self._max_diffusion_step = max_diffusion_step
         self._supports = []
@@ -73,8 +62,6 @@ class DCGRUCell(torch.nn.RNN):
             supports.append(utils.calculate_scaled_laplacian(adj_mx))
         for support in supports:
             self._supports.append(self._build_sparse_matrix(support))
-
-        self._proj_weights = torch.nn.Parameter(torch.randn(self._num_units, self._num_proj))
         self._fc_params = LayerParams(self, 'fc')
         self._gconv_params = LayerParams(self, 'gconv')
 
@@ -89,7 +76,7 @@ class DCGRUCell(torch.nn.RNN):
             output_size = self._num_nodes * self._num_proj
         return output_size
 
-    def forward(self, input: Tensor, hx: Optional[Tensor] = ...):
+    def forward(self, inputs, hx):
         """Gated recurrent unit (GRU) with Graph Convolution.
         :param input: (B, num_nodes * input_dim)
 
@@ -104,23 +91,18 @@ class DCGRUCell(torch.nn.RNN):
             fn = self._gconv
         else:
             fn = self._fc
-        value = torch.sigmoid(fn(input, hx, output_size, bias_start=1.0))
+        value = torch.sigmoid(fn(inputs, hx, output_size, bias_start=1.0))
         value = torch.reshape(value, (-1, self._num_nodes, output_size))
         r, u = torch.split(tensor=value, split_size_or_sections=2, dim=-1)
         r = torch.reshape(r, (-1, self._num_nodes * self._num_units))
         u = torch.reshape(u, (-1, self._num_nodes * self._num_units))
 
-        c = self._gconv(input, r * hx, self._num_units)
+        c = self._gconv(inputs, r * hx, self._num_units)
         if self._activation is not None:
             c = self._activation(c)
 
-        output = new_state = u * hx + (1 - u) * c
-        if self._num_proj is not None:
-            batch_size = input.shape[0]
-            output = torch.reshape(new_state, shape=(-1, self._num_units))
-            output = torch.reshape(torch.matmul(output, self._proj_weights),
-                                   shape=(batch_size, self.output_size))
-        return output, new_state
+        new_state = u * hx + (1.0 - u) * c
+        return new_state
 
     @staticmethod
     def _concat(x, x_):