czzhangheng
/
DCRNN
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Rough implementation complete - could forward pass it through the network

This commit is contained in:
Chintan Shah 2019-10-06 13:24:37 -04:00
parent b65df994e4
commit 2e1836df40
3 changed files with 27 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
model/pytorch/dcrnn_cell.py
View File

@ -1,3 +1,4 @@
import numpy as np
import torch import torch
from lib import utils from lib import utils
@ -12,7 +13,7 @@ class LayerParams:
def get_weights(self, shape): def get_weights(self, shape):
if shape not in self._params_dict: 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._params_dict[shape] = nn_param
self._rnn_network.register_parameter('{}_weight_{}'.format(self._type, str(shape)), self._rnn_network.register_parameter('{}_weight_{}'.format(self._type, str(shape)),
nn_param) nn_param)
@ -20,7 +21,7 @@ class LayerParams:
def get_biases(self, length, bias_start=0.0): def get_biases(self, length, bias_start=0.0):
if length not in self._biases_dict: 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._biases_dict[length] = biases
self._rnn_network.register_parameter('{}_biases_{}'.format(self._type, str(length)), self._rnn_network.register_parameter('{}_biases_{}'.format(self._type, str(length)),
biases) biases)
@ -65,16 +66,13 @@ class DCGRUCell(torch.nn.Module):
self._fc_params = LayerParams(self, 'fc') self._fc_params = LayerParams(self, 'fc')
self._gconv_params = LayerParams(self, 'gconv') 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 L = torch.sparse_coo_tensor(indices.T, L.data, L.shape)
def output_size(self): return L
output_size = self._num_nodes * self._num_units # return torch.sparse.sparse_reorder(L)
if self._num_proj is not None:
output_size = self._num_nodes * self._num_proj
return output_size
def forward(self, inputs, hx): def forward(self, inputs, hx):
"""Gated recurrent unit (GRU) with Graph Convolution. """Gated recurrent unit (GRU) with Graph Convolution.
@ -86,14 +84,13 @@ class DCGRUCell(torch.nn.Module):
the arity and shapes of `state` the arity and shapes of `state`
""" """
output_size = 2 * self._num_units 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: if self._use_gc_for_ru:
fn = self._gconv fn = self._gconv
else: else:
fn = self._fc fn = self._fc
value = torch.sigmoid(fn(inputs, 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)) 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)) r = torch.reshape(r, (-1, self._num_nodes * self._num_units))
u = torch.reshape(u, (-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)) inputs = torch.reshape(inputs, (batch_size, self._num_nodes, -1))
state = torch.reshape(state, (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) 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 dtype = inputs.dtype
x = inputs_and_state x = inputs_and_state
@ -147,8 +144,7 @@ class DCGRUCell(torch.nn.Module):
pass pass
else: else:
for support in self._supports: for support in self._supports:
# https://discuss.pytorch.org/t/sparse-x-dense-dense-matrix-multiplication/6116/7 x1 = torch.sparse.mm(support, x0) # this is not reordered, does this work - todo
x1 = torch.mm(support, x0)
x = self._concat(x, x1) x = self._concat(x, x1)
for k in range(2, self._max_diffusion_step + 1): for k in range(2, self._max_diffusion_step + 1):

32
model/pytorch/dcrnn_model.py
View File

@ -2,6 +2,8 @@ import numpy as np
import torch import torch
import torch.nn as nn import torch.nn as nn
from model.pytorch.dcrnn_cell import DCGRUCell
class Seq2SeqAttrs: class Seq2SeqAttrs:
def __init__(self, adj_mx, **model_kwargs): 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.max_diffusion_step = int(model_kwargs.get('max_diffusion_step', 2))
self.cl_decay_steps = int(model_kwargs.get('cl_decay_steps', 1000)) self.cl_decay_steps = int(model_kwargs.get('cl_decay_steps', 1000))
self.filter_type = model_kwargs.get('filter_type', 'laplacian') 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_nodes = int(model_kwargs.get('num_nodes', 1))
self.num_rnn_layers = int(model_kwargs.get('num_rnn_layers', 1)) self.num_rnn_layers = int(model_kwargs.get('num_rnn_layers', 1))
self.rnn_units = int(model_kwargs.get('rnn_units')) self.rnn_units = int(model_kwargs.get('rnn_units'))
@ -18,19 +19,13 @@ class Seq2SeqAttrs:
class EncoderModel(nn.Module, Seq2SeqAttrs): class EncoderModel(nn.Module, Seq2SeqAttrs):
def __init__(self, adj_mx, **model_kwargs): 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) nn.Module.__init__(self)
Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs) Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs)
self.input_dim = int(model_kwargs.get('input_dim', 1)) self.input_dim = int(model_kwargs.get('input_dim', 1))
self.seq_len = int(model_kwargs.get('seq_len')) # for the encoder 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): def forward(self, inputs, hidden_state=None):
""" """
@ -63,14 +58,10 @@ class DecoderModel(nn.Module, Seq2SeqAttrs):
Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs) Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs)
self.output_dim = int(model_kwargs.get('output_dim', 1)) self.output_dim = int(model_kwargs.get('output_dim', 1))
self.horizon = int(model_kwargs.get('horizon', 1)) # for the decoder 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): def forward(self, inputs, hidden_state=None):
""" """
@ -90,7 +81,10 @@ class DecoderModel(nn.Module, Seq2SeqAttrs):
hidden_states.append(next_hidden_state) hidden_states.append(next_hidden_state)
output = 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): class DCRNNModel(nn.Module, Seq2SeqAttrs):

1
model/pytorch/dcrnn_supervisor.py
View File

@ -36,6 +36,7 @@ class DCRNNSupervisor:
# setup model # setup model
dcrnn_model = DCRNNModel(adj_mx, self._logger, **self._model_kwargs) 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.dcrnn_model = dcrnn_model.cuda() if torch.cuda.is_available() else dcrnn_model
self._logger.info("Model created") self._logger.info("Model created")