Cleanup and Model Improvements

.gitignore

@@ -0,0 +1,8 @@
+utest_experiment/*
+heatmap.png
+last_test_metrics.csv
+preds.csv
+.ipynb_checkpoints/
+data/
+.DS_Store
+*.pkl

README.md

@@ -1,7 +1,9 @@
 # Graph WaveNet for Deep Spatial-Temporal Graph Modeling
 
 This is the original pytorch implementation of Graph WaveNet in the following paper: 
-[Graph WaveNet for Deep Spatial-Temporal Graph Modeling, IJCAI 2019] (https://arxiv.org/abs/1906.00121).
+[Graph WaveNet for Deep Spatial-Temporal Graph Modeling, IJCAI 2019] (https://arxiv.org/abs/1906.00121),
+with modifications presented in [Incrementally Improving Graph WaveNet Performance on Traffic Prediction] (https://arxiv.org/abs/1912.07390):
+
 
 <p align="center">
   <img width="350" height="400" src=./fig/model.png>
@@ -14,9 +16,9 @@ This is the original pytorch implementation of Graph WaveNet in the following pa
 
 ## Data Preparation
 
-### Step1: Download METR-LA and PEMS-BAY data from [Google Drive](https://drive.google.com/open?id=10FOTa6HXPqX8Pf5WRoRwcFnW9BrNZEIX) or [Baidu Yun](https://pan.baidu.com/s/14Yy9isAIZYdU__OYEQGa_g) links provided by [DCRNN](https://github.com/liyaguang/DCRNN).
+1) Download METR-LA and PEMS-BAY data from [Google Drive](https://drive.google.com/open?id=10FOTa6HXPqX8Pf5WRoRwcFnW9BrNZEIX) or [Baidu Yun](https://pan.baidu.com/s/14Yy9isAIZYdU__OYEQGa_g) links provided by [DCRNN](https://github.com/liyaguang/DCRNN).
 
-### Step2:
+2)
 
 ```
 # Create data directories
@@ -29,40 +31,42 @@ python generate_training_data.py --output_dir=data/METR-LA --traffic_df_filename
 python generate_training_data.py --output_dir=data/PEMS-BAY --traffic_df_filename=data/pems-bay.h5
 
 ```
-## Experiments
-Train models configured in Table 3 of the paper.
 
+## Train Commands
+Note: train.py saves metrics to a directory specified by the `--save` arg in metrics.csv and test_metrics.csv
+
+Model that gets (3.00 - 3.02 Test MAE, ~2.73 Validation MAE)
+```
+python train.py --cat_feat_gc --fill_zeroes --do_graph_conv --addaptadj  --randomadj --es_patience 20 --save logs/baseline_v2
+```
+
+Finetuning (2.99 - 3.00 MAE)
+```
+python generate_training_data.py --seq_length_y 6 --output_dir data/METR-LA_12_6
+python train.py --data  data/METR-LA_12_6 --cat_feat_gc --fill_zeroes --do_graph_conv --addaptadj  --randomadj --es_patience 20 --save logs/front_6
+python train.py --checkpoint  logs/front_6/best_model.pth --cat_feat_gc --fill_zeroes --do_graph_conv --addaptadj  --randomadj --es_patience 20 --save logs/finetuned
+
+```
+Original Graph Wavenet Model (3.04-3.07 MAE)
+```
+python train.py --clip 5 --lr_decay_rate 1. --nhid 32 --do_graph_conv --addaptadj  --randomadj --save logs/baseline
 ```
-ep=100
-dv=cuda:0
-mkdir experiment
-mkdir experiment/metr
-
-#identity
-expid=1
-python train.py --device $dv --gcn_bool --adjtype identity  --epoch $ep --expid $expid  --save ./experiment/metr/metr > ./experiment/metr/train-$expid.log
-rm ./experiment/metr/metr_epoch*
-
-#forward-only
-expid=2
-python train.py --device $dv --gcn_bool --adjtype transition --epoch $ep --expid $expid  --save ./experiment/metr/metr > ./experiment/metr/train-$expid.log
-rm ./experiment/metr/metr_epoch*
-
-#adaptive-only
-expid=3
-python train.py --device $dv --gcn_bool --adjtype transition --aptonly  --addaptadj --randomadj --epoch $ep --expid $expid  --save ./experiment/metr/metr > ./experiment/metr/train-$expid.log
-rm ./experiment/metr/metr_epoch*
-
-#forward-backward
-expid=4
-python train.py --device $dv --gcn_bool --adjtype doubletransition  --epoch $ep --expid $expid  --save ./experiment/metr/metr > ./experiment/metr/train-$expid.log
-rm ./experiment/metr/metr_epoch*
-
-#forward-backward-adaptive
-expid=5
-python train.py --device $dv --gcn_bool --adjtype doubletransition --addaptadj  --randomadj  --epoch $ep --expid $expid  --save ./experiment/metr/metr > ./experiment/metr/train-$expid.log
-rm ./experiment/metr/metr_epoch*
 
+You can also train from a jupyter notebook with
+```{python}
+from train import main
+from durbango import pickle_load
+args = pickle_load('baseline_args.pkl') # manipulate these in python
+args.lr_decay_rate = .97
+args.clip = 3
+args.save = 'logs/from_jupyter'
+main(args) # takes roughly an hour depending on nhid, and early_stopping
 ```
 
+Train models configured in Table 3 of the original GraphWavenet paper by using the `--adjtype, --addaptadj, --aptonly` command line argument.
+These flags are (somewhat) documented in util.py.
+
+Run unitests with `pytest`
 
+### Possible Improvements
+* move redundant `.transpose(1,3)` to dataloader or `load_dataset`

engine.py

@@ -1,43 +1,42 @@
 import torch.optim as optim
 from model import *
 import util
-class trainer():
-    def __init__(self, scaler, in_dim, seq_length, num_nodes, nhid , dropout, lrate, wdecay, device, supports, gcn_bool, addaptadj, aptinit):
-        self.model = gwnet(device, num_nodes, dropout, supports=supports, gcn_bool=gcn_bool, addaptadj=addaptadj, aptinit=aptinit, in_dim=in_dim, out_dim=seq_length, residual_channels=nhid, dilation_channels=nhid, skip_channels=nhid * 8, end_channels=nhid * 16)
-        self.model.to(device)
+
+class Trainer():
+    def __init__(self, model: GWNet, scaler, lrate, wdecay, clip=3, lr_decay_rate=.97):
+        self.model = model
+
         self.optimizer = optim.Adam(self.model.parameters(), lr=lrate, weight_decay=wdecay)
-        self.loss = util.masked_mae
         self.scaler = scaler
-        self.clip = 5
+        self.clip = clip
+        self.scheduler = optim.lr_scheduler.LambdaLR(
+            self.optimizer, lr_lambda=lambda epoch: lr_decay_rate ** epoch)
+
+    @classmethod
+    def from_args(cls, model, scaler, args):
+        return cls(model, scaler, args.learning_rate, args.weight_decay, clip=args.clip,
+                   lr_decay_rate=args.lr_decay_rate)
 
     def train(self, input, real_val):
         self.model.train()
         self.optimizer.zero_grad()
         input = nn.functional.pad(input,(1,0,0,0))
-        output = self.model(input)
-        output = output.transpose(1,3)
-        #output = [batch_size,12,num_nodes,1]
-        real = torch.unsqueeze(real_val,dim=1)
+        output = self.model(input).transpose(1,3)  # now, output = [batch_size,1,num_nodes, seq_length]
         predict = self.scaler.inverse_transform(output)
-
-        loss = self.loss(predict, real, 0.0)
-        loss.backward()
+        assert predict.shape[1] == 1
+        mae, mape, rmse = util.calc_metrics(predict.squeeze(1), real_val, null_val=0.0)
+        mae.backward()
         if self.clip is not None:
             torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip)
         self.optimizer.step()
-        mape = util.masked_mape(predict,real,0.0).item()
-        rmse = util.masked_rmse(predict,real,0.0).item()
-        return loss.item(),mape,rmse
+        return mae.item(),mape.item(),rmse.item()
 
     def eval(self, input, real_val):
         self.model.eval()
         input = nn.functional.pad(input,(1,0,0,0))
-        output = self.model(input)
-        output = output.transpose(1,3)
-        #output = [batch_size,12,num_nodes,1]
+        output = self.model(input).transpose(1,3) #  [batch_size,seq_length,num_nodes,1]
         real = torch.unsqueeze(real_val,dim=1)
         predict = self.scaler.inverse_transform(output)
-        loss = self.loss(predict, real, 0.0)
-        mape = util.masked_mape(predict,real,0.0).item()
-        rmse = util.masked_rmse(predict,real,0.0).item()
-        return loss.item(),mape,rmse
+        predict = torch.clamp(predict, min=0., max=70.)
+        mae, mape, rmse = [x.item() for x in util.calc_metrics(predict, real, null_val=0.0)]
+        return mae, mape, rmse

exp_results.py

@@ -0,0 +1,38 @@
+"""Utilities for comparing metrics saved by train.py"""
+import pandas as pd
+import os
+from glob import glob
+import matplotlib.pyplot as plt
+
+
+def summary(d):
+    try:
+        tr_val = pd.read_csv(f'{d}/metrics.csv', index_col=0)
+        tr_ser = tr_val.loc[tr_val.valid_loss.idxmin()]
+        tr_ser['best_epoch'] = tr_val.valid_loss.idxmin()
+        tr_ser['min_train_loss'] = tr_val.train_loss.min()
+    except FileNotFoundError:
+        tr_ser = pd.Series()
+    try:
+        tmet = pd.read_csv(f'{d}/test_metrics.csv', index_col=0)
+        tmean = tmet.add_prefix('test_').mean()
+
+    except FileNotFoundError:
+        tmean = pd.Series()
+    tab = pd.concat([tr_ser, tmean]).round(3)
+    return tab
+
+def loss_curve(d):
+    if 'logs' not in d: d =  f'logs/{d}'
+    tr_val = pd.read_csv(f'{d}/metrics.csv', index_col=0)
+    return tr_val[['train_loss', 'valid_loss']]
+
+
+def plot_loss_curve(log_dir):
+    d = loss_curve(log_dir)
+    ax = d.plot()
+    plt.axhline(d.valid_loss.min())
+    print(d.valid_loss.idxmin())
+
+def make_results_table():
+    return pd.DataFrame({os.path.basename(c): summary(c) for c in glob('logs/*')}).T.sort_values('valid_loss')

gen_adj_mx.py

@@ -0,0 +1,63 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import numpy as np
+import pandas as pd
+import pickle
+
+
+def get_adjacency_matrix(distance_df, sensor_ids, normalized_k=0.1):
+    """
+
+    :param distance_df: data frame with three columns: [from, to, distance].
+    :param sensor_ids: list of sensor ids.
+    :param normalized_k: entries that become lower than normalized_k after normalization are set to zero for sparsity.
+    :return:
+    """
+    num_sensors = len(sensor_ids)
+    dist_mx = np.zeros((num_sensors, num_sensors), dtype=np.float32)
+    dist_mx[:] = np.inf
+    # Builds sensor id to index map.
+    sensor_id_to_ind = {}
+    for i, sensor_id in enumerate(sensor_ids):
+        sensor_id_to_ind[sensor_id] = i
+
+    # Fills cells in the matrix with distances.
+    for row in distance_df.values:
+        if row[0] not in sensor_id_to_ind or row[1] not in sensor_id_to_ind:
+            continue
+        dist_mx[sensor_id_to_ind[row[0]], sensor_id_to_ind[row[1]]] = row[2]
+
+    # Calculates the standard deviation as theta.
+    distances = dist_mx[~np.isinf(dist_mx)].flatten()
+    std = distances.std()
+    adj_mx = np.exp(-np.square(dist_mx / std))
+    # Make the adjacent matrix symmetric by taking the max.
+    # adj_mx = np.maximum.reduce([adj_mx, adj_mx.T])
+
+    # Sets entries that lower than a threshold, i.e., k, to zero for sparsity.
+    adj_mx[adj_mx < normalized_k] = 0
+    return sensor_ids, sensor_id_to_ind, adj_mx
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--sensor_ids_filename', type=str, default='data/sensor_graph/graph_sensor_ids.txt',
+                        help='File containing sensor ids separated by comma.')
+    parser.add_argument('--distances_filename', type=str, default='data/sensor_graph/distances_la_2012.csv',
+                        help='CSV file containing sensor distances with three columns: [from, to, distance].')
+    parser.add_argument('--normalized_k', type=float, default=0.1,
+                        help='Entries that become lower than normalized_k after normalization are set to zero for sparsity.')
+    parser.add_argument('--output_pkl_filename', type=str, default='data/sensor_graph/adj_mat.pkl',
+                        help='Path of the output file.')
+    args = parser.parse_args()
+
+    with open(args.sensor_ids_filename) as f:
+        sensor_ids = f.read().strip().split(',')
+    distance_df = pd.read_csv(args.distances_filename, dtype={'from': 'str', 'to': 'str'})
+    _, sensor_id_to_ind, adj_mx = get_adjacency_matrix(distance_df, sensor_ids, args.normalized_k)
+    # Save to pickle file.
+    with open(args.output_pkl_filename, 'wb') as f:
+        pickle.dump([sensor_ids, sensor_id_to_ind, adj_mx], f, protocol=2)