From 69288460b1da9f016a99989cfb74b9a42a3031b9 Mon Sep 17 00:00:00 2001 From: Chintan Shah Date: Tue, 1 Oct 2019 10:07:34 -0400 Subject: [PATCH] Simplified encoder decoder model and moved curriculum learning outside --- model/pytorch/dcrnn_model.py | 114 ++++++++++------------------------- 1 file changed, 33 insertions(+), 81 deletions(-) diff --git a/model/pytorch/dcrnn_model.py b/model/pytorch/dcrnn_model.py index 5612b1d..cf05659 100644 --- a/model/pytorch/dcrnn_model.py +++ b/model/pytorch/dcrnn_model.py @@ -1,12 +1,10 @@ -import numpy as np import torch import torch.nn as nn -from abc import ABC, abstractmethod device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -class DCRNNModel(metaclass=ABC): +class DCRNNModel: def __init__(self, is_training, scale_factor, adj_mx, **model_kwargs): super().__init__() self.adj_mx = adj_mx @@ -22,56 +20,6 @@ class DCRNNModel(metaclass=ABC): self.input_dim = int(model_kwargs.get('input_dim', 1)) self.hidden_state_size = self.num_nodes * self.rnn_units - @abstractmethod - @property - def dcgru_layers(self): - pass - - def _forward_cell(self, cell_input, prev_hidden_states): - """ - Runs for 1 time step through all layers. - :param cell_input: shape (batch_size, input feature size) - :param prev_hidden_states: (num_layers, batch_size, hidden size) - :return: output of cell: shape(batch_size, hidden size) - all hidden states from layer: shape(num_layers, batch_size, hidden size) - """ - hidden_states = [] - output = cell_input - for layer_num, dcgru_layer in enumerate(self.dcgru_layers): - hidden_state = dcgru_layer(output, prev_hidden_states[layer_num]) - hidden_states.append(hidden_state) - output = hidden_state - - return output, torch.cat(hidden_states, dim=1) # runs in O(num_layers) so not too slow #todo: check dim - - def _forward_impl(self, inputs, hidden_state): - """ - forward pass. - - :param inputs: shape (batch_size, timesteps, input_size) - :param hidden_state: (num_layers, batch_size, self.hidden_state_size) -> optional, zeros if not provided - :return: output: # shape (timesteps, batch_size, self.hidden_state_size) - hidden_state # shape (num_layers, batch_size, self.hidden_state_size) (lower indices mean lower layers) - """ - batch_size, timesteps, _ = inputs.size() - if hidden_state is None: - hidden_state = torch.zeros((self.num_rnn_layers, batch_size, self.hidden_state_size), - device=device) - output = torch.empty((timesteps, batch_size, self.hidden_state_size)) - for t in range(timesteps): - hidden_state = self.t_step_forward_pass(hidden_state, inputs, output, t) - - output = output.permute(1, 0, 2) - return output, hidden_state - - @abstractmethod - def t_step_forward_pass(self, hidden_state, inputs, output, t): - """ - Implements the forward pass for timestep t. - """ - # this is to accommodate curriculum learning - pass - class EncoderModel(nn.Module, DCRNNModel): def __init__(self, is_training, scaler, adj_mx, **model_kwargs): @@ -79,14 +27,9 @@ class EncoderModel(nn.Module, DCRNNModel): # https://pytorch.org/docs/stable/nn.html#gru self.seq_len = int(model_kwargs.get('seq_len')) # for the encoder - def t_step_forward_pass(self, hidden_state, inputs, output, t): - cell_input = inputs[:, t, :] # (batch_size, input_size) - cell_output, hidden_state = self._forward_cell(cell_input, hidden_state) - output[t] = cell_output - return hidden_state - + @property def dcgru_layers(self): - # input shape is supposed to be Input (batch_size, timesteps, num_sensor*input_dim) + # input shape is supposed to be Input (batch_size, num_sensor*input_dim) # first layer takes input shape and subsequent layer take input from the first layer return [nn.GRUCell(input_size=self.num_nodes * self.input_dim, hidden_size=self.hidden_state_size, @@ -99,12 +42,23 @@ class EncoderModel(nn.Module, DCRNNModel): """ Encoder forward pass. - :param inputs: shape (batch_size, timesteps, num_nodes * 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 - :return: output: # shape (timesteps, 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) """ - return self._forward_impl(inputs, hidden_state) + 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 = [] + output = inputs + for layer_num, dcgru_layer in enumerate(self.dcgru_layers): + hidden_state = dcgru_layer(output, hidden_state) + hidden_states.append(hidden_state) + output = hidden_state + + return output, torch.cat(hidden_states, dim=1) # runs in O(num_layers) so not too slow # todo: check dim class DecoderModel(nn.Module, DCRNNModel): @@ -115,6 +69,7 @@ class DecoderModel(nn.Module, DCRNNModel): 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) + @property def dcgru_layers(self): return [nn.GRUCell(input_size=self.num_nodes * self.output_dim, hidden_size=self.hidden_state_size, @@ -123,28 +78,25 @@ class DecoderModel(nn.Module, DCRNNModel): bias=True) for _ in range(self.num_rnn_layers - 1)] - def t_step_forward_pass(self, hidden_state, inputs, output, t): - cell_input = inputs[:, t, :] # (batch_size, input_size) - - if self.is_training: - if t > 0 and self.use_curriculum_learning: - c = np.random.uniform(0, 1) - if c >= self._compute_sampling_threshold(): #todo - cell_input = output[ - t - 1] # todo: this won't work because the linear layer is applied after forward_impl - - cell_output, hidden_state = self._forward_cell(cell_input, hidden_state) - output[t] = cell_output - return hidden_state - def forward(self, inputs, hidden_state=None): """ Decoder forward pass. - :param inputs: shape (batch_size, timesteps, num_nodes * input_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 - :return: output: # shape (timesteps, 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) """ - output, hidden = self._forward_impl(inputs, hidden_state) - return self.projection_layer(output), hidden_state + 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 = [] + output = inputs + for layer_num, dcgru_layer in enumerate(self.dcgru_layers): + hidden_state = dcgru_layer(output, hidden_state) + hidden_states.append(hidden_state) + output = hidden_state + + return self.projection_layer(output), torch.cat(hidden_states, + dim=1) # runs in O(num_layers) so not too slow #todo: check dim