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

Cleaned up code, fixed bugs in implementation, seems like it starts training with GRU

This commit is contained in:
Chintan Shah 2019-10-02 22:20:43 -04:00
parent f96a8c0d59
commit 9834b12d5a
2 changed files with 63 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

90
model/pytorch/dcrnn_model.py
View File

@ -4,25 +4,8 @@ import torch.nn as nn
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# class DCRNNModel: class DCRNNModel:
# def __init__(self, is_training, adj_mx, **model_kwargs):
# self.adj_mx = adj_mx
# self.is_training = is_training
# 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'))
# self.hidden_state_size = self.num_nodes * self.rnn_units
class EncoderModel(nn.Module):
def __init__(self, is_training, adj_mx, **model_kwargs): def __init__(self, is_training, adj_mx, **model_kwargs):
# super().__init__(is_training, adj_mx, **model_kwargs)
# https://pytorch.org/docs/stable/nn.html#gru
super().__init__()
self.adj_mx = adj_mx self.adj_mx = adj_mx
self.is_training = is_training self.is_training = is_training
self.max_diffusion_step = int(model_kwargs.get('max_diffusion_step', 2)) self.max_diffusion_step = int(model_kwargs.get('max_diffusion_step', 2))
@ -33,22 +16,34 @@ class EncoderModel(nn.Module):
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'))
self.hidden_state_size = self.num_nodes * self.rnn_units self.hidden_state_size = self.num_nodes * self.rnn_units
class EncoderModel(nn.Module, DCRNNModel):
def __init__(self, is_training, adj_mx, **model_kwargs):
# super().__init__(is_training, adj_mx, **model_kwargs)
# https://pytorch.org/docs/stable/nn.html#gru
nn.Module.__init__(self)
DCRNNModel.__init__(self, is_training, 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.GRUCell(input_size=self.num_nodes * self.input_dim, self.dcgru_layers = nn.ModuleList([nn.GRUCell(input_size=self.num_nodes * self.input_dim,
hidden_size=self.hidden_state_size, hidden_size=self.hidden_state_size,
bias=True)] + [nn.GRUCell(input_size=self.hidden_state_size, bias=True)] + [
hidden_size=self.hidden_state_size, nn.GRUCell(input_size=self.hidden_state_size,
bias=True) for _ in range(self.num_rnn_layers - 1)] 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):
""" """
Encoder forward pass. Encoder forward pass.
:param inputs: shape (batch_size, self.num_nodes * self.input_dim) :param inputs: shape (batch_size, self.num_nodes * self.input_dim)
:param hidden_state: (num_layers, batch_size, self.hidden_state_size) -> optional, zeros if not provided :param hidden_state: (num_layers, batch_size, self.hidden_state_size)
optional, zeros if not provided
:return: output: # shape (batch_size, self.hidden_state_size) :return: output: # shape (batch_size, self.hidden_state_size)
hidden_state # shape (num_layers, batch_size, self.hidden_state_size) (lower indices mean lower layers) hidden_state # shape (num_layers, batch_size, self.hidden_state_size)
(lower indices mean lower layers)
""" """
batch_size, _ = inputs.size() batch_size, _ = inputs.size()
if hidden_state is None: if hidden_state is None:
@ -57,17 +52,18 @@ class EncoderModel(nn.Module):
hidden_states = [] hidden_states = []
output = inputs output = inputs
for layer_num, dcgru_layer in enumerate(self.dcgru_layers): for layer_num, dcgru_layer in enumerate(self.dcgru_layers):
hidden_state = dcgru_layer(output, hidden_state) next_hidden_state = dcgru_layer(output, hidden_state[layer_num])
hidden_states.append(hidden_state) hidden_states.append(next_hidden_state)
output = hidden_state output = next_hidden_state
return output, torch.cat(hidden_states, dim=1) # runs in O(num_layers) so not too slow # todo: check dim return output, torch.stack(hidden_states) # runs in O(num_layers) so not too slow
class DecoderModel(nn.Module): class DecoderModel(nn.Module, DCRNNModel):
def __init__(self, is_training, adj_mx, **model_kwargs): def __init__(self, is_training, adj_mx, **model_kwargs):
# super().__init__(is_training, adj_mx, **model_kwargs) # super().__init__(is_training, adj_mx, **model_kwargs)
super().__init__() nn.Module.__init__(self)
DCRNNModel.__init__(self, is_training, adj_mx, **model_kwargs)
self.adj_mx = adj_mx self.adj_mx = adj_mx
self.is_training = is_training self.is_training = is_training
self.max_diffusion_step = int(model_kwargs.get('max_diffusion_step', 2)) self.max_diffusion_step = int(model_kwargs.get('max_diffusion_step', 2))
@ -82,32 +78,30 @@ class DecoderModel(nn.Module):
self.use_curriculum_learning = bool(model_kwargs.get('use_curriculum_learning', False)) self.use_curriculum_learning = bool(model_kwargs.get('use_curriculum_learning', False))
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.hidden_state_size, self.num_nodes * self.output_dim)
self.dcgru_layers = [nn.GRUCell(input_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, hidden_size=self.hidden_state_size,
bias=True)] + [nn.GRUCell(input_size=self.hidden_state_size, bias=True)] + [
hidden_size=self.hidden_state_size, nn.GRUCell(input_size=self.hidden_state_size,
bias=True) for _ in hidden_size=self.hidden_state_size,
range(self.num_rnn_layers - 1)] bias=True) for _ in
range(self.num_rnn_layers - 1)])
def forward(self, inputs, hidden_state=None): def forward(self, inputs, hidden_state=None):
""" """
Decoder forward pass. Decoder forward pass.
:param inputs: shape (batch_size, self.num_nodes * self.output_dim) :param inputs: shape (batch_size, self.num_nodes * self.output_dim)
:param hidden_state: (num_layers, batch_size, self.hidden_state_size) -> optional, zeros if not provided :param hidden_state: (num_layers, batch_size, self.hidden_state_size)
optional, zeros if not provided
:return: output: # shape (batch_size, self.num_nodes * self.output_dim) :return: output: # shape (batch_size, self.num_nodes * self.output_dim)
hidden_state # shape (num_layers, batch_size, self.hidden_state_size) (lower indices mean lower layers) hidden_state # shape (num_layers, batch_size, self.hidden_state_size)
(lower indices mean lower layers)
""" """
batch_size, _ = inputs.size()
if hidden_state is None:
hidden_state = torch.zeros((self.num_rnn_layers, batch_size, self.hidden_state_size),
device=device)
hidden_states = [] hidden_states = []
output = inputs output = inputs
for layer_num, dcgru_layer in enumerate(self.dcgru_layers): for layer_num, dcgru_layer in enumerate(self.dcgru_layers):
hidden_state = dcgru_layer(output, hidden_state) next_hidden_state = dcgru_layer(output, hidden_state[layer_num])
hidden_states.append(hidden_state) hidden_states.append(next_hidden_state)
output = hidden_state output = next_hidden_state
return self.projection_layer(output), torch.cat(hidden_states, return self.projection_layer(output), torch.stack(hidden_states)
dim=1) # runs in O(num_layers) so not too slow #todo: check dim

27
model/pytorch/dcrnn_supervisor.py
View File

@ -7,6 +7,8 @@ import torch
from lib import utils from lib import utils
from model.pytorch.dcrnn_model import EncoderModel, DecoderModel from model.pytorch.dcrnn_model import EncoderModel, DecoderModel
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class DCRNNSupervisor: class DCRNNSupervisor:
def __init__(self, adj_mx, **kwargs): def __init__(self, adj_mx, **kwargs):
@ -87,7 +89,8 @@ class DCRNNSupervisor:
batch_size = inputs.size(1) batch_size = inputs.size(1)
inputs = inputs.view(self.seq_len, batch_size, self.num_nodes * self.input_dim) inputs = inputs.view(self.seq_len, batch_size, self.num_nodes * self.input_dim)
labels = labels.view(self.horizon, batch_size, self.num_nodes * self.output_dim) labels = labels[..., :self.output_dim].view(self.horizon, batch_size,
self.num_nodes * self.output_dim)
loss = 0 loss = 0
@ -95,6 +98,7 @@ class DCRNNSupervisor:
for t in range(self.seq_len): for t in range(self.seq_len):
_, encoder_hidden_state = self.encoder_model.forward(inputs[t], encoder_hidden_state) _, encoder_hidden_state = self.encoder_model.forward(inputs[t], encoder_hidden_state)
self._logger.info("Encoder complete, starting decoder")
go_symbol = torch.zeros((batch_size, self.num_nodes * self.output_dim)) go_symbol = torch.zeros((batch_size, self.num_nodes * self.output_dim))
decoder_hidden_state = encoder_hidden_state decoder_hidden_state = encoder_hidden_state
@ -113,6 +117,7 @@ class DCRNNSupervisor:
loss += criterion(self.standard_scaler.inverse_transform(decoder_output), loss += criterion(self.standard_scaler.inverse_transform(decoder_output),
self.standard_scaler.inverse_transform(labels[t])) self.standard_scaler.inverse_transform(labels[t]))
self._logger.info("Decoder complete, starting backprop")
loss.backward() loss.backward()
encoder_optimizer.step() encoder_optimizer.step()
decoder_optimizer.step() decoder_optimizer.step()
@ -135,16 +140,26 @@ class DCRNNSupervisor:
start_time = time.time() start_time = time.time()
for x, y in train_iterator: for _, (x, y) in enumerate(train_iterator):
loss = self._train_one_batch(x, y, batches_seen, encoder_optimizer, decoder_optimizer, criterion) x = torch.from_numpy(x).float()
y = torch.from_numpy(y).float()
self._logger.debug("X: {}".format(x.size()))
self._logger.debug("y: {}".format(y.size()))
x = x.permute(1, 0, 2, 3)
y = y.permute(1, 0, 2, 3)
loss = self._train_one_batch(x, y, batches_seen, encoder_optimizer,
decoder_optimizer, criterion)
losses.append(loss) losses.append(loss)
batches_seen += 1 batches_seen += 1
end_time = time.time() end_time = time.time()
if epoch_num % log_every == 0: if epoch_num % log_every == 0:
message = 'Epoch [{}/{}] ({}) train_mae: {:.4f}, val_mae: {:.4f} lr:{:.6f} {:.1f}s'.format( message = 'Epoch [{}/{}] ({}) train_mae: {:.4f}, val_mae: {:.4f} ' \
epoch_num, epochs, batches_seen, np.mean(losses), 0.0, 0.0, (end_time - start_time)) 'lr:{:.6f} {:.1f}s'.format(epoch_num, epochs, batches_seen,
np.mean(losses), 0.0,
0.0, (end_time - start_time))
self._logger.info(message) self._logger.info(message)
def _compute_sampling_threshold(self, batches_seen): def _compute_sampling_threshold(self, batches_seen):
return self.cl_decay_steps / (self.cl_decay_steps + np.exp(batches_seen / self.cl_decay_steps)) return self.cl_decay_steps / (
self.cl_decay_steps + np.exp(batches_seen / self.cl_decay_steps))