squash! Added dcrnn_cell

model/pytorch/dcrnn_cell.py

@@ -1,7 +1,10 @@
+import numpy as np
 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):
@@ -12,7 +15,8 @@ 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.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)
@@ -20,7 +24,8 @@ 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.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)
@@ -65,16 +70,14 @@ class DCGRUCell(torch.nn.Module):
         self._fc_params = LayerParams(self, 'fc')
         self._gconv_params = LayerParams(self, 'gconv')
 
-    @property
+    @staticmethod
-    def state_size(self):
+    def _build_sparse_matrix(L):
-        return self._num_nodes * self._num_units
+        L = L.tocoo()
-
+        indices = np.column_stack((L.row, L.col))
-    @property
+        # this is to ensure row-major ordering to equal torch.sparse.sparse_reorder(L)
-    def output_size(self):
+        indices = indices[np.lexsort((indices[:, 0], indices[:, 1]))]
-        output_size = self._num_nodes * self._num_units
+        L = torch.sparse_coo_tensor(indices.T, L.data, L.shape, device=device)
-        if self._num_proj is not None:
+        return L
-            output_size = self._num_nodes * self._num_proj
-        return output_size
 
     def forward(self, inputs, hx):
         """Gated recurrent unit (GRU) with Graph Convolution.
@@ -86,14 +89,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,8 +137,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
         x0 = x.permute(1, 2, 0)  # (num_nodes, total_arg_size, batch_size)
@@ -147,12 +148,11 @@ 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.sparse.mm(support, x0)
-                x1 = torch.mm(support, x0)
                 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
 

model/pytorch/dcrnn_model.py

@@ -2,6 +2,10 @@ import numpy as np
 import torch
 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):
@@ -9,7 +13,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 +21,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,
+        self.dcgru_layers = nn.ModuleList(
-                                                      hidden_size=self.hidden_state_size,
+            [DCGRUCell(self.rnn_units, adj_mx, self.max_diffusion_step, self.num_nodes,
-                                                      bias=True)] + [
+                       filter_type=self.filter_type) for _ in range(self.num_rnn_layers)])
-                                              nn.GRUCell(input_size=self.hidden_state_size,
-                                                         hidden_size=self.hidden_state_size,
-                                                         bias=True) for _ in
-                                              range(self.num_rnn_layers - 1)])
 
     def forward(self, inputs, hidden_state=None):
         """
@@ -45,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):
@@ -63,14 +61,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.projection_layer = nn.Linear(self.rnn_units, self.output_dim)
-        self.dcgru_layers = nn.ModuleList([nn.GRUCell(input_size=self.num_nodes * self.output_dim,
+        self.dcgru_layers = nn.ModuleList(
-                                                      hidden_size=self.hidden_state_size,
+            [DCGRUCell(self.rnn_units, adj_mx, self.max_diffusion_step, self.num_nodes,
-                                                      bias=True)] + [
+                       filter_type=self.filter_type) for _ in range(self.num_rnn_layers)])
-                                              nn.GRUCell(input_size=self.hidden_state_size,
-                                                         hidden_size=self.hidden_state_size,
-                                                         bias=True) for _ in
-                                              range(self.num_rnn_layers - 1)])
 
     def forward(self, inputs, hidden_state=None):
         """
@@ -90,7 +84,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):
@@ -128,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
 
@@ -155,7 +153,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

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):
@@ -75,7 +77,7 @@ class DCRNNSupervisor:
         config['model_state_dict'] = self.dcrnn_model.state_dict()
         config['epoch'] = epoch
         torch.save(config, self._log_dir + 'models/epo%d.tar' % epoch)
-        self._logger.info("Loaded model at {}".format(epoch))
+        self._logger.info("Saved model at {}".format(epoch))
         return self._log_dir + 'models/epo%d.tar' % epoch
 
     def load_model(self, epoch):
@@ -102,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._prepare_data(x, y)
-                x, y = self._get_x_y_in_correct_dims(x, y)
 
                 output = self.dcrnn_model(x)
                 loss = self._compute_loss(y, output, criterion)
@@ -128,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 = []
@@ -137,12 +139,13 @@ class DCRNNSupervisor:
             for _, (x, y) in enumerate(train_iterator):
                 optimizer.zero_grad()
 
-                x, y = self._get_x_y(x, y)
+                x, y = self._prepare_data(x, y)
-                x, y = self._get_x_y_in_correct_dims(x, y)
 
                 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
@@ -152,40 +155,46 @@ 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:
-                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(),
+                                           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)
 
             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:
                     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)