From b65df994e4370e261f5f7f52f73c6675b416f58b Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 11:55:02 -0400
Subject: [PATCH 01/14] Added dcrnn_cell

---
 model/pytorch/dcrnn_cell.py | 42 +++++++++++--------------------------
 1 file changed, 12 insertions(+), 30 deletions(-)

diff --git a/model/pytorch/dcrnn_cell.py b/model/pytorch/dcrnn_cell.py
index d8f045c..5863c99 100644
--- a/model/pytorch/dcrnn_cell.py
+++ b/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_):

From 2e1836df4081be823a245c861c47921ac283516c Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 13:24:37 -0400
Subject: [PATCH 02/14] Rough implementation complete - could forward pass it
 through the network

---
 model/pytorch/dcrnn_cell.py       | 30 +++++++++++++----------------
 model/pytorch/dcrnn_model.py      | 32 +++++++++++++------------------
 model/pytorch/dcrnn_supervisor.py |  1 +
 3 files changed, 27 insertions(+), 36 deletions(-)

diff --git a/model/pytorch/dcrnn_cell.py b/model/pytorch/dcrnn_cell.py
index 5863c99..1bbdf20 100644
--- a/model/pytorch/dcrnn_cell.py
+++ b/model/pytorch/dcrnn_cell.py
@@ -1,3 +1,4 @@
+import numpy as np
 import torch
 
 from lib import utils
@@ -12,7 +13,7 @@ class LayerParams:
 
     def get_weights(self, shape):
         if shape not in self._params_dict:
-            nn_param = torch.nn.init.xavier_normal(torch.empty(*shape))
+            nn_param = torch.nn.Parameter(torch.nn.init.xavier_normal(torch.empty(*shape)))
             self._params_dict[shape] = nn_param
             self._rnn_network.register_parameter('{}_weight_{}'.format(self._type, str(shape)),
                                                  nn_param)
@@ -20,7 +21,7 @@ class LayerParams:
 
     def get_biases(self, length, bias_start=0.0):
         if length not in self._biases_dict:
-            biases = torch.nn.init.constant(torch.empty(length), bias_start)
+            biases = torch.nn.Parameter(torch.nn.init.constant(torch.empty(length), bias_start))
             self._biases_dict[length] = biases
             self._rnn_network.register_parameter('{}_biases_{}'.format(self._type, str(length)),
                                                  biases)
@@ -65,16 +66,13 @@ class DCGRUCell(torch.nn.Module):
         self._fc_params = LayerParams(self, 'fc')
         self._gconv_params = LayerParams(self, 'gconv')
 
-    @property
-    def state_size(self):
-        return self._num_nodes * self._num_units
-
-    @property
-    def output_size(self):
-        output_size = self._num_nodes * self._num_units
-        if self._num_proj is not None:
-            output_size = self._num_nodes * self._num_proj
-        return output_size
+    @staticmethod
+    def _build_sparse_matrix(L):
+        L = L.tocoo()
+        indices = np.column_stack((L.row, L.col))
+        L = torch.sparse_coo_tensor(indices.T, L.data, L.shape)
+        return L
+        # return torch.sparse.sparse_reorder(L)
 
     def forward(self, inputs, hx):
         """Gated recurrent unit (GRU) with Graph Convolution.
@@ -86,14 +84,13 @@ class DCGRUCell(torch.nn.Module):
             the arity and shapes of `state`
         """
         output_size = 2 * self._num_units
-        # We start with bias of 1.0 to not reset and not update.
         if self._use_gc_for_ru:
             fn = self._gconv
         else:
             fn = self._fc
         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, u = torch.split(tensor=value, split_size_or_sections=self._num_units, dim=-1)
         r = torch.reshape(r, (-1, self._num_nodes * self._num_units))
         u = torch.reshape(u, (-1, self._num_nodes * self._num_units))
 
@@ -135,7 +132,7 @@ class DCGRUCell(torch.nn.Module):
         inputs = torch.reshape(inputs, (batch_size, self._num_nodes, -1))
         state = torch.reshape(state, (batch_size, self._num_nodes, -1))
         inputs_and_state = torch.cat([inputs, state], dim=2)
-        input_size = inputs_and_state.shape[2].value
+        input_size = inputs_and_state.size(2)
         dtype = inputs.dtype
 
         x = inputs_and_state
@@ -147,8 +144,7 @@ class DCGRUCell(torch.nn.Module):
             pass
         else:
             for support in self._supports:
-                # https://discuss.pytorch.org/t/sparse-x-dense-dense-matrix-multiplication/6116/7
-                x1 = torch.mm(support, x0)
+                x1 = torch.sparse.mm(support, x0)  # this is not reordered, does this work - todo
                 x = self._concat(x, x1)
 
                 for k in range(2, self._max_diffusion_step + 1):
diff --git a/model/pytorch/dcrnn_model.py b/model/pytorch/dcrnn_model.py
index 20ca2d8..36196a3 100644
--- a/model/pytorch/dcrnn_model.py
+++ b/model/pytorch/dcrnn_model.py
@@ -2,6 +2,8 @@ import numpy as np
 import torch
 import torch.nn as nn
 
+from model.pytorch.dcrnn_cell import DCGRUCell
+
 
 class Seq2SeqAttrs:
     def __init__(self, adj_mx, **model_kwargs):
@@ -9,7 +11,6 @@ class Seq2SeqAttrs:
         self.max_diffusion_step = int(model_kwargs.get('max_diffusion_step', 2))
         self.cl_decay_steps = int(model_kwargs.get('cl_decay_steps', 1000))
         self.filter_type = model_kwargs.get('filter_type', 'laplacian')
-        # self.max_grad_norm = float(model_kwargs.get('max_grad_norm', 5.0))
         self.num_nodes = int(model_kwargs.get('num_nodes', 1))
         self.num_rnn_layers = int(model_kwargs.get('num_rnn_layers', 1))
         self.rnn_units = int(model_kwargs.get('rnn_units'))
@@ -18,19 +19,13 @@ class Seq2SeqAttrs:
 
 class EncoderModel(nn.Module, Seq2SeqAttrs):
     def __init__(self, adj_mx, **model_kwargs):
-        # super().__init__(is_training, adj_mx, **model_kwargs)
-        # https://pytorch.org/docs/stable/nn.html#gru
         nn.Module.__init__(self)
         Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs)
         self.input_dim = int(model_kwargs.get('input_dim', 1))
         self.seq_len = int(model_kwargs.get('seq_len'))  # for the encoder
-        self.dcgru_layers = nn.ModuleList([nn.GRUCell(input_size=self.num_nodes * self.input_dim,
-                                                      hidden_size=self.hidden_state_size,
-                                                      bias=True)] + [
-                                              nn.GRUCell(input_size=self.hidden_state_size,
-                                                         hidden_size=self.hidden_state_size,
-                                                         bias=True) for _ in
-                                              range(self.num_rnn_layers - 1)])
+        self.dcgru_layers = nn.ModuleList(
+            [DCGRUCell(self.rnn_units, adj_mx, self.max_diffusion_step, self.num_nodes,
+                       filter_type=self.filter_type) for _ in range(self.num_rnn_layers)])
 
     def forward(self, inputs, hidden_state=None):
         """
@@ -63,14 +58,10 @@ class DecoderModel(nn.Module, Seq2SeqAttrs):
         Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs)
         self.output_dim = int(model_kwargs.get('output_dim', 1))
         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)
-        self.dcgru_layers = nn.ModuleList([nn.GRUCell(input_size=self.num_nodes * self.output_dim,
-                                                      hidden_size=self.hidden_state_size,
-                                                      bias=True)] + [
-                                              nn.GRUCell(input_size=self.hidden_state_size,
-                                                         hidden_size=self.hidden_state_size,
-                                                         bias=True) for _ in
-                                              range(self.num_rnn_layers - 1)])
+        self.projection_layer = nn.Linear(self.rnn_units, self.output_dim)
+        self.dcgru_layers = nn.ModuleList(
+            [DCGRUCell(self.rnn_units, adj_mx, self.max_diffusion_step, self.num_nodes,
+                       filter_type=self.filter_type) for _ in range(self.num_rnn_layers)])
 
     def forward(self, inputs, hidden_state=None):
         """
@@ -90,7 +81,10 @@ class DecoderModel(nn.Module, Seq2SeqAttrs):
             hidden_states.append(next_hidden_state)
             output = next_hidden_state
 
-        return self.projection_layer(output), torch.stack(hidden_states)
+        projected = self.projection_layer(output.view(-1, self.rnn_units))
+        output = projected.view(-1, self.num_nodes * self.output_dim)
+
+        return output, torch.stack(hidden_states)
 
 
 class DCRNNModel(nn.Module, Seq2SeqAttrs):
diff --git a/model/pytorch/dcrnn_supervisor.py b/model/pytorch/dcrnn_supervisor.py
index 512d370..cadcf24 100644
--- a/model/pytorch/dcrnn_supervisor.py
+++ b/model/pytorch/dcrnn_supervisor.py
@@ -36,6 +36,7 @@ class DCRNNSupervisor:
 
         # setup model
         dcrnn_model = DCRNNModel(adj_mx, self._logger, **self._model_kwargs)
+        print(dcrnn_model)
         self.dcrnn_model = dcrnn_model.cuda() if torch.cuda.is_available() else dcrnn_model
         self._logger.info("Model created")
 

From 9454fd91a2991e599829854efd03617a12da5d44 Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 13:44:55 -0400
Subject: [PATCH 03/14] Ensured sparse mm for readability, logging sparsely as
 well

---
 model/pytorch/dcrnn_cell.py       | 3 +--
 model/pytorch/dcrnn_model.py      | 4 ++--
 model/pytorch/dcrnn_supervisor.py | 4 +++-
 3 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/model/pytorch/dcrnn_cell.py b/model/pytorch/dcrnn_cell.py
index 1bbdf20..0779a72 100644
--- a/model/pytorch/dcrnn_cell.py
+++ b/model/pytorch/dcrnn_cell.py
@@ -133,7 +133,6 @@ class DCGRUCell(torch.nn.Module):
         state = torch.reshape(state, (batch_size, self._num_nodes, -1))
         inputs_and_state = torch.cat([inputs, state], dim=2)
         input_size = inputs_and_state.size(2)
-        dtype = inputs.dtype
 
         x = inputs_and_state
         x0 = x.permute(1, 2, 0)  # (num_nodes, total_arg_size, batch_size)
@@ -148,7 +147,7 @@ class DCGRUCell(torch.nn.Module):
                 x = self._concat(x, x1)
 
                 for k in range(2, self._max_diffusion_step + 1):
-                    x2 = 2 * torch.mm(support, x1) - x0
+                    x2 = 2 * torch.sparse.mm(support, x1) - x0
                     x = self._concat(x, x2)
                     x1, x0 = x2, x1
 
diff --git a/model/pytorch/dcrnn_model.py b/model/pytorch/dcrnn_model.py
index 36196a3..4a14334 100644
--- a/model/pytorch/dcrnn_model.py
+++ b/model/pytorch/dcrnn_model.py
@@ -149,7 +149,7 @@ class DCRNNModel(nn.Module, Seq2SeqAttrs):
         :return: output: (self.horizon, batch_size, self.num_nodes * self.output_dim)
         """
         encoder_hidden_state = self.encoder(inputs)
-        self._logger.info("Encoder complete, starting decoder")
+        self._logger.debug("Encoder complete, starting decoder")
         outputs = self.decoder(encoder_hidden_state, labels, batches_seen=batches_seen)
-        self._logger.info("Decoder complete")
+        self._logger.debug("Decoder complete")
         return outputs
diff --git a/model/pytorch/dcrnn_supervisor.py b/model/pytorch/dcrnn_supervisor.py
index cadcf24..67c8440 100644
--- a/model/pytorch/dcrnn_supervisor.py
+++ b/model/pytorch/dcrnn_supervisor.py
@@ -143,7 +143,9 @@ class DCRNNSupervisor:
 
                 output = self.dcrnn_model(x, y, batches_seen)
                 loss = self._compute_loss(y, output, criterion)
-                self._logger.info(loss.item())
+
+                self._logger.debug(loss.item())
+
                 losses.append(loss.item())
 
                 batches_seen += 1

From ba304e9f04bd03100b597a628102c9803e460dd8 Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 14:00:54 -0400
Subject: [PATCH 04/14] moving tensors to GPU

---
 model/pytorch/dcrnn_cell.py       | 8 ++++++--
 model/pytorch/dcrnn_supervisor.py | 5 +++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/model/pytorch/dcrnn_cell.py b/model/pytorch/dcrnn_cell.py
index 0779a72..2a0b85f 100644
--- a/model/pytorch/dcrnn_cell.py
+++ b/model/pytorch/dcrnn_cell.py
@@ -3,6 +3,8 @@ import torch
 
 from lib import utils
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 
 class LayerParams:
     def __init__(self, rnn_network: torch.nn.Module, layer_type: str):
@@ -13,7 +15,8 @@ class LayerParams:
 
     def get_weights(self, shape):
         if shape not in self._params_dict:
-            nn_param = torch.nn.Parameter(torch.nn.init.xavier_normal(torch.empty(*shape)))
+            nn_param = torch.nn.Parameter(torch.empty(*shape, device=device))
+            torch.nn.init.xavier_normal_(nn_param)
             self._params_dict[shape] = nn_param
             self._rnn_network.register_parameter('{}_weight_{}'.format(self._type, str(shape)),
                                                  nn_param)
@@ -21,7 +24,8 @@ class LayerParams:
 
     def get_biases(self, length, bias_start=0.0):
         if length not in self._biases_dict:
-            biases = torch.nn.Parameter(torch.nn.init.constant(torch.empty(length), bias_start))
+            biases = torch.nn.Parameter(torch.empty(length, device=device))
+            torch.nn.init.constant_(biases, bias_start)
             self._biases_dict[length] = biases
             self._rnn_network.register_parameter('{}_biases_{}'.format(self._type, str(length)),
                                                  biases)
diff --git a/model/pytorch/dcrnn_supervisor.py b/model/pytorch/dcrnn_supervisor.py
index 67c8440..6f3cdd2 100644
--- a/model/pytorch/dcrnn_supervisor.py
+++ b/model/pytorch/dcrnn_supervisor.py
@@ -7,6 +7,8 @@ import torch
 from lib import utils
 from model.pytorch.dcrnn_model import DCRNNModel
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 
 class DCRNNSupervisor:
     def __init__(self, adj_mx, **kwargs):
@@ -36,7 +38,6 @@ class DCRNNSupervisor:
 
         # setup model
         dcrnn_model = DCRNNModel(adj_mx, self._logger, **self._model_kwargs)
-        print(dcrnn_model)
         self.dcrnn_model = dcrnn_model.cuda() if torch.cuda.is_available() else dcrnn_model
         self._logger.info("Model created")
 
@@ -215,7 +216,7 @@ class DCRNNSupervisor:
         x = x.view(self.seq_len, batch_size, self.num_nodes * self.input_dim)
         y = y[..., :self.output_dim].view(self.horizon, batch_size,
                                           self.num_nodes * self.output_dim)
-        return x, y
+        return x.to(device), y.to(device)
 
     def _compute_loss(self, y_true, y_predicted, criterion):
         loss = 0

From 017ec70783f8247e1f70a5ac203970530290423a Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 14:10:20 -0400
Subject: [PATCH 05/14] moving tensors to GPU [v2]

---
 model/pytorch/dcrnn_model.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/model/pytorch/dcrnn_model.py b/model/pytorch/dcrnn_model.py
index 4a14334..9390c6e 100644
--- a/model/pytorch/dcrnn_model.py
+++ b/model/pytorch/dcrnn_model.py
@@ -4,6 +4,8 @@ import torch.nn as nn
 
 from model.pytorch.dcrnn_cell import DCGRUCell
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 
 class Seq2SeqAttrs:
     def __init__(self, adj_mx, **model_kwargs):
@@ -40,7 +42,8 @@ class EncoderModel(nn.Module, Seq2SeqAttrs):
         """
         batch_size, _ = inputs.size()
         if hidden_state is None:
-            hidden_state = torch.zeros((self.num_rnn_layers, batch_size, self.hidden_state_size))
+            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):
@@ -122,7 +125,8 @@ class DCRNNModel(nn.Module, Seq2SeqAttrs):
         :return: output: (self.horizon, batch_size, self.num_nodes * self.output_dim)
         """
         batch_size = encoder_hidden_state.size(1)
-        go_symbol = torch.zeros((batch_size, self.num_nodes * self.decoder_model.output_dim))
+        go_symbol = torch.zeros((batch_size, self.num_nodes * self.decoder_model.output_dim),
+                                device=device)
         decoder_hidden_state = encoder_hidden_state
         decoder_input = go_symbol
 

From e563e1bf37182290866a1a1b1f5a83e8bf56efd5 Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 14:13:02 -0400
Subject: [PATCH 06/14] moving tensors to GPU [v3]

---
 model/pytorch/dcrnn_cell.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/model/pytorch/dcrnn_cell.py b/model/pytorch/dcrnn_cell.py
index 2a0b85f..7550501 100644
--- a/model/pytorch/dcrnn_cell.py
+++ b/model/pytorch/dcrnn_cell.py
@@ -74,7 +74,7 @@ class DCGRUCell(torch.nn.Module):
     def _build_sparse_matrix(L):
         L = L.tocoo()
         indices = np.column_stack((L.row, L.col))
-        L = torch.sparse_coo_tensor(indices.T, L.data, L.shape)
+        L = torch.sparse_coo_tensor(indices.T, L.data, L.shape, device=device)
         return L
         # return torch.sparse.sparse_reorder(L)
 

From 31acadedce9a17fe97148b3dee783d48e989654f Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 14:29:28 -0400
Subject: [PATCH 07/14] logging and refactor

---
 model/pytorch/dcrnn_supervisor.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/model/pytorch/dcrnn_supervisor.py b/model/pytorch/dcrnn_supervisor.py
index 6f3cdd2..592454b 100644
--- a/model/pytorch/dcrnn_supervisor.py
+++ b/model/pytorch/dcrnn_supervisor.py
@@ -104,8 +104,7 @@ class DCRNNSupervisor:
             criterion = torch.nn.L1Loss()
 
             for _, (x, y) in enumerate(val_iterator):
-                x, y = self._get_x_y(x, y)
-                x, y = self._get_x_y_in_correct_dims(x, y)
+                x, y = self._prepare_data(x, y)
 
                 output = self.dcrnn_model(x)
                 loss = self._compute_loss(y, output, criterion)
@@ -139,13 +138,12 @@ class DCRNNSupervisor:
             for _, (x, y) in enumerate(train_iterator):
                 optimizer.zero_grad()
 
-                x, y = self._get_x_y(x, y)
-                x, y = self._get_x_y_in_correct_dims(x, y)
+                x, y = self._prepare_data(x, y)
 
                 output = self.dcrnn_model(x, y, batches_seen)
                 loss = self._compute_loss(y, output, criterion)
 
-                self._logger.debug(loss.item())
+                self._logger.info(loss.item())
 
                 losses.append(loss.item())
 
@@ -156,6 +154,7 @@ class DCRNNSupervisor:
                 torch.nn.utils.clip_grad_norm_(self.dcrnn_model.parameters(), self.max_grad_norm)
 
                 optimizer.step()
+                self._logger.info("finished one batch in {}".format(time.time() - start_time))
 
             lr_scheduler.step()
 
@@ -190,6 +189,11 @@ class DCRNNSupervisor:
                     self._logger.warning('Early stopping at epoch: %d' % epoch_num)
                     break
 
+    def _prepare_data(self, x, y):
+        x, y = self._get_x_y(x, y)
+        x, y = self._get_x_y_in_correct_dims(x, y)
+        return x.to(device), y.to(device)
+
     def _get_x_y(self, x, y):
         """
         :param x: shape (batch_size, seq_len, num_sensor, input_dim)
@@ -216,7 +220,7 @@ class DCRNNSupervisor:
         x = x.view(self.seq_len, batch_size, self.num_nodes * self.input_dim)
         y = y[..., :self.output_dim].view(self.horizon, batch_size,
                                           self.num_nodes * self.output_dim)
-        return x.to(device), y.to(device)
+        return x, y
 
     def _compute_loss(self, y_true, y_predicted, criterion):
         loss = 0

From 036e552bf6c4df64307751165350e22a06819b03 Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 14:31:46 -0400
Subject: [PATCH 08/14] logging and refactor

---
 model/pytorch/dcrnn_supervisor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/model/pytorch/dcrnn_supervisor.py b/model/pytorch/dcrnn_supervisor.py
index 592454b..84d0938 100644
--- a/model/pytorch/dcrnn_supervisor.py
+++ b/model/pytorch/dcrnn_supervisor.py
@@ -154,7 +154,7 @@ class DCRNNSupervisor:
                 torch.nn.utils.clip_grad_norm_(self.dcrnn_model.parameters(), self.max_grad_norm)
 
                 optimizer.step()
-                self._logger.info("finished one batch in {}".format(time.time() - start_time))
+                self._logger.info("finished one batch in {:.1f}s".format(time.time() - start_time))
 
             lr_scheduler.step()
 

From b6a2b3fe8eeb3d9573a5cdf80dbac7ec772f8750 Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 14:34:58 -0400
Subject: [PATCH 09/14] logging and refactor

---
 model/pytorch/dcrnn_supervisor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/model/pytorch/dcrnn_supervisor.py b/model/pytorch/dcrnn_supervisor.py
index 84d0938..f593c50 100644
--- a/model/pytorch/dcrnn_supervisor.py
+++ b/model/pytorch/dcrnn_supervisor.py
@@ -129,6 +129,7 @@ class DCRNNSupervisor:
         self.dcrnn_model = self.dcrnn_model.train()
 
         self._logger.info('Start training ...')
+        self._logger.info("num_batches:".format(self._data['train_loader'].num_batch))
         for epoch_num in range(epochs):
             train_iterator = self._data['train_loader'].get_iterator()
             losses = []
@@ -154,7 +155,6 @@ class DCRNNSupervisor:
                 torch.nn.utils.clip_grad_norm_(self.dcrnn_model.parameters(), self.max_grad_norm)
 
                 optimizer.step()
-                self._logger.info("finished one batch in {:.1f}s".format(time.time() - start_time))
 
             lr_scheduler.step()
 

From 02c4681ad9742be1eaad12d9a18f18e245473051 Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 14:38:44 -0400
Subject: [PATCH 10/14] logging and refactor

---
 model/pytorch/dcrnn_supervisor.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/model/pytorch/dcrnn_supervisor.py b/model/pytorch/dcrnn_supervisor.py
index f593c50..f1692be 100644
--- a/model/pytorch/dcrnn_supervisor.py
+++ b/model/pytorch/dcrnn_supervisor.py
@@ -129,7 +129,7 @@ class DCRNNSupervisor:
         self.dcrnn_model = self.dcrnn_model.train()
 
         self._logger.info('Start training ...')
-        self._logger.info("num_batches:".format(self._data['train_loader'].num_batch))
+        self._logger.info("num_batches:{}".format(self._data['train_loader'].num_batch))
         for epoch_num in range(epochs):
             train_iterator = self._data['train_loader'].get_iterator()
             losses = []
@@ -144,7 +144,7 @@ class DCRNNSupervisor:
                 output = self.dcrnn_model(x, y, batches_seen)
                 loss = self._compute_loss(y, output, criterion)
 
-                self._logger.info(loss.item())
+                self._logger.debug(loss.item())
 
                 losses.append(loss.item())
 

From 55a087ac9f39795b3dec1361c4548c4c09c943f9 Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 14:40:53 -0400
Subject: [PATCH 11/14] logging and refactor

---
 model/pytorch/dcrnn_supervisor.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/model/pytorch/dcrnn_supervisor.py b/model/pytorch/dcrnn_supervisor.py
index f1692be..6a8462a 100644
--- a/model/pytorch/dcrnn_supervisor.py
+++ b/model/pytorch/dcrnn_supervisor.py
@@ -155,9 +155,9 @@ class DCRNNSupervisor:
                 torch.nn.utils.clip_grad_norm_(self.dcrnn_model.parameters(), self.max_grad_norm)
 
                 optimizer.step()
-
+            self._logger.info("epoch complete")
             lr_scheduler.step()
-
+            self._logger.info("evaluating now!")
             val_loss = self.evaluate(dataset='val')
             end_time = time.time()
             if epoch_num % log_every == 0:

From ec5d9555a52ff61b7feed315ea78cdad358b6987 Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 15:15:11 -0400
Subject: [PATCH 12/14] logging and refactor

---
 model/pytorch/dcrnn_supervisor.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/model/pytorch/dcrnn_supervisor.py b/model/pytorch/dcrnn_supervisor.py
index 6a8462a..c92dad3 100644
--- a/model/pytorch/dcrnn_supervisor.py
+++ b/model/pytorch/dcrnn_supervisor.py
@@ -163,15 +163,15 @@ class DCRNNSupervisor:
             if epoch_num % log_every == 0:
                 message = 'Epoch [{}/{}] ({}) train_mae: {:.4f}, val_mae: {:.4f}, lr: {:.6f}' \
                           '{:.1f}s'.format(epoch_num, epochs, batches_seen,
-                                           np.mean(losses), val_loss, lr_scheduler.get_lr(),
+                                           np.mean(losses), val_loss, lr_scheduler.get_lr()[0],
                                            (end_time - start_time))
                 self._logger.info(message)
 
             if epoch_num % test_every_n_epochs == 0:
                 test_loss = self.evaluate(dataset='test')
-                message = 'Epoch [{}/{}] ({}) train_mae: {:.4f}, test_mae: {:.4f},  lr: {:.6f} ' \
+                message = 'Epoch [{}/{}] ({}) train_mae: {:.4f}, test_mae: {:.4f},  lr: {:.6f}, ' \
                           '{:.1f}s'.format(epoch_num, epochs, batches_seen,
-                                           np.mean(losses), test_loss, lr_scheduler.get_lr(),
+                                           np.mean(losses), test_loss, lr_scheduler.get_lr()[0],
                                            (end_time - start_time))
                 self._logger.info(message)
 

From 6331173f443603742128908ba96c15227109c7cf Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 15:22:57 -0400
Subject: [PATCH 13/14] logging and refactor

---
 model/pytorch/dcrnn_supervisor.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/model/pytorch/dcrnn_supervisor.py b/model/pytorch/dcrnn_supervisor.py
index c92dad3..a44e30d 100644
--- a/model/pytorch/dcrnn_supervisor.py
+++ b/model/pytorch/dcrnn_supervisor.py
@@ -161,7 +161,7 @@ class DCRNNSupervisor:
             val_loss = self.evaluate(dataset='val')
             end_time = time.time()
             if epoch_num % log_every == 0:
-                message = 'Epoch [{}/{}] ({}) train_mae: {:.4f}, val_mae: {:.4f}, lr: {:.6f}' \
+                message = 'Epoch [{}/{}] ({}) train_mae: {:.4f}, val_mae: {:.4f}, lr: {:.6f}, ' \
                           '{:.1f}s'.format(epoch_num, epochs, batches_seen,
                                            np.mean(losses), val_loss, lr_scheduler.get_lr()[0],
                                            (end_time - start_time))
@@ -177,12 +177,13 @@ class DCRNNSupervisor:
 
             if val_loss < min_val_loss:
                 wait = 0
-                min_val_loss = val_loss
                 if save_model:
                     model_file_name = self.save_model(epoch_num)
                     self._logger.info(
                         'Val loss decrease from {:.4f} to {:.4f}, '
                         'saving to {}'.format(min_val_loss, val_loss, model_file_name))
+                min_val_loss = val_loss
+
             elif val_loss >= min_val_loss:
                 wait += 1
                 if wait == patience:

From d46b605a65fd7b72f99c825b87cc237b292da36c Mon Sep 17 00:00:00 2001
From: Chintan Shah <chnsh24@gmail.com>
Date: Sun, 6 Oct 2019 15:53:14 -0400
Subject: [PATCH 14/14] ensured row major ordering

---
 model/pytorch/dcrnn_cell.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/model/pytorch/dcrnn_cell.py b/model/pytorch/dcrnn_cell.py
index 7550501..b790730 100644
--- a/model/pytorch/dcrnn_cell.py
+++ b/model/pytorch/dcrnn_cell.py
@@ -74,9 +74,10 @@ class DCGRUCell(torch.nn.Module):
     def _build_sparse_matrix(L):
         L = L.tocoo()
         indices = np.column_stack((L.row, L.col))
+        # this is to ensure row-major ordering to equal torch.sparse.sparse_reorder(L)
+        indices = indices[np.lexsort((indices[:, 0], indices[:, 1]))]
         L = torch.sparse_coo_tensor(indices.T, L.data, L.shape, device=device)
         return L
-        # return torch.sparse.sparse_reorder(L)
 
     def forward(self, inputs, hx):
         """Gated recurrent unit (GRU) with Graph Convolution.