From 66fb202d217d66b28aa547d0b6d9c42248012e9d Mon Sep 17 00:00:00 2001 From: Chintan Shah Date: Sun, 29 Sep 2019 12:51:49 -0400 Subject: [PATCH] Implemented Encoder with GRU - should swap GRU with DCGRU --- dcrnn_train.py | 2 +- model/pytorch/__init__.py | 0 model/pytorch/dcrnn_model.py | 41 ++++++++++++++++++++++++++++++ model/tf/__init__.py | 0 model/{ => tf}/dcrnn_cell.py | 0 model/{ => tf}/dcrnn_model.py | 3 +-- model/{ => tf}/dcrnn_supervisor.py | 2 +- requirements.txt | 1 + run_demo.py | 2 +- 9 files changed, 46 insertions(+), 5 deletions(-) create mode 100644 model/pytorch/__init__.py create mode 100644 model/pytorch/dcrnn_model.py create mode 100644 model/tf/__init__.py rename model/{ => tf}/dcrnn_cell.py (100%) rename model/{ => tf}/dcrnn_model.py (98%) rename model/{ => tf}/dcrnn_supervisor.py (99%) diff --git a/dcrnn_train.py b/dcrnn_train.py index de75465..0c28dfb 100644 --- a/dcrnn_train.py +++ b/dcrnn_train.py @@ -7,7 +7,7 @@ import tensorflow as tf import yaml from lib.utils import load_graph_data -from model.dcrnn_supervisor import DCRNNSupervisor +from model.tf.dcrnn_supervisor import DCRNNSupervisor def main(args): diff --git a/model/pytorch/__init__.py b/model/pytorch/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/model/pytorch/dcrnn_model.py b/model/pytorch/dcrnn_model.py new file mode 100644 index 0000000..0403a14 --- /dev/null +++ b/model/pytorch/dcrnn_model.py @@ -0,0 +1,41 @@ +import torch +import torch.nn as nn + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +class DCRNNModel(nn.Module): + def __init__(self, is_training, scale_factor, adj_mx, **model_kwargs): + super().__init__() + self.adj_mx = adj_mx + self.is_training = is_training + self.scale_factor = scale_factor + 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.horizon = int(model_kwargs.get('horizon', 1)) # for the decoder + # 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.seq_len = int(model_kwargs.get('seq_len')) # for the encoder + self.use_curriculum_learning = bool(model_kwargs.get('use_curriculum_learning', False)) + self.input_dim = int(model_kwargs.get('input_dim', 1)) + self.output_dim = int(model_kwargs.get('output_dim', 1)) + + +class EncoderModel(DCRNNModel): + def __init__(self, is_training, scaler, adj_mx, **model_kwargs): + super().__init__(is_training, scaler, adj_mx, **model_kwargs) + + # https://pytorch.org/docs/stable/nn.html#gru + + # since input shape is Input (batch_size, timesteps, num_sensor*input_dim),batch_first=True + self.dcgru = nn.GRU(input_size=self.num_nodes * self.input_dim, + hidden_size=self.rnn_units, + num_layers=self.num_rnn_layers, + batch_first=True) + + def forward(self, inputs, hidden_state=None): + # is None okay? + return self.dcgru(inputs, hidden_state) diff --git a/model/tf/__init__.py b/model/tf/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/model/dcrnn_cell.py b/model/tf/dcrnn_cell.py similarity index 100% rename from model/dcrnn_cell.py rename to model/tf/dcrnn_cell.py diff --git a/model/dcrnn_model.py b/model/tf/dcrnn_model.py similarity index 98% rename from model/dcrnn_model.py rename to model/tf/dcrnn_model.py index 43797fa..16fa379 100644 --- a/model/dcrnn_model.py +++ b/model/tf/dcrnn_model.py @@ -6,8 +6,7 @@ import tensorflow as tf from tensorflow.contrib import legacy_seq2seq -from lib.metrics import masked_mae_loss -from model.dcrnn_cell import DCGRUCell +from model.tf.dcrnn_cell import DCGRUCell class DCRNNModel(object): diff --git a/model/dcrnn_supervisor.py b/model/tf/dcrnn_supervisor.py similarity index 99% rename from model/dcrnn_supervisor.py rename to model/tf/dcrnn_supervisor.py index 18a399e..8024628 100644 --- a/model/dcrnn_supervisor.py +++ b/model/tf/dcrnn_supervisor.py @@ -13,7 +13,7 @@ from lib import utils, metrics from lib.AMSGrad import AMSGrad from lib.metrics import masked_mae_loss -from model.dcrnn_model import DCRNNModel +from model.tf.dcrnn_model import DCRNNModel class DCRNNSupervisor(object): diff --git a/requirements.txt b/requirements.txt index f577b89..eb89a38 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +torch scipy>=0.19.0 numpy>=1.12.1 pandas>=0.19.2 diff --git a/run_demo.py b/run_demo.py index ecbbe86..05b617f 100644 --- a/run_demo.py +++ b/run_demo.py @@ -6,7 +6,7 @@ import tensorflow as tf import yaml from lib.utils import load_graph_data -from model.dcrnn_supervisor import DCRNNSupervisor +from model.tf.dcrnn_supervisor import DCRNNSupervisor def run_dcrnn(args):