From 2e1836df4081be823a245c861c47921ac283516c Mon Sep 17 00:00:00 2001 From: Chintan Shah Date: Sun, 6 Oct 2019 13:24:37 -0400 Subject: [PATCH] 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")