Rough implementation complete - could forward pass it through the network
This commit is contained in:
Chintan Shah
parent
b65df994e4
commit
2e1836df40
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue