├── lib ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-36.pyc │ └── __init__.cpython-37.pyc ├── metrics.py └── utils_HIAM.py ├── models ├── __init__.py ├── __pycache__ │ ├── Net.cpython-36.pyc │ ├── rgcn.cpython-36.pyc │ ├── DO_Net.cpython-36.pyc │ ├── OD_Net.cpython-36.pyc │ ├── GGRUCell.cpython-36.pyc │ └── __init__.cpython-36.pyc ├── DualInfoTransformer.py ├── rgcn.py ├── GGRUCell.py ├── DO_Net.py ├── OD_Net.py └── Net.py ├── SOTA ├── DCRNN │ ├── lib │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── utils.cpython-36.pyc │ │ │ ├── __init__.cpython-36.pyc │ │ │ └── metrics.cpython-36.pyc │ │ ├── metrics.py │ │ └── utils.py │ ├── model │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ └── dcrnn_supervisor.cpython-36.pyc │ │ ├── dcrnn_cell.py │ │ ├── dcrnn_model.py │ │ └── dcrnn_supervisor.py │ ├── script.sh │ ├── data │ │ └── config │ │ │ ├── dcrnn_do_sh_96.yaml │ │ │ ├── dcrnn_od_hz_96.yaml │ │ │ ├── dcrnn_od_sh_96.yaml │ │ │ └── dcrnn_do_hz_96.yaml │ └── dcrnn_train_pytorch.py └── Graph-WaveNet-master │ ├── metrics.py │ ├── scripts.sh │ ├── engine1.py │ ├── util1.py │ ├── model1.py │ └── train1.py ├── data ├── prediction_formula.png ├── config │ ├── train_hz_dim26_units96_h4c512.yaml │ ├── train_sh_dim76_units96_h4c512.yaml │ ├── eval_hz_dim26_units96_h4c512.yaml │ └── eval_sh_dim76_units96_h4c512.yaml └── README.md ├── .idea ├── misc.xml ├── vcs.xml ├── inspectionProfiles │ ├── profiles_settings.xml │ └── Project_Default.xml ├── .gitignore ├── modules.xml └── HIAM.iml ├── README.md ├── evaluate.py └── train.py /lib/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /SOTA/DCRNN/lib/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /SOTA/DCRNN/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data/prediction_formula.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HCPLab-SYSU/HIAM/HEAD/data/prediction_formula.png -------------------------------------------------------------------------------- /models/__pycache__/Net.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HCPLab-SYSU/HIAM/HEAD/models/__pycache__/Net.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/rgcn.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HCPLab-SYSU/HIAM/HEAD/models/__pycache__/rgcn.cpython-36.pyc -------------------------------------------------------------------------------- /lib/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HCPLab-SYSU/HIAM/HEAD/lib/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /lib/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HCPLab-SYSU/HIAM/HEAD/lib/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/DO_Net.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HCPLab-SYSU/HIAM/HEAD/models/__pycache__/DO_Net.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/OD_Net.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HCPLab-SYSU/HIAM/HEAD/models/__pycache__/OD_Net.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/GGRUCell.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HCPLab-SYSU/HIAM/HEAD/models/__pycache__/GGRUCell.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HCPLab-SYSU/HIAM/HEAD/models/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /SOTA/DCRNN/lib/__pycache__/utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HCPLab-SYSU/HIAM/HEAD/SOTA/DCRNN/lib/__pycache__/utils.cpython-36.pyc -------------------------------------------------------------------------------- /SOTA/DCRNN/lib/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HCPLab-SYSU/HIAM/HEAD/SOTA/DCRNN/lib/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /SOTA/DCRNN/lib/__pycache__/metrics.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HCPLab-SYSU/HIAM/HEAD/SOTA/DCRNN/lib/__pycache__/metrics.cpython-36.pyc -------------------------------------------------------------------------------- /SOTA/DCRNN/model/__pycache__/dcrnn_supervisor.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HCPLab-SYSU/HIAM/HEAD/SOTA/DCRNN/model/__pycache__/dcrnn_supervisor.cpython-36.pyc -------------------------------------------------------------------------------- /.idea/misc.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /.idea/vcs.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /.idea/inspectionProfiles/profiles_settings.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 6 | -------------------------------------------------------------------------------- /.idea/.gitignore: -------------------------------------------------------------------------------- 1 | # Default ignored files 2 | /shelf/ 3 | /workspace.xml 4 | # Datasource local storage ignored files 5 | /../../../../:\PycharmProjects\HIAM\.idea/dataSources/ 6 | /dataSources.local.xml 7 | # Editor-based HTTP Client requests 8 | /httpRequests/ 9 | -------------------------------------------------------------------------------- /.idea/modules.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /.idea/HIAM.iml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 12 | -------------------------------------------------------------------------------- /SOTA/DCRNN/script.sh: -------------------------------------------------------------------------------- 1 | #! /usr/bin/bash 2 | 3 | 4 | ### hangzhou 5 | ## od 6 | CUDA_VISIBLE_DEVICES=1 OMP_NUM_THREADS=1 python dcrnn_train_pytorch.py \ 7 | --config 8 | ## do 9 | CUDA_VISIBLE_DEVICES=5 OMP_NUM_THREADS=1 python dcrnn_train_pytorch.py \ 10 | --config data/config/dcrnn_do_hz_96.yaml 11 | 12 | ### shanghai 13 | ## od 14 | CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 python dcrnn_train_pytorch.py \ 15 | --config data/model/dcrnn_od_sh.yaml 16 | ## do 17 | CUDA_VISIBLE_DEVICES=5 OMP_NUM_THREADS=1 python dcrnn_train_pytorch.py \ 18 | --config data/config/dcrnn_do_sh_96.yaml -------------------------------------------------------------------------------- /lib/metrics.py: -------------------------------------------------------------------------------- 1 | # part of this code are copied from DCRNN 2 | import numpy as np 3 | 4 | def masked_rmse_np(preds, labels): 5 | return np.sqrt(masked_mse_np(preds=preds, labels=labels)) 6 | 7 | 8 | def masked_mse_np(preds, labels): 9 | rmse = np.square(np.subtract(preds, labels)).astype('float32') 10 | rmse = np.nan_to_num(rmse) 11 | return np.mean(rmse) 12 | 13 | 14 | def masked_mae_np(preds, labels): 15 | mae = np.abs(np.subtract(preds, labels)).astype('float32') 16 | mae = np.nan_to_num(mae) 17 | return np.mean(mae) 18 | 19 | 20 | def masked_mape_np(preds, labels): 21 | mape_net = np.divide(np.sum(np.abs(np.subtract(preds, labels)).astype('float32')), np.sum(labels)) 22 | mape_net = np.nan_to_num(mape_net) 23 | return mape_net 24 | -------------------------------------------------------------------------------- /SOTA/DCRNN/lib/metrics.py: -------------------------------------------------------------------------------- 1 | # part of this code are copied from DCRNN 2 | import numpy as np 3 | 4 | def masked_rmse_np(preds, labels): 5 | return np.sqrt(masked_mse_np(preds=preds, labels=labels)) 6 | 7 | 8 | def masked_mse_np(preds, labels): 9 | rmse = np.square(np.subtract(preds, labels)).astype('float32') 10 | rmse = np.nan_to_num(rmse) 11 | return np.mean(rmse) 12 | 13 | 14 | def masked_mae_np(preds, labels): 15 | mae = np.abs(np.subtract(preds, labels)).astype('float32') 16 | mae = np.nan_to_num(mae) 17 | return np.mean(mae) 18 | 19 | 20 | def masked_mape_np(preds, labels): 21 | mape_net = np.divide(np.sum(np.abs(np.subtract(preds, labels)).astype('float32')), np.sum(labels)) 22 | mape_net = np.nan_to_num(mape_net) 23 | return mape_net -------------------------------------------------------------------------------- /SOTA/Graph-WaveNet-master/metrics.py: -------------------------------------------------------------------------------- 1 | # part of this code are copied from DCRNN 2 | import numpy as np 3 | 4 | def masked_rmse_np(preds, labels): 5 | return np.sqrt(masked_mse_np(preds=preds, labels=labels)) 6 | 7 | 8 | def masked_mse_np(preds, labels): 9 | rmse = np.square(np.subtract(preds, labels)).astype('float32') 10 | rmse = np.nan_to_num(rmse) 11 | return np.mean(rmse) 12 | 13 | 14 | def masked_mae_np(preds, labels): 15 | mae = np.abs(np.subtract(preds, labels)).astype('float32') 16 | mae = np.nan_to_num(mae) 17 | return np.mean(mae) 18 | 19 | 20 | def masked_mape_np(preds, labels): 21 | mape_net = np.divide(np.sum(np.abs(np.subtract(preds, labels)).astype('float32')), np.sum(labels)) 22 | mape_net = np.nan_to_num(mape_net) 23 | return mape_net -------------------------------------------------------------------------------- /SOTA/DCRNN/data/config/dcrnn_do_sh_96.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | base_dir: data/checkpoint/sh 3 | log_level: INFO 4 | data: 5 | batch_size: 8 6 | dataset_dir: data/shanghai/DO/DO_76 7 | test_batch_size: 8 8 | val_batch_size: 8 9 | graph_pkl_filename: data/shanghai/graph_sh_conn.pkl 10 | ds_type: do 11 | 12 | model: 13 | cl_decay_steps: 200 14 | filter_type: dual_random_walk 15 | horizon: 4 16 | input_dim: 76 17 | l1_decay: 0 18 | max_diffusion_step: 2 19 | num_nodes: 288 20 | num_rnn_layers: 2 21 | output_dim: 76 22 | rnn_units: 96 23 | seq_len: 4 24 | use_curriculum_learning: true 25 | 26 | train: 27 | base_lr: 0.001 28 | dropout: 0 29 | epoch: 0 30 | epochs: 300 31 | epsilon: 1.0e-4 32 | global_step: 0 33 | lr_decay_ratio: 0.2 34 | max_grad_norm: 5 35 | max_to_keep: 100 36 | min_learning_rate: 1.0e-12 37 | optimizer: adam 38 | patience: 100 39 | steps: [60, 80, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280] 40 | test_every_n_epochs: 1 41 | -------------------------------------------------------------------------------- /SOTA/DCRNN/data/config/dcrnn_od_hz_96.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | base_dir: data/checkpoint/hz 3 | log_level: INFO 4 | data: 5 | batch_size: 32 6 | dataset_dir: data/hangzhou/OD/OD_26 7 | test_batch_size: 32 8 | val_batch_size: 32 9 | graph_pkl_filename: data/hangzhou/graph_hz_conn.pkl 10 | ds_type: od 11 | 12 | model: 13 | cl_decay_steps: 200 14 | filter_type: dual_random_walk 15 | horizon: 4 16 | input_dim: 26 17 | l1_decay: 0 18 | max_diffusion_step: 2 19 | num_nodes: 80 20 | num_rnn_layers: 2 21 | output_dim: 26 22 | rnn_units: 96 23 | seq_len: 4 24 | use_curriculum_learning: true 25 | 26 | train: 27 | base_lr: 0.001 28 | dropout: 0 29 | epoch: 0 30 | epochs: 300 31 | epsilon: 1.0e-4 32 | global_step: 0 33 | lr_decay_ratio: 0.5 34 | max_grad_norm: 5 35 | max_to_keep: 100 36 | min_learning_rate: 1.0e-12 37 | optimizer: adam 38 | patience: 100 39 | steps: [60, 80, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280] 40 | test_every_n_epochs: 1 41 | -------------------------------------------------------------------------------- /SOTA/DCRNN/data/config/dcrnn_od_sh_96.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | base_dir: data/checkpoint/sh 3 | log_level: INFO 4 | data: 5 | batch_size: 8 6 | dataset_dir: data/shanghai/OD/OD_76 7 | test_batch_size: 8 8 | val_batch_size: 8 9 | graph_pkl_filename: data/shanghai/graph_sh_conn.pkl 10 | ds_type: od 11 | 12 | model: 13 | cl_decay_steps: 200 14 | filter_type: dual_random_walk 15 | horizon: 4 16 | input_dim: 76 17 | l1_decay: 0 18 | max_diffusion_step: 2 19 | num_nodes: 288 20 | num_rnn_layers: 2 21 | output_dim: 76 22 | rnn_units: 96 23 | seq_len: 4 24 | use_curriculum_learning: true 25 | 26 | train: 27 | base_lr: 0.001 28 | dropout: 0 29 | epoch: 0 30 | epochs: 300 31 | epsilon: 1.0e-4 32 | global_step: 0 33 | lr_decay_ratio: 0.2 34 | max_grad_norm: 5 35 | max_to_keep: 100 36 | min_learning_rate: 1.0e-12 37 | optimizer: adam 38 | patience: 100 39 | steps: [60, 80, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280] 40 | test_every_n_epochs: 1 41 | -------------------------------------------------------------------------------- /SOTA/DCRNN/data/config/dcrnn_do_hz_96.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | base_dir: data/checkpoint/hz_do 3 | log_level: INFO 4 | data: 5 | batch_size: 32 6 | dataset_dir: data/hangzhou/DO/DO_26 7 | test_batch_size: 32 8 | val_batch_size: 32 9 | graph_pkl_filename: data/hangzhou/graph_hz_conn.pkl 10 | ds_type: do 11 | 12 | model: 13 | cl_decay_steps: 200 14 | filter_type: dual_random_walk 15 | horizon: 4 16 | input_dim: 26 17 | l1_decay: 0 18 | max_diffusion_step: 2 19 | num_nodes: 80 20 | num_rnn_layers: 2 21 | output_dim: 26 22 | rnn_units: 96 23 | seq_len: 4 24 | use_curriculum_learning: true 25 | 26 | train: 27 | base_lr: 0.001 28 | dropout: 0 29 | epoch: 0 30 | epochs: 300 31 | epsilon: 1.0e-4 32 | global_step: 0 33 | lr_decay_ratio: 0.5 34 | max_grad_norm: 5 35 | max_to_keep: 100 36 | min_learning_rate: 1.0e-12 37 | optimizer: adam 38 | patience: 100 39 | steps: [60, 80, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280] 40 | test_every_n_epochs: 1 41 | -------------------------------------------------------------------------------- /.idea/inspectionProfiles/Project_Default.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 24 | -------------------------------------------------------------------------------- /data/config/train_hz_dim26_units96_h4c512.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | base_dir: data/checkpoint/hz 3 | log_level: INFO 4 | data: 5 | batch_size: 32 6 | test_batch_size: 32 7 | dataset_dir: data/hangzhou/OD/ 8 | do_dataset_dir: data/hangzhou/DO/ 9 | graph_pkl_filename: [ data/hangzhou/graph_hz_conn.pkl] 10 | 11 | model: 12 | cl_decay_steps: 200 13 | horizon: 4 14 | input_dim: 26 15 | output_dim: 26 16 | output_type: fc 17 | num_nodes: 80 18 | num_rnn_layers: 2 19 | rnn_units: 96 20 | seq_len: 4 21 | head: 4 22 | channel: 512 23 | l1_decay: 0 24 | use_curriculum_learning: true 25 | dropout_type: none 26 | dropout_prop: 0.05 27 | use_input: true 28 | num_relations: 1 29 | num_bases: 1 30 | K: 2 31 | norm: True 32 | global_fusion: false 33 | 34 | 35 | train: 36 | base_lr: 0.001 37 | epoch: 0 38 | epochs: 300 39 | epsilon: 1.0e-4 40 | global_step: 0 41 | lr_decay_ratio: 0.5 42 | max_grad_norm: 5 43 | min_learning_rate: 1.0e-12 44 | optimizer: adam 45 | patience: 100 46 | steps: [60, 80, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280] 47 | test_every_n_epochs: 1 48 | save_every_n_epochs: 1 49 | -------------------------------------------------------------------------------- /data/config/train_sh_dim76_units96_h4c512.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | base_dir: data/checkpoint/sh 3 | log_level: INFO 4 | data: 5 | batch_size: 8 6 | test_batch_size: 8 7 | dataset_dir: data/shanghai/OD/ 8 | do_dataset_dir: data/shanghai/DO/ 9 | graph_pkl_filename: [ data/shanghai/graph_sh_conn.pkl] 10 | 11 | model: 12 | cl_decay_steps: 200 13 | horizon: 4 14 | input_dim: 76 15 | output_dim: 76 16 | output_type: fc 17 | num_nodes: 288 18 | num_rnn_layers: 2 19 | rnn_units: 96 20 | seq_len: 4 21 | head: 4 22 | channel: 512 23 | l1_decay: 0 24 | use_curriculum_learning: true 25 | dropout_type: none 26 | dropout_prop: 0.05 27 | use_input: true 28 | num_relations: 1 29 | num_bases: 1 30 | K: 2 31 | norm: True 32 | global_fusion: false 33 | 34 | 35 | train: 36 | base_lr: 0.001 37 | epoch: 0 38 | epochs: 300 39 | epsilon: 1.0e-4 40 | global_step: 0 41 | lr_decay_ratio: 0.2 42 | max_grad_norm: 5 43 | min_learning_rate: 1.0e-12 44 | optimizer: adam 45 | patience: 100 46 | steps: [60, 80, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280] 47 | test_every_n_epochs: 1 48 | save_every_n_epochs: 1 49 | -------------------------------------------------------------------------------- /data/config/eval_hz_dim26_units96_h4c512.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | base_dir: data/checkpoint/hz 3 | log_level: INFO 4 | data: 5 | batch_size: 32 6 | test_batch_size: 32 7 | dataset_dir: data/hangzhou/OD/ 8 | do_dataset_dir: data/hangzhou/DO/ 9 | graph_pkl_filename: [ data/hangzhou/graph_hz_conn.pkl] 10 | 11 | model: 12 | cl_decay_steps: 200 13 | horizon: 4 14 | input_dim: 26 15 | output_dim: 26 16 | output_type: fc 17 | num_nodes: 80 18 | num_rnn_layers: 2 19 | rnn_units: 96 20 | seq_len: 4 21 | head: 4 22 | channel: 512 23 | l1_decay: 0 24 | use_curriculum_learning: true 25 | dropout_type: none 26 | dropout_prop: 0.05 27 | use_input: true 28 | num_relations: 1 29 | num_bases: 1 30 | K: 2 31 | norm: True 32 | global_fusion: false 33 | save_path: data/checkpoint/hz_dim26_units96_h4c512.pt 34 | 35 | train: 36 | base_lr: 0.001 37 | epoch: 0 38 | epochs: 300 39 | epsilon: 1.0e-4 40 | global_step: 0 41 | lr_decay_ratio: 0.5 42 | max_grad_norm: 5 43 | min_learning_rate: 1.0e-12 44 | optimizer: adam 45 | patience: 100 46 | steps: [60, 80, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280] 47 | test_every_n_epochs: 1 48 | save_every_n_epochs: 1 49 | -------------------------------------------------------------------------------- /data/config/eval_sh_dim76_units96_h4c512.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | base_dir: data/checkpoint/sh 3 | log_level: INFO 4 | data: 5 | batch_size: 8 6 | test_batch_size: 8 7 | dataset_dir: data/shanghai/OD/ 8 | do_dataset_dir: data/shanghai/DO/ 9 | graph_pkl_filename: [ data/shanghai/graph_sh_conn.pkl] 10 | 11 | model: 12 | cl_decay_steps: 200 13 | horizon: 4 14 | input_dim: 76 15 | output_dim: 76 16 | output_type: fc 17 | num_nodes: 288 18 | num_rnn_layers: 2 19 | rnn_units: 96 20 | seq_len: 4 21 | head: 4 22 | channel: 512 23 | l1_decay: 0 24 | use_curriculum_learning: true 25 | dropout_type: none 26 | dropout_prop: 0.05 27 | use_input: true 28 | num_relations: 1 29 | num_bases: 1 30 | K: 2 31 | norm: True 32 | global_fusion: false 33 | save_path: data/checkpoint/sh_dim76_units96_h4c512.pt 34 | 35 | 36 | train: 37 | base_lr: 0.001 38 | epoch: 0 39 | epochs: 300 40 | epsilon: 1.0e-4 41 | global_step: 0 42 | lr_decay_ratio: 0.2 43 | max_grad_norm: 5 44 | min_learning_rate: 1.0e-12 45 | optimizer: adam 46 | patience: 100 47 | steps: [60, 80, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280] 48 | test_every_n_epochs: 1 49 | save_every_n_epochs: 1 50 | -------------------------------------------------------------------------------- /SOTA/DCRNN/dcrnn_train_pytorch.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import argparse 6 | import yaml 7 | import pickle 8 | import numpy as np 9 | 10 | from lib.utils import load_graph_data 11 | from model.dcrnn_supervisor import DCRNNSupervisor 12 | 13 | def load_pickle(pickle_file): 14 | try: 15 | with open(pickle_file, 'rb') as f: 16 | pickle_data = pickle.load(f) 17 | except UnicodeDecodeError as e: 18 | with open(pickle_file, 'rb') as f: 19 | pickle_data = pickle.load(f, encoding='latin1') 20 | except Exception as e: 21 | print('Unable to load data ', pickle_file, ':', e) 22 | raise 23 | return pickle_data 24 | 25 | def main(args): 26 | with open(args.config_filename) as f: 27 | supervisor_config = yaml.load(f) 28 | 29 | graph_pkl_filename = supervisor_config['data'].get('graph_pkl_filename') 30 | # sensor_ids, sensor_id_to_ind, adj_mx = load_graph_data(graph_pkl_filename) 31 | adj_mx = load_pickle(graph_pkl_filename).astype(np.float32) 32 | 33 | for i in range(len(adj_mx)): # 构建邻接矩阵 34 | adj_mx[i, i] = 0 35 | 36 | supervisor = DCRNNSupervisor(adj_mx=adj_mx, **supervisor_config) 37 | 38 | supervisor.train() 39 | 40 | 41 | if __name__ == '__main__': 42 | parser = argparse.ArgumentParser() 43 | parser.add_argument('--config_filename', default=None, type=str, 44 | help='Configuration filename for restoring the model.') 45 | parser.add_argument('--use_cpu_only', default=False, type=bool, help='Set to true to only use cpu.') 46 | args = parser.parse_args() 47 | main(args) 48 | -------------------------------------------------------------------------------- /SOTA/Graph-WaveNet-master/scripts.sh: -------------------------------------------------------------------------------- 1 | #! /usr/bin/bash 2 | 3 | ### hangzhou OD 4 | CUDA_VISIBLE_DEVICES=7 OMP_NUM_THREADS=1 python train1.py \ 5 | --data data/hangzhou/OD/OD_26 \ 6 | --adjdata data/hangzhou/graph_hz_conn.pkl \ 7 | --exp_base "data" \ 8 | --runs "hz_od" \ 9 | --dropout 0.001 \ 10 | --epochs 400 \ 11 | --learning_rate 0.008 \ 12 | --nhid 96 \ 13 | --out_dim 26 \ 14 | --in_dim 26 \ 15 | --seq_length 4 \ 16 | --num_nodes 80 \ 17 | --gcn_bool \ 18 | --addaptadj \ 19 | --train_tyep "od" \ 20 | --randomadj 21 | 22 | ### hangzhou DO 23 | CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 python train1.py \ 24 | --data data/hangzhou/DO/DO_26 \ 25 | --adjdata data/hangzhou/graph_hz_conn.pkl \ 26 | --exp_base "data" \ 27 | --runs "hz_do" \ 28 | --dropout 0.001 \ 29 | --epochs 400 \ 30 | --learning_rate 0.008 \ 31 | --nhid 96 \ 32 | --out_dim 26 \ 33 | --in_dim 26 \ 34 | --seq_length 4 \ 35 | --num_nodes 80 \ 36 | --gcn_bool \ 37 | --addaptadj \ 38 | --train_type "do" \ 39 | --randomadj 40 | 41 | 42 | ###shanghai OD 43 | CUDA_VISIBLE_DEVICES=5 OMP_NUM_THREADS=1 python train1.py \ 44 | --data data/shanghai/OD/OD_76 \ 45 | --adjdata data/shanghai/graph_sh_conn.pkl \ 46 | --exp_base "data" \ 47 | --runs "sh" \ 48 | --dropout 0.001 \ 49 | --epochs 300 \ 50 | --learning_rate 0.01 \ 51 | --nhid 96 \ 52 | --in_dim 76 \ 53 | --out_dim 76 \ 54 | --num_nodes 288 \ 55 | --seq_length 4 \ 56 | --gcn_bool \ 57 | --addaptadj \ 58 | --randomadj 59 | 60 | ###shanghai DO 61 | CUDA_VISIBLE_DEVICES=1 OMP_NUM_THREADS=1 python train1.py \ 62 | --data data/shanghai/DO/DO_76 \ 63 | --adjdata data/shanghai/graph_sh_conn.pkl \ 64 | --exp_base "data" \ 65 | --runs "sh_do" \ 66 | --dropout 0.001 \ 67 | --epochs 300 \ 68 | --learning_rate 0.01 \ 69 | --nhid 96 \ 70 | --in_dim 76 \ 71 | --out_dim 76 \ 72 | --num_nodes 288 \ 73 | --seq_length 4 \ 74 | --gcn_bool \ 75 | --addaptadj \ 76 | --train_type "do" \ 77 | --randomadj 78 | 79 | tar --exclude='./data' -zcvf /backup/filename.tgz . -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Online Metro Origin-Destination Prediction via Heterogeneous Information Aggregation 2 | This is a PyTorch implementation of **Online Metro Origin-Destination Prediction via Heterogeneous Information Aggregation**. 3 | 4 | With high trip efficiencies and cheap ticket charges, metro has recently become a popular travel mode for urban residents. Due to its significant applications, metro ridership prediction has recently attracted extensive attention in both academic and industrial communities. However, most of conventional works were merely proposed for station-level prediction, i.e., forecasting the inflow and outflow of each metro station. Such information of inflow/outflow ridership is too coarse to reflect the mobility of passengers. 5 | 6 | We propose a novel **Heterogeneous Information Aggregation Machine** to facilitate the online metro origin-destination prediction. To the best of our knowledge, our **HIAM** is the first deep learning based approach that fully aggregates heterogeneous information to jointly forecast the future OD ridership and DO ridership. 7 | 8 | If you use this code for your research, please cite our papers. 9 | ``` 10 | @article{liu2022online, 11 | title={Online Metro Origin-Destination Prediction via Heterogeneous Information Aggregation}, 12 | author={Liu, Lingbo and Zhu, Yuying and Li, Guanbin and Wu, Ziyi and Bai, Lei and Lin, Liang}, 13 | journal={IEEE Transactions on Pattern Analysis and Machine Intelligence}, 14 | year={2022}, 15 | publisher={IEEE} 16 | } 17 | ``` 18 | 19 | ### Requirements 20 | - python3 21 | - numpy 22 | - yaml 23 | - pytorch 24 | - torch_geometric 25 | ### Extract dataset 26 | Please download the dataset and extract it to `data` folder. 27 | - [Dropbox link](https://www.dropbox.com/sh/4pgk4uez7g200fg/AACHN6wMhjq_v0R2ZZ8ZeI6ma?dl=0) 28 | - [Baidu Netdisk，password：q6e0 ](https://pan.baidu.com/s/1PHN8SNT3jTroX0sTWHsrXw) 29 | 30 | ## Train 31 | - SHMOD 32 | ``` 33 | python train.py --config data/config/train_sh_dim76_units96_h4c512.yaml 34 | ``` 35 | 36 | - HZMOD 37 | ``` 38 | python train.py --config data/config/train_hz_dim26_units96_h4c512.yaml 39 | ``` 40 | ## Test 41 | First of all, download the trained model and extract it to the path:`data/checkpoint`. 42 | 43 | - SHMOD 44 | ``` 45 | python evaluate.py --config data/checkpoint/eval_sh_dim76_units96_h4c512.yaml 46 | ``` 47 | - HZMOD 48 | ``` 49 | python evaluate.py --config data/checkpoint/eval_hz_dim26_units96_h4c512.yaml 50 | ``` 51 | 52 | -------------------------------------------------------------------------------- /data/README.md: -------------------------------------------------------------------------------- 1 | # SHMOD & HZMOD Dataset 2 | In this work, we focus on the metro ridership of 5:15 - 23:30 and utilize the metro OD/DO distribution of the previous four time intervals (15minutes x 4 = 60minutes) to predict the metro OD/DO distribution of future four time intervals (15minutes x 4 = 60minutes). 3 | 4 | Specifically, the data of input sequence consists of incomplete OD matrix(IOD), unfinished order vector(U), and DO matrix(DO) and our online origin-destination prediction could be denoted as follows. 5 | 6 | ​ ![image](https://github.com/GillianZhu/HIAM/blob/master/data/prediction_formula.png) 7 | 8 | ​ where i=1,2,...,n and j=1,2,...,m. n=4, m=4 9 | 10 | We release two datasets named SHMOD and HZMOD, respectively. Each dataset is divided into a training set, a validation set and a testing set, and also contains corresponding metro graph information. For each set, we release four `pkl` files, three for metro OD ridership distribution data, and one for metro DO ridership distribution data. 11 | 12 | ### 1. Metro OD&DO Matrix 13 | 14 | - train.pkl: the OD training set. It is a `dict` that consists of 5 `ndarray`: 15 | 16 | ``` 17 | (1) finished: the incomplete OD matrix of the previous four time intervals. Its shape is [T, n, N, D]. 18 | (2) unfinished: the unfinished order vector of the previous four time intervals. Its shape is [T, n, N, 1]. 19 | (3) y: the complete OD matrix of the next four time intervals. Its shape is [T, n, N, D]. 20 | (4) xtime: the timestamps of x. Its shape is [T, n]. 21 | (5) ytime: the timestamps of y. Its shape is [T, m]. 22 | 23 | T = the number of time slices 24 | N = the number of metro stations, i.e, 80/288 (HZMOD/SHMOD) in our work 25 | n = the length of the input sequence, i.e, 4 time intervals in our work 26 | m = the length of the output sequence, i.e, 4 time intervals in our work 27 | D = the data dimension of each station, i.e, 26/76 (HZMOD/SHMOD) in our work 28 | ``` 29 | 30 | - train_history_long.pkl: the long-term historical information of the training set. It is a `dict` that consists of 1 `ndarray`: 31 | 32 | ``` 33 | (1) history: the long-term destination distribution matrix of the previous four time intervals. Its shape is [T, n, N, D]. 34 | ``` 35 | 36 | Note that val_history_long.pkl and test_history_long.pkl are also built based on the long-term historical data on the training set. 37 | 38 | - train_history_short.pkl: the short-term historical information of the training set. It is a `dict` that consists of 1 `ndarray`: 39 | 40 | ``` 41 | (1) history: the short-term destination distribution matrix of the previous four time intervals. Its shape is [T, n, N, D]. 42 | ``` 43 | 44 | - train_do.pkl: the DO training set. It is a `dict` that consists of 4 `ndarray`: 45 | 46 | ``` 47 | (1) do_x: the DO matrix of the previous four time intervals. Its shape is [T, n, N, D]. 48 | (2) do_y: the DO matrix of the next four time intervals. Its shape is [T, m, N, D]. 49 | (3) xtime: the timestamps of x. Its shape is [T, n]. 50 | (4) ytime: the timestamps of y. Its shape is [T, m]. 51 | ``` 52 | 53 | Besides, each pkl file has its counterparts on testing set and validation set with similar structure. 54 | 55 | ### 2. Graph Information 56 | 57 | - graph_conn.pkl: the physical topology information of metro system. 58 | -------------------------------------------------------------------------------- /SOTA/Graph-WaveNet-master/engine1.py: -------------------------------------------------------------------------------- 1 | import torch.optim as optim 2 | from model1 import * 3 | from torch.optim.lr_scheduler import MultiStepLR 4 | import torch.nn as nn 5 | import torch 6 | 7 | class StepLR2(MultiStepLR): 8 | """StepLR with min_lr""" 9 | 10 | def __init__(self, 11 | optimizer, 12 | milestones, 13 | gamma=0.1, 14 | last_epoch=-1, 15 | min_lr=2.0e-6): 16 | """ 17 | 18 | :optimizer: TODO 19 | :milestones: TODO 20 | :gamma: TODO 21 | :last_epoch: TODO 22 | :min_lr: TODO 23 | 24 | """ 25 | self.optimizer = optimizer 26 | self.milestones = milestones 27 | self.gamma = gamma 28 | self.last_epoch = last_epoch 29 | self.min_lr = min_lr 30 | super(StepLR2, self).__init__(optimizer, milestones, gamma) 31 | 32 | def get_lr(self): 33 | lr_candidate = super(StepLR2, self).get_lr() 34 | if isinstance(lr_candidate, list): 35 | for i in range(len(lr_candidate)): 36 | lr_candidate[i] = max(self.min_lr, lr_candidate[i]) 37 | 38 | else: 39 | lr_candidate = max(self.min_lr, lr_candidate) 40 | 41 | return lr_candidate 42 | 43 | class trainer(): 44 | def __init__(self, scaler, in_dim, seq_length, num_nodes, out_dim, nhid , dropout, lrate, wdecay, device, supports, gcn_bool, addaptadj, aptinit): 45 | self.model = gwnet( 46 | device, 47 | num_nodes, 48 | dropout, 49 | supports=supports, 50 | gcn_bool=gcn_bool, 51 | addaptadj=addaptadj, 52 | aptinit=aptinit, 53 | in_dim=in_dim, # input feature 54 | out_dim=out_dim, # output feature dim 55 | residual_channels=nhid, 56 | dilation_channels=nhid, 57 | skip_channels=nhid * 8, # TODO: skip channels ? 58 | end_channels=nhid * 16, 59 | blocks=4, 60 | ) 61 | self.model.to(device) 62 | self.optimizer = optim.Adam(self.model.parameters(), lr=lrate, weight_decay=wdecay) 63 | self.scheduler = MultiStepLR(optimizer=self.optimizer, milestones=[200, 350], gamma=0.5) 64 | self.loss = nn.L1Loss(reduction='mean') 65 | self.scaler = scaler 66 | self.clip = 5 67 | self.device =device 68 | 69 | def train(self, input, real_val): 70 | self.model.train() 71 | self.optimizer.zero_grad() 72 | input = input.to(self.device) 73 | real_val = real_val.to(self.device) 74 | output = self.model(input) 75 | real = real_val 76 | predict = self.scaler.inverse_transform(output) 77 | loss = self.loss(predict, real) 78 | loss.backward() 79 | if self.clip is not None: 80 | torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip) 81 | self.optimizer.step() 82 | self.scheduler.step() 83 | 84 | return loss.item(), predict, real 85 | 86 | def eval(self, input, real_val): 87 | self.model.eval() 88 | input = nn.functional.pad(input,(1,0,0,0)).to(self.device) 89 | output = self.model(input) 90 | output = output.transpose(1,3) 91 | real = real_val 92 | predict = self.scaler.inverse_transform(output) 93 | loss = self.loss(predict, real,self.device, 0.0) 94 | return loss.item(), predict, real -------------------------------------------------------------------------------- /models/DualInfoTransformer.py: -------------------------------------------------------------------------------- 1 | from torch_geometric import nn as gnn 2 | from torch import nn 3 | import torch 4 | import math 5 | import sys 6 | import os 7 | import copy 8 | from torch.nn import functional as F 9 | sys.path.insert(0, os.path.abspath('..')) 10 | 11 | 12 | class DualInfoTransformer(nn.Module): 13 | def __init__(self, h=4, d_nodes=288, d_channel=512, d_model=96): 14 | "Take in model size and number of heads." 15 | super(DualInfoTransformer, self).__init__() 16 | assert d_model % h == 0 17 | 18 | self.d_nodes = d_nodes 19 | self.d_model = d_model 20 | self.d_channel = d_channel 21 | self.d_k = d_channel // h 22 | self.h = h 23 | 24 | self.od_linears = self.clones( 25 | nn.Sequential( 26 | nn.Conv1d(in_channels=d_model, out_channels=d_channel, kernel_size=1), 27 | nn.PReLU(d_channel), 28 | nn.Conv1d(in_channels=d_channel, out_channels=d_channel, kernel_size=1), 29 | nn.PReLU(d_channel)), 3) 30 | self.do_linears = self.clones( 31 | nn.Sequential( 32 | nn.Conv1d(in_channels=d_model, out_channels=d_channel, kernel_size=1), 33 | nn.PReLU(d_channel), 34 | nn.Conv1d(in_channels=d_channel, out_channels=d_channel, kernel_size=1), 35 | nn.PReLU(d_channel)), 3) 36 | self.od_conv = nn.Sequential( 37 | nn.Conv1d(in_channels=d_channel, out_channels=d_channel, kernel_size=1), 38 | nn.PReLU(d_channel), 39 | nn.Conv1d(in_channels=d_channel, out_channels=d_model, kernel_size=1), 40 | nn.PReLU(d_model) 41 | ) 42 | self.do_conv = nn.Sequential( 43 | nn.Conv1d(in_channels=d_channel, out_channels=d_channel, kernel_size=1), 44 | nn.PReLU(d_channel), 45 | nn.Conv1d(in_channels=d_channel, out_channels=d_model, kernel_size=1), 46 | nn.PReLU(d_model) 47 | ) 48 | 49 | def clones(self, module, N): 50 | "Produce N identical layers." 51 | return nn.ModuleList([copy.deepcopy(module) for _ in range(N)]) 52 | 53 | def attention(self, query, key, value): 54 | "Compute 'Scaled Dot Product Attention'" 55 | d_k = query.size(-1) # batch, heads, num_nodes, num_units / heads 56 | scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k) # batch, heads, num_nodes, num_nodes 57 | p_attn = F.softmax(scores, dim=-1) # batch, heads, num_nodes, num_nodes 58 | return torch.matmul(p_attn, value) # batch, heads, num_nodes, num_units / heads 59 | 60 | def MultiHeadedAttention(self, hid_od, hid_do): 61 | hid_od = hid_od.view(-1, self.d_model, self.d_nodes) 62 | hid_do = hid_do.view(-1, self.d_model, self.d_nodes) 63 | nbatches = hid_od.size(0) 64 | 65 | # 1) Do all the linear projections in batch from d_model => h x d_k 66 | odquery, odkey, odvalue = \ 67 | [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2) 68 | for l, x in zip(self.od_linears, (hid_od, hid_od, hid_od))] # batch, heads, num_nodes, num_units / heads 69 | doquery, dokey, dovalue = \ 70 | [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2) 71 | for l, x in zip(self.do_linears, (hid_do, hid_do, hid_do))] 72 | 73 | # 2) Apply attention on all the projected vectors in batch. 74 | attn_od = self.od_conv( 75 | self.attention(query=odquery, 76 | key=dokey, 77 | value=dovalue).transpose(-2, -1).contiguous().view(-1, self.d_channel, self.d_nodes)) # batch, heads, 1, num_units / heads 78 | attn_do = self.do_conv( 79 | self.attention(query=doquery, 80 | key=odkey, 81 | value=odvalue).transpose(-2, -1).contiguous().view(-1, self.d_channel, self.d_nodes)) 82 | return attn_od.view(-1, self.d_model), attn_do.view(-1, self.d_model) 83 | 84 | def forward(self, hidden_states_od, hidden_states_do): 85 | return self.MultiHeadedAttention(hidden_states_od, hidden_states_do) 86 | 87 | 88 | 89 | -------------------------------------------------------------------------------- /models/rgcn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.nn import Parameter as Param 3 | from torch_geometric.nn.conv import MessagePassing 4 | 5 | from torch_geometric.nn.inits import uniform 6 | from torch.nn.init import xavier_normal_, calculate_gain 7 | 8 | 9 | class RGCNConv(MessagePassing): 10 | r"""The relational graph convolutional operator from the `"Modeling 11 | Relational Data with Graph Convolutional Networks" 12 | `_ paper 13 | 14 | .. math:: 15 | \mathbf{x}^{\prime}_i = \mathbf{\Theta}_0 \cdot \mathbf{x}_i + 16 | \sum_{r \in \mathcal{R}} \sum_{j \in \mathcal{N}_r(i)} 17 | \frac{1}{|\mathcal{N}_r(i)|} \mathbf{\Theta}_r \cdot \mathbf{x}_j, 18 | 19 | where :math:`\mathcal{R}` denotes the set of relations, *i.e.* edge types. 20 | Edge type needs to be a one-dimensional :obj:`torch.long` tensor which 21 | stores a relation identifier 22 | :math:`\in \{ 0, \ldots, |\mathcal{R}| - 1\}` for each edge. 23 | 24 | Args: 25 | in_channels (int): Size of each input sample. 26 | out_channels (int): Size of each output sample. 27 | num_relations (int): Number of relations. 28 | num_bases (int): Number of bases used for basis-decomposition. 29 | bias (bool, optional): If set to :obj:`False`, the layer will not learn 30 | an additive bias. (default: :obj:`True`) 31 | **kwargs (optional): Additional arguments of 32 | :class:`torch_geometric.nn.conv.MessagePassing`. 33 | """ 34 | 35 | def __init__(self, 36 | in_channels, 37 | out_channels, 38 | num_relations, 39 | num_bases, 40 | bias=True, 41 | **kwargs): 42 | super(RGCNConv, self).__init__(aggr='add', **kwargs) 43 | 44 | self.in_channels = in_channels 45 | self.out_channels = out_channels 46 | self.num_relations = num_relations 47 | self.num_bases = num_bases 48 | 49 | self.basis = Param(torch.Tensor(num_bases, in_channels, out_channels)) 50 | self.att = Param(torch.Tensor(num_relations, num_bases)) 51 | self.root = Param(torch.Tensor(in_channels, out_channels)) 52 | 53 | if bias: 54 | self.bias = Param(torch.Tensor(out_channels)) 55 | else: 56 | self.register_parameter('bias', None) 57 | 58 | self.reset_parameters() 59 | 60 | def reset_parameters(self): 61 | # size = self.num_bases * self.in_channels 62 | # uniform(size, self.basis) 63 | # uniform(size, self.att) 64 | # uniform(size, self.root) 65 | # xavier_normal_(self.basis, gain=calculate_gain('relu')) 66 | # xavier_normal_(self.root, gain=calculate_gain('relu')) 67 | torch.nn.init.xavier_uniform_(self.basis) 68 | torch.nn.init.xavier_uniform_(self.att) 69 | torch.nn.init.xavier_uniform_(self.root) 70 | 71 | def forward(self, x, edge_index, edge_attr, edge_norm=None): 72 | """""" 73 | # print("forward:", x.shape, edge_attr.shape) 74 | return self.propagate( 75 | edge_index, x=x, edge_attr=edge_attr, edge_norm=edge_norm) 76 | 77 | def message(self, x_j, edge_index_j, edge_attr, edge_norm): 78 | w = torch.matmul(self.att, self.basis.view(self.num_bases, -1)) 79 | w = w.to(x_j.device) # 加cuda 80 | w = w.view(self.num_relations, self.in_channels, self.out_channels) # (num_relations, in, out) 81 | out = torch.einsum('bi,rio->bro', x_j, w) # x_j: (batch, in) w:(num_relations, in, out) out:(batch, num_relations, out) 先转置w再相乘 82 | # print("message1:", x_j.size(), w.size(), out.size()) 83 | out = (out * edge_attr.unsqueeze(2)).sum(dim=1) 84 | # print("message2:", out.size(), edge_attr.unsqueeze(2).size()) 85 | # out = torch.bmm(x_j.unsqueeze(1), w).squeeze(-2) 86 | 87 | return out if edge_norm is None else out * edge_norm.view(-1, 1) 88 | 89 | def update(self, aggr_out, x): 90 | if x is None: 91 | out = aggr_out + self.root 92 | else: 93 | out = aggr_out + torch.matmul(x, self.root) 94 | 95 | if self.bias is not None: 96 | out = out + self.bias 97 | 98 | # print("update:", out.size()) 99 | return out 100 | 101 | def __repr__(self): 102 | return '{}({}, {}, num_relations={})'.format( 103 | self.__class__.__name__, self.in_channels, self.out_channels, 104 | self.num_relations) 105 | -------------------------------------------------------------------------------- /models/GGRUCell.py: -------------------------------------------------------------------------------- 1 | from torch_geometric import nn as gnn 2 | from torch import nn 3 | from torch.nn import functional as F 4 | from torch.nn import init, Parameter 5 | import torch 6 | import random 7 | import math 8 | import sys 9 | import os 10 | import copy 11 | # from lib.utils_unfinished import softmax 12 | from torch.nn import functional as F 13 | sys.path.insert(0, os.path.abspath('..')) 14 | 15 | from models.rgcn import RGCNConv 16 | 17 | 18 | def zoneout(prev_h, next_h, rate, training=True): 19 | """TODO: Docstring for zoneout. 20 | 21 | :prev_h: TODO 22 | :next_h: TODO 23 | 24 | :p: when p = 1, all new elements should be droped 25 | when p = 0, all new elements should be maintained 26 | 27 | :returns: TODO 28 | 29 | """ 30 | from torch.nn.functional import dropout 31 | if training: 32 | # bernoulli: draw a value 1. 33 | # p = 1 -> d = 1 -> return prev_h 34 | # p = 0 -> d = 0 -> return next_h 35 | # d = torch.zeros_like(next_h).bernoulli_(p) 36 | # return (1 - d) * next_h + d * prev_h 37 | next_h = (1 - rate) * dropout(next_h - prev_h, rate) + prev_h 38 | else: 39 | next_h = rate * prev_h + (1 - rate) * next_h 40 | 41 | return next_h 42 | 43 | 44 | class KStepRGCN(nn.Module): 45 | """docstring for KStepRGCN""" 46 | 47 | def __init__( 48 | self, 49 | in_channels, 50 | out_channels, 51 | num_relations, 52 | num_bases, 53 | K, 54 | bias, 55 | ): 56 | super(KStepRGCN, self).__init__() 57 | self.in_channels = in_channels 58 | self.out_channels = out_channels 59 | self.num_relations = num_relations 60 | self.num_bases = num_bases 61 | self.K = K 62 | self.rgcn_layers = nn.ModuleList([ 63 | RGCNConv(in_channels, 64 | out_channels, 65 | num_relations, 66 | num_bases, 67 | bias) 68 | ] + [ 69 | RGCNConv(out_channels, 70 | out_channels, 71 | num_relations, 72 | num_bases, 73 | bias) for _ in range(self.K - 1) 74 | ]) 75 | self.mediate_activation = nn.PReLU() 76 | 77 | def forward(self, x, edge_index, edge_attr): 78 | for i in range(self.K): 79 | x = self.rgcn_layers[i](x=x, 80 | edge_index=edge_index, 81 | edge_attr=edge_attr, 82 | edge_norm=None) 83 | # not final layer, add relu 84 | if i != self.K - 1: 85 | #x = torch.relu(x) 86 | x = self.mediate_activation(x) 87 | return x 88 | 89 | 90 | class GGRUCell(nn.Module): 91 | """Docstring for GGRUCell. """ 92 | 93 | def __init__(self, 94 | in_channels, 95 | out_channels, 96 | dropout_type=None, 97 | dropout_prob=0.0, 98 | num_relations=3, 99 | num_bases=3, 100 | K=1, 101 | num_nodes=80, 102 | global_fusion=False): 103 | """TODO: to be defined1. """ 104 | super(GGRUCell, self).__init__() 105 | self.num_chunks = 3 106 | self.in_channels = in_channels 107 | self.out_channels = out_channels 108 | self.num_relations = num_relations 109 | self.num_bases = num_bases 110 | self.num_nodes = num_nodes 111 | self.global_fusion = global_fusion 112 | self.cheb_i = KStepRGCN(in_channels, 113 | out_channels * self.num_chunks, 114 | num_relations=num_relations, 115 | num_bases=num_bases, 116 | K=K, 117 | bias=False) 118 | self.cheb_h = KStepRGCN(out_channels, 119 | out_channels * self.num_chunks, 120 | num_relations=num_relations, 121 | num_bases=num_bases, 122 | K=K, 123 | bias=False) 124 | 125 | self.bias_i = Parameter(torch.Tensor(self.out_channels)) 126 | self.bias_r = Parameter(torch.Tensor(self.out_channels)) 127 | self.bias_n = Parameter(torch.Tensor(self.out_channels)) 128 | self.dropout_prob = dropout_prob 129 | self.dropout_type = dropout_type 130 | 131 | self.reset_parameters() 132 | 133 | def reset_parameters(self): 134 | init.ones_(self.bias_i) 135 | init.ones_(self.bias_r) 136 | init.ones_(self.bias_n) 137 | 138 | if self.global_fusion is True: 139 | init.ones_(self.bias_i_g) 140 | init.ones_(self.bias_r_g) 141 | init.ones_(self.bias_n_g) 142 | 143 | def forward(self, inputs, edge_index, edge_attr, hidden=None): 144 | """TODO: Docstring for forward. 145 | 146 | :inputs: TODO 147 | :hidden: TODO 148 | :returns: TODO 149 | 150 | """ 151 | if hidden is None: 152 | hidden = torch.zeros(inputs.size(0), 153 | self.out_channels, 154 | dtype=inputs.dtype, 155 | device=inputs.device) 156 | gi = self.cheb_i(inputs, edge_index=edge_index, edge_attr=edge_attr) 157 | gh = self.cheb_h(hidden, edge_index=edge_index, edge_attr=edge_attr) 158 | # print("cheb_i:", gi.shape) 159 | # print("cheb_h:", gh.shape) 160 | i_r, i_i, i_n = gi.chunk(3, 1) 161 | h_r, h_i, h_n = gh.chunk(3, 1) 162 | 163 | resetgate = torch.sigmoid(i_r + h_r + self.bias_r) 164 | inputgate = torch.sigmoid(i_i + h_i + self.bias_i) 165 | newgate = torch.tanh(i_n + resetgate * h_n + self.bias_n) 166 | next_hidden = (1 - inputgate) * newgate + inputgate * hidden 167 | 168 | output = next_hidden 169 | 170 | if self.dropout_type == 'zoneout': 171 | next_hidden = zoneout(prev_h=hidden, 172 | next_h=next_hidden, 173 | rate=self.dropout_prob, 174 | training=self.training) 175 | 176 | elif self.dropout_type == 'dropout': 177 | next_hidden = F.dropout(next_hidden, 178 | self.dropout_prob, 179 | self.training) 180 | 181 | return output, next_hidden -------------------------------------------------------------------------------- /models/DO_Net.py: -------------------------------------------------------------------------------- 1 | from torch_geometric import nn as gnn 2 | from torch import nn 3 | import torch 4 | import sys 5 | import os 6 | sys.path.insert(0, os.path.abspath('..')) 7 | 8 | from models.GGRUCell import GGRUCell 9 | 10 | class DONet(torch.nn.Module): 11 | 12 | def __init__(self, cfg, logger): 13 | super(DONet, self).__init__() 14 | self.logger = logger 15 | self.cfg = cfg 16 | 17 | self.num_nodes = cfg['model']['num_nodes'] 18 | self.num_output_dim = cfg['model']['output_dim'] 19 | self.num_units = cfg['model']['rnn_units'] 20 | self.num_finished_input_dim = cfg['model']['input_dim'] 21 | self.num_rnn_layers = cfg['model']['num_rnn_layers'] 22 | self.seq_len = cfg['model']['seq_len'] 23 | self.horizon = cfg['model']['horizon'] 24 | self.num_relations = cfg['model'].get('num_relations', 1) 25 | self.K = cfg['model'].get('K', 2) 26 | self.num_bases = cfg['model'].get('num_bases', 1) 27 | 28 | self.use_curriculum_learning = self.cfg['model'][ 29 | 'use_curriculum_learning'] 30 | self.cl_decay_steps = torch.FloatTensor( 31 | data=[self.cfg['model']['cl_decay_steps']]) 32 | self.dropout_type = cfg['model'].get('dropout_type', None) 33 | self.dropout_prob = cfg['model'].get('dropout_prob', 0.0) 34 | self.use_input = cfg['model'].get('use_input', True) 35 | 36 | self.global_fusion = cfg['model'].get('global_fusion', False) 37 | 38 | self.encoder_first_cells = GGRUCell(self.num_finished_input_dim, 39 | self.num_units, 40 | self.dropout_type, 41 | self.dropout_prob, 42 | self.num_relations, 43 | num_bases=self.num_bases, 44 | K=self.K, 45 | num_nodes=self.num_nodes, 46 | global_fusion=self.global_fusion) 47 | self.encoder_second_cells = nn.ModuleList([GGRUCell(self.num_units, 48 | self.num_units, 49 | self.dropout_type, 50 | self.dropout_prob, 51 | self.num_relations, 52 | num_bases=self.num_bases, 53 | K=self.K, 54 | num_nodes=self.num_nodes, 55 | global_fusion=self.global_fusion) 56 | for _ in range(self.num_rnn_layers - 1)]) 57 | 58 | self.decoder_first_cells = GGRUCell(self.num_finished_input_dim, 59 | self.num_units, 60 | self.dropout_type, 61 | self.dropout_prob, 62 | self.num_relations, 63 | num_bases=self.num_bases, 64 | K=self.K, 65 | num_nodes=self.num_nodes, 66 | global_fusion=self.global_fusion) 67 | 68 | self.decoder_second_cells = nn.ModuleList([GGRUCell(self.num_units, 69 | self.num_units, 70 | self.dropout_type, 71 | self.dropout_prob, 72 | self.num_relations, 73 | self.K, 74 | num_nodes=self.num_nodes, 75 | global_fusion=self.global_fusion) 76 | for _ in range(self.num_rnn_layers - 1)]) 77 | self.output_type = cfg['model'].get('output_type', 'fc') 78 | if self.output_type == 'fc': 79 | self.output_layer = nn.Linear(self.num_units, self.num_output_dim) 80 | 81 | 82 | def encoder_first_layer(self, 83 | batch, 84 | enc_first_hidden, 85 | edge_index, 86 | edge_attr=None): 87 | enc_first_out, enc_first_hidden = self.encoder_first_cells(inputs=batch.x_do, 88 | edge_index=edge_index, 89 | edge_attr=edge_attr, 90 | hidden=enc_first_hidden) 91 | return enc_first_out, enc_first_hidden 92 | 93 | def encoder_second_layer(self, 94 | index, 95 | encoder_first_out, 96 | enc_second_hidden, 97 | edge_index, 98 | edge_attr): 99 | enc_second_out, enc_second_hidden = self.encoder_second_cells[index](inputs=encoder_first_out, 100 | hidden=enc_second_hidden, 101 | edge_index=edge_index, 102 | edge_attr=edge_attr) 103 | return enc_second_out, enc_second_hidden 104 | 105 | def decoder_first_layer(self, 106 | decoder_input, 107 | dec_first_hidden, 108 | edge_index, 109 | edge_attr=None): 110 | dec_first_out, dec_first_hidden = self.decoder_first_cells(inputs=decoder_input, 111 | hidden=dec_first_hidden, 112 | edge_index=edge_index, 113 | edge_attr=edge_attr) 114 | 115 | return dec_first_out, dec_first_hidden 116 | 117 | def decoder_second_layer(self, 118 | index, 119 | decoder_first_out, 120 | dec_second_hidden, 121 | edge_index, 122 | edge_attr=None): 123 | dec_second_out, dec_second_hidden = self.decoder_second_cells[index](inputs=decoder_first_out, 124 | hidden=dec_second_hidden, 125 | edge_index=edge_index, 126 | edge_attr=edge_attr) 127 | return dec_second_out, dec_second_hidden 128 | -------------------------------------------------------------------------------- /SOTA/DCRNN/model/dcrnn_cell.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | from lib import utils_HIAM 5 | 6 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 7 | 8 | 9 | class LayerParams: 10 | def __init__(self, rnn_network: torch.nn.Module, layer_type: str): 11 | self._rnn_network = rnn_network 12 | self._params_dict = {} 13 | self._biases_dict = {} 14 | self._type = layer_type 15 | 16 | def get_weights(self, shape): 17 | if shape not in self._params_dict: 18 | nn_param = torch.nn.Parameter(torch.empty(*shape, device=device)) 19 | torch.nn.init.xavier_normal_(nn_param) 20 | self._params_dict[shape] = nn_param 21 | self._rnn_network.register_parameter('{}_weight_{}'.format(self._type, str(shape)), 22 | nn_param) 23 | return self._params_dict[shape] 24 | 25 | def get_biases(self, length, bias_start=0.0): 26 | if length not in self._biases_dict: 27 | biases = torch.nn.Parameter(torch.empty(length, device=device)) 28 | torch.nn.init.constant_(biases, bias_start) 29 | self._biases_dict[length] = biases 30 | self._rnn_network.register_parameter('{}_biases_{}'.format(self._type, str(length)), 31 | biases) 32 | 33 | return self._biases_dict[length] 34 | 35 | 36 | class DCGRUCell(torch.nn.Module): 37 | def __init__(self, num_units, adj_mx, max_diffusion_step, num_nodes, nonlinearity='tanh', 38 | filter_type="laplacian", use_gc_for_ru=True): 39 | """ 40 | 41 | :param num_units: 42 | :param adj_mx: 43 | :param max_diffusion_step: 44 | :param num_nodes: 45 | :param nonlinearity: 46 | :param filter_type: "laplacian", "random_walk", "dual_random_walk". 47 | :param use_gc_for_ru: whether to use Graph convolution to calculate the reset and update gates. 48 | """ 49 | 50 | super().__init__() 51 | self._activation = torch.tanh if nonlinearity == 'tanh' else torch.relu 52 | # support other nonlinearities up here? 53 | self._num_nodes = num_nodes 54 | self._num_units = num_units 55 | self._max_diffusion_step = max_diffusion_step 56 | self._supports = [] 57 | self._use_gc_for_ru = use_gc_for_ru 58 | supports = [] 59 | if filter_type == "laplacian": 60 | supports.append(utils_HIAM.calculate_scaled_laplacian(adj_mx, lambda_max=None)) 61 | elif filter_type == "random_walk": 62 | supports.append(utils_HIAM.calculate_random_walk_matrix(adj_mx).T) 63 | elif filter_type == "dual_random_walk": 64 | supports.append(utils_HIAM.calculate_random_walk_matrix(adj_mx).T) 65 | supports.append(utils_HIAM.calculate_random_walk_matrix(adj_mx.T).T) 66 | else: 67 | supports.append(utils_HIAM.calculate_scaled_laplacian(adj_mx)) 68 | for support in supports: 69 | self._supports.append(self._build_sparse_matrix(support)) 70 | 71 | self._fc_params = LayerParams(self, 'fc') 72 | self._gconv_params = LayerParams(self, 'gconv') 73 | 74 | @staticmethod 75 | def _build_sparse_matrix(L): 76 | L = L.tocoo() 77 | indices = np.column_stack((L.row, L.col)) 78 | # this is to ensure row-major ordering to equal torch.sparse.sparse_reorder(L) 79 | indices = indices[np.lexsort((indices[:, 0], indices[:, 1]))] 80 | L = torch.sparse_coo_tensor(indices.T, L.data, L.shape, device=device) 81 | return L 82 | 83 | def forward(self, inputs, hx): 84 | """Gated recurrent unit (GRU) with Graph Convolution. 85 | :param inputs: (B, num_nodes * input_dim) 86 | :param hx: (B, num_nodes * rnn_units) 87 | 88 | :return 89 | - Output: A `2-D` tensor with shape `(B, num_nodes * rnn_units)`. 90 | """ 91 | output_size = 2 * self._num_units 92 | if self._use_gc_for_ru: 93 | fn = self._gconv 94 | else: 95 | fn = self._fc 96 | value = torch.sigmoid(fn(inputs, hx, output_size, bias_start=1.0)) 97 | value = torch.reshape(value, (-1, self._num_nodes, output_size)) 98 | r, u = torch.split(tensor=value, split_size_or_sections=self._num_units, dim=-1) 99 | r = torch.reshape(r, (-1, self._num_nodes * self._num_units)) 100 | u = torch.reshape(u, (-1, self._num_nodes * self._num_units)) 101 | 102 | c = self._gconv(inputs, r * hx, self._num_units) 103 | if self._activation is not None: 104 | c = self._activation(c) 105 | 106 | new_state = u * hx + (1.0 - u) * c 107 | return new_state 108 | 109 | @staticmethod 110 | def _concat(x, x_): 111 | x_ = x_.unsqueeze(0) 112 | return torch.cat([x, x_], dim=0) 113 | 114 | def _fc(self, inputs, state, output_size, bias_start=0.0): 115 | batch_size = inputs.shape[0] 116 | inputs = torch.reshape(inputs, (batch_size * self._num_nodes, -1)) 117 | state = torch.reshape(state, (batch_size * self._num_nodes, -1)) 118 | inputs_and_state = torch.cat([inputs, state], dim=-1) 119 | input_size = inputs_and_state.shape[-1] 120 | weights = self._fc_params.get_weights((input_size, output_size)) 121 | value = torch.sigmoid(torch.matmul(inputs_and_state, weights)) 122 | biases = self._fc_params.get_biases(output_size, bias_start) 123 | value += biases 124 | return value 125 | 126 | def _gconv(self, inputs, state, output_size, bias_start=0.0): 127 | # Reshape input and state to (batch_size, num_nodes, input_dim/state_dim) 128 | batch_size = inputs.shape[0] 129 | inputs = torch.reshape(inputs, (batch_size, self._num_nodes, -1)) 130 | state = torch.reshape(state, (batch_size, self._num_nodes, -1)) 131 | inputs_and_state = torch.cat([inputs, state], dim=2) 132 | input_size = inputs_and_state.size(2) 133 | 134 | x = inputs_and_state 135 | x0 = x.permute(1, 2, 0) # (num_nodes, total_arg_size, batch_size) 136 | x0 = torch.reshape(x0, shape=[self._num_nodes, input_size * batch_size]) 137 | x = torch.unsqueeze(x0, 0) 138 | 139 | if self._max_diffusion_step == 0: 140 | pass 141 | else: 142 | for support in self._supports: 143 | x1 = torch.sparse.mm(support, x0) 144 | x = self._concat(x, x1) 145 | 146 | for k in range(2, self._max_diffusion_step + 1): 147 | x2 = 2 * torch.sparse.mm(support, x1) - x0 148 | x = self._concat(x, x2) 149 | x1, x0 = x2, x1 150 | 151 | num_matrices = len(self._supports) * self._max_diffusion_step + 1 # Adds for x itself. 152 | x = torch.reshape(x, shape=[num_matrices, self._num_nodes, input_size, batch_size]) 153 | x = x.permute(3, 1, 2, 0) # (batch_size, num_nodes, input_size, order) 154 | x = torch.reshape(x, shape=[batch_size * self._num_nodes, input_size * num_matrices]) 155 | 156 | weights = self._gconv_params.get_weights((input_size * num_matrices, output_size)) 157 | x = torch.matmul(x, weights) # (batch_size * self._num_nodes, output_size) 158 | 159 | biases = self._gconv_params.get_biases(output_size, bias_start) 160 | x += biases 161 | # Reshape res back to 2D: (batch_size, num_node, state_dim) -> (batch_size, num_node * state_dim) 162 | return torch.reshape(x, [batch_size, self._num_nodes * output_size]) 163 | -------------------------------------------------------------------------------- /SOTA/DCRNN/model/dcrnn_model.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | 5 | from .dcrnn_cell import DCGRUCell 6 | 7 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 8 | 9 | 10 | def count_parameters(model): 11 | return sum(p.numel() for p in model.parameters() if p.requires_grad) 12 | 13 | 14 | class Seq2SeqAttrs: 15 | def __init__(self, adj_mx, **model_kwargs): 16 | self.adj_mx = adj_mx 17 | self.max_diffusion_step = int(model_kwargs.get('max_diffusion_step', 2)) 18 | self.cl_decay_steps = int(model_kwargs.get('cl_decay_steps', 1000)) 19 | self.filter_type = model_kwargs.get('filter_type', 'laplacian') 20 | self.num_nodes = int(model_kwargs.get('num_nodes', 1)) 21 | self.num_rnn_layers = int(model_kwargs.get('num_rnn_layers', 1)) 22 | self.rnn_units = int(model_kwargs.get('rnn_units')) 23 | self.hidden_state_size = self.num_nodes * self.rnn_units 24 | 25 | 26 | class EncoderModel(nn.Module, Seq2SeqAttrs): 27 | def __init__(self, adj_mx, **model_kwargs): 28 | nn.Module.__init__(self) 29 | Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs) 30 | self.input_dim = int(model_kwargs.get('input_dim', 1)) 31 | self.seq_len = int(model_kwargs.get('seq_len')) # for the encoder 32 | self.dcgru_layers = nn.ModuleList( 33 | [DCGRUCell(self.rnn_units, adj_mx, self.max_diffusion_step, self.num_nodes, 34 | filter_type=self.filter_type) for _ in range(self.num_rnn_layers)]) 35 | 36 | def forward(self, inputs, hidden_state=None): 37 | """ 38 | Encoder forward pass. 39 | 40 | :param inputs: shape (batch_size, self.num_nodes * self.input_dim) 41 | :param hidden_state: (num_layers, batch_size, self.hidden_state_size) 42 | optional, zeros if not provided 43 | :return: output: # shape (batch_size, self.hidden_state_size) 44 | hidden_state # shape (num_layers, batch_size, self.hidden_state_size) 45 | (lower indices mean lower layers) 46 | """ 47 | batch_size, _ = inputs.size() 48 | if hidden_state is None: 49 | hidden_state = torch.zeros((self.num_rnn_layers, batch_size, self.hidden_state_size), 50 | device=device) 51 | hidden_states = [] 52 | output = inputs 53 | for layer_num, dcgru_layer in enumerate(self.dcgru_layers): 54 | next_hidden_state = dcgru_layer(output, hidden_state[layer_num]) 55 | hidden_states.append(next_hidden_state) 56 | output = next_hidden_state 57 | 58 | return output, torch.stack(hidden_states) # runs in O(num_layers) so not too slow 59 | 60 | 61 | class DecoderModel(nn.Module, Seq2SeqAttrs): 62 | def __init__(self, adj_mx, **model_kwargs): 63 | # super().__init__(is_training, adj_mx, **model_kwargs) 64 | nn.Module.__init__(self) 65 | Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs) 66 | self.output_dim = int(model_kwargs.get('output_dim', 1)) 67 | self.horizon = int(model_kwargs.get('horizon', 1)) # for the decoder 68 | self.projection_layer = nn.Linear(self.rnn_units, self.output_dim) 69 | self.dcgru_layers = nn.ModuleList( 70 | [DCGRUCell(self.rnn_units, adj_mx, self.max_diffusion_step, self.num_nodes, 71 | filter_type=self.filter_type) for _ in range(self.num_rnn_layers)]) 72 | 73 | def forward(self, inputs, hidden_state=None): 74 | """ 75 | Decoder forward pass. 76 | 77 | :param inputs: shape (batch_size, self.num_nodes * self.output_dim) 78 | :param hidden_state: (num_layers, batch_size, self.hidden_state_size) 79 | optional, zeros if not provided 80 | :return: output: # shape (batch_size, self.num_nodes * self.output_dim) 81 | hidden_state # shape (num_layers, batch_size, self.hidden_state_size) 82 | (lower indices mean lower layers) 83 | """ 84 | hidden_states = [] 85 | output = inputs 86 | for layer_num, dcgru_layer in enumerate(self.dcgru_layers): 87 | next_hidden_state = dcgru_layer(output, hidden_state[layer_num]) 88 | hidden_states.append(next_hidden_state) 89 | output = next_hidden_state 90 | 91 | projected = self.projection_layer(output.view(-1, self.rnn_units)) 92 | output = projected.view(-1, self.num_nodes * self.output_dim) 93 | 94 | return output, torch.stack(hidden_states) 95 | 96 | 97 | class DCRNNModel(nn.Module, Seq2SeqAttrs): 98 | def __init__(self, adj_mx, logger, **model_kwargs): 99 | super().__init__() 100 | Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs) 101 | self.encoder_model = EncoderModel(adj_mx, **model_kwargs) 102 | self.decoder_model = DecoderModel(adj_mx, **model_kwargs) 103 | self.cl_decay_steps = int(model_kwargs.get('cl_decay_steps', 1000)) 104 | self.use_curriculum_learning = bool(model_kwargs.get('use_curriculum_learning', False)) 105 | self._logger = logger 106 | 107 | def _compute_sampling_threshold(self, batches_seen): 108 | return self.cl_decay_steps / ( 109 | self.cl_decay_steps + np.exp(batches_seen / self.cl_decay_steps)) 110 | 111 | def encoder(self, inputs): 112 | """ 113 | encoder forward pass on t time steps 114 | :param inputs: shape (seq_len, batch_size, num_sensor * input_dim) 115 | :return: encoder_hidden_state: (num_layers, batch_size, self.hidden_state_size) 116 | """ 117 | encoder_hidden_state = None 118 | for t in range(self.encoder_model.seq_len): 119 | _, encoder_hidden_state = self.encoder_model(inputs[t], encoder_hidden_state) 120 | 121 | return encoder_hidden_state 122 | 123 | def decoder(self, encoder_hidden_state, labels=None, batches_seen=None): 124 | """ 125 | Decoder forward pass 126 | :param encoder_hidden_state: (num_layers, batch_size, self.hidden_state_size) 127 | :param labels: (self.horizon, batch_size, self.num_nodes * self.output_dim) [optional, not exist for inference] 128 | :param batches_seen: global step [optional, not exist for inference] 129 | :return: output: (self.horizon, batch_size, self.num_nodes * self.output_dim) 130 | """ 131 | batch_size = encoder_hidden_state.size(1) 132 | go_symbol = torch.zeros((batch_size, self.num_nodes * self.decoder_model.output_dim), 133 | device=device) 134 | decoder_hidden_state = encoder_hidden_state 135 | decoder_input = go_symbol 136 | 137 | outputs = [] 138 | 139 | for t in range(self.decoder_model.horizon): 140 | decoder_output, decoder_hidden_state = self.decoder_model(decoder_input, 141 | decoder_hidden_state) 142 | decoder_input = decoder_output 143 | outputs.append(decoder_output) 144 | if self.training and self.use_curriculum_learning: 145 | c = np.random.uniform(0, 1) 146 | if c < self._compute_sampling_threshold(batches_seen): 147 | decoder_input = labels[t] 148 | outputs = torch.stack(outputs) 149 | return outputs 150 | 151 | def forward(self, inputs, labels=None, batches_seen=None): 152 | """ 153 | seq2seq forward pass 154 | :param inputs: shape (seq_len, batch_size, num_sensor * input_dim) 155 | :param labels: shape (horizon, batch_size, num_sensor * output) 156 | :param batches_seen: batches seen till now 157 | :return: output: (self.horizon, batch_size, self.num_nodes * self.output_dim) 158 | """ 159 | encoder_hidden_state = self.encoder(inputs) 160 | self._logger.debug("Encoder complete, starting decoder") 161 | outputs = self.decoder(encoder_hidden_state, labels, batches_seen=batches_seen) 162 | self._logger.debug("Decoder complete") 163 | if batches_seen == 0: 164 | self._logger.info( 165 | "Total trainable parameters {}".format(count_parameters(self)) 166 | ) 167 | return outputs 168 | -------------------------------------------------------------------------------- /SOTA/DCRNN/lib/utils.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import numpy as np 3 | import os 4 | import pickle 5 | import scipy.sparse as sp 6 | import sys 7 | # import tensorflow as tf 8 | 9 | from scipy.sparse import linalg 10 | 11 | 12 | class DataLoader(object): 13 | def __init__(self, xs, ys, batch_size, pad_with_last_sample=True, shuffle=False): 14 | """ 15 | 16 | :param xs: 17 | :param ys: 18 | :param batch_size: 19 | :param pad_with_last_sample: pad with the last sample to make number of samples divisible to batch_size. 20 | """ 21 | self.batch_size = batch_size 22 | self.current_ind = 0 23 | if pad_with_last_sample: 24 | num_padding = (batch_size - (len(xs) % batch_size)) % batch_size 25 | x_padding = np.repeat(xs[-1:], num_padding, axis=0) 26 | y_padding = np.repeat(ys[-1:], num_padding, axis=0) 27 | xs = np.concatenate([xs, x_padding], axis=0) 28 | ys = np.concatenate([ys, y_padding], axis=0) 29 | self.size = len(xs) 30 | self.num_batch = int(self.size // self.batch_size) 31 | if shuffle: 32 | permutation = np.random.permutation(self.size) 33 | xs, ys = xs[permutation], ys[permutation] 34 | self.xs = xs 35 | self.ys = ys 36 | 37 | def get_iterator(self): 38 | self.current_ind = 0 39 | 40 | def _wrapper(): 41 | while self.current_ind < self.num_batch: 42 | start_ind = self.batch_size * self.current_ind 43 | end_ind = min(self.size, self.batch_size * (self.current_ind + 1)) 44 | x_i = self.xs[start_ind: end_ind, ...] 45 | y_i = self.ys[start_ind: end_ind, ...] 46 | yield (x_i, y_i) 47 | self.current_ind += 1 48 | 49 | return _wrapper() 50 | 51 | 52 | class StandardScaler: 53 | """ 54 | Standard the input 55 | """ 56 | 57 | def __init__(self, mean, std): 58 | self.mean = mean 59 | self.std = std 60 | 61 | def transform(self, data): 62 | return (data - self.mean) / self.std * 1.0 63 | 64 | def inverse_transform(self, data): 65 | return (data * self.std) + self.mean 66 | 67 | def calculate_normalized_laplacian(adj): 68 | """ 69 | # L = D^-1/2 (D-A) D^-1/2 = I - D^-1/2 A D^-1/2 70 | # D = diag(A 1) 71 | :param adj: 72 | :return: 73 | """ 74 | adj = sp.coo_matrix(adj) 75 | d = np.array(adj.sum(1)) 76 | d_inv_sqrt = np.power(d, -0.5).flatten() 77 | d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0. 78 | d_mat_inv_sqrt = sp.diags(d_inv_sqrt) 79 | normalized_laplacian = sp.eye(adj.shape[0]) - adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt).tocoo() 80 | return normalized_laplacian 81 | 82 | 83 | def calculate_random_walk_matrix(adj_mx): 84 | adj_mx = sp.coo_matrix(adj_mx) 85 | d = np.array(adj_mx.sum(1)) 86 | d_inv = np.power(d, -1).flatten() 87 | d_inv[np.isinf(d_inv)] = 0. 88 | d_mat_inv = sp.diags(d_inv) 89 | random_walk_mx = d_mat_inv.dot(adj_mx).tocoo() 90 | return random_walk_mx 91 | 92 | 93 | def calculate_reverse_random_walk_matrix(adj_mx): 94 | return calculate_random_walk_matrix(np.transpose(adj_mx)) 95 | 96 | 97 | def calculate_scaled_laplacian(adj_mx, lambda_max=2, undirected=True): 98 | if undirected: 99 | adj_mx = np.maximum.reduce([adj_mx, adj_mx.T]) 100 | L = calculate_normalized_laplacian(adj_mx) 101 | if lambda_max is None: 102 | lambda_max, _ = linalg.eigsh(L, 1, which='LM') 103 | lambda_max = lambda_max[0] 104 | L = sp.csr_matrix(L) 105 | M, _ = L.shape 106 | I = sp.identity(M, format='csr', dtype=L.dtype) 107 | L = (2 / lambda_max * L) - I 108 | return L.astype(np.float32) 109 | 110 | 111 | def config_logging(log_dir, log_filename='info.log', level=logging.INFO): 112 | # Add file handler and stdout handler 113 | formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') 114 | # Create the log directory if necessary. 115 | try: 116 | os.makedirs(log_dir) 117 | except OSError: 118 | pass 119 | file_handler = logging.FileHandler(os.path.join(log_dir, log_filename)) 120 | file_handler.setFormatter(formatter) 121 | file_handler.setLevel(level=level) 122 | # Add console handler. 123 | console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') 124 | console_handler = logging.StreamHandler(sys.stdout) 125 | console_handler.setFormatter(console_formatter) 126 | console_handler.setLevel(level=level) 127 | logging.basicConfig(handlers=[file_handler, console_handler], level=level) 128 | 129 | 130 | def get_logger(log_dir, name, log_filename='info.log', level=logging.INFO): 131 | logger = logging.getLogger(name) 132 | logger.setLevel(level) 133 | # Add file handler and stdout handler 134 | formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') 135 | file_handler = logging.FileHandler(os.path.join(log_dir, log_filename)) 136 | file_handler.setFormatter(formatter) 137 | # Add console handler. 138 | console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') 139 | console_handler = logging.StreamHandler(sys.stdout) 140 | console_handler.setFormatter(console_formatter) 141 | logger.addHandler(file_handler) 142 | logger.addHandler(console_handler) 143 | # Add google cloud log handler 144 | logger.info('Log directory: %s', log_dir) 145 | return logger 146 | 147 | 148 | def get_total_trainable_parameter_size(): 149 | """ 150 | Calculates the total number of trainable parameters in the current graph. 151 | :return: 152 | """ 153 | total_parameters = 0 154 | for variable in tf.trainable_variables(): 155 | # shape is an array of tf.Dimension 156 | total_parameters += np.product([x.value for x in variable.get_shape()]) 157 | return total_parameters 158 | 159 | # return data 160 | def load_dataset(dataset_dir, 161 | batch_size, 162 | test_batch_size=None, 163 | ds_type = "", 164 | scaler_axis=(0, 165 | 1, 166 | 2, 167 | 3), 168 | **kwargs): 169 | data = {} 170 | if ds_type == "od": 171 | for category in ['train', 'val', 'test']: 172 | cat_data = load_pickle(os.path.join(dataset_dir, category + '.pkl')) 173 | data['x_' + category] = cat_data['finished'] 174 | data['y_' + category] = cat_data['y'] 175 | elif ds_type == "do": 176 | for category in ['train', 'val', 'test']: 177 | cat_data = load_pickle(os.path.join(dataset_dir, category + '_do.pkl')) 178 | data['x_' + category] = cat_data['do_x'] / 1.0 179 | data['y_' + category] = cat_data['do_y'] / 1.0 180 | 181 | scaler = StandardScaler(mean=data['x_train'].mean(axis=scaler_axis), 182 | std=data['x_train'].std(axis=scaler_axis)) 183 | # Data format 184 | for category in ['train', 'val', 'test']: 185 | data['x_' + category] = scaler.transform(data['x_' + category]) 186 | data['y_' + category] = scaler.transform(data['y_' + category]) 187 | data['train_loader'] = DataLoader(data['x_train'], 188 | data['y_train'], 189 | batch_size, 190 | shuffle=True) 191 | data['val_loader'] = DataLoader(data['x_val'], 192 | data['y_val'], 193 | test_batch_size, 194 | shuffle=False) 195 | data['test_loader'] = DataLoader(data['x_test'], 196 | data['y_test'], 197 | test_batch_size, 198 | shuffle=False) 199 | data['scaler'] = scaler 200 | return data 201 | 202 | def load_graph_data(pkl_filename): 203 | sensor_ids, sensor_id_to_ind, adj_mx = load_pickle(pkl_filename) 204 | return sensor_ids, sensor_id_to_ind, adj_mx 205 | 206 | 207 | def load_pickle(pickle_file): 208 | try: 209 | with open(pickle_file, 'rb') as f: 210 | pickle_data = pickle.load(f) 211 | except UnicodeDecodeError as e: 212 | with open(pickle_file, 'rb') as f: 213 | pickle_data = pickle.load(f, encoding='latin1') 214 | except Exception as e: 215 | print('Unable to load data ', pickle_file, ':', e) 216 | raise 217 | return pickle_data 218 | -------------------------------------------------------------------------------- /SOTA/Graph-WaveNet-master/util1.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | import numpy as np 3 | import os 4 | import scipy.sparse as sp 5 | import torch 6 | from scipy.sparse import linalg 7 | import logging 8 | import sys 9 | 10 | 11 | class DataLoader(object): 12 | def __init__(self, xs, ys, batch_size, pad_with_last_sample=True): 13 | """ 14 | :param xs: 15 | :param ys: 16 | :param batch_size: 17 | :param pad_with_last_sample: pad with the last sample to make number of samples divisible to batch_size. 18 | """ 19 | self.batch_size = batch_size 20 | self.current_ind = 0 21 | if pad_with_last_sample: 22 | num_padding = (batch_size - (len(xs) % batch_size)) % batch_size 23 | x_padding = np.repeat(xs[-1:], num_padding, axis=0) 24 | y_padding = np.repeat(ys[-1:], num_padding, axis=0) 25 | xs = np.concatenate([xs, x_padding], axis=0) 26 | ys = np.concatenate([ys, y_padding], axis=0) 27 | self.size = len(xs) 28 | self.num_batch = int(self.size // self.batch_size) 29 | self.xs = xs 30 | self.ys = ys 31 | 32 | def shuffle(self): 33 | permutation = np.random.permutation(self.size) 34 | xs, ys = self.xs[permutation], self.ys[permutation] 35 | self.xs = xs 36 | self.ys = ys 37 | 38 | def get_iterator(self): 39 | self.current_ind = 0 40 | 41 | def _wrapper(): 42 | while self.current_ind < self.num_batch: 43 | start_ind = self.batch_size * self.current_ind 44 | end_ind = min(self.size, self.batch_size * (self.current_ind + 1)) 45 | x_i = self.xs[start_ind: end_ind, ...] 46 | y_i = self.ys[start_ind: end_ind, ...] 47 | yield (x_i, y_i) 48 | self.current_ind += 1 49 | 50 | return _wrapper() 51 | 52 | 53 | class StandardScaler(): 54 | """ 55 | Standard the input 56 | """ 57 | 58 | def __init__(self, mean, std): 59 | self.mean = mean 60 | self.std = std 61 | 62 | def transform(self, data): 63 | return (data - self.mean) / self.std 64 | 65 | def inverse_transform(self, data): 66 | return (data * self.std) + self.mean 67 | 68 | 69 | def sym_adj(adj): 70 | """Symmetrically normalize adjacency matrix.""" 71 | adj = sp.coo_matrix(adj) 72 | rowsum = np.array(adj.sum(1)) 73 | d_inv_sqrt = np.power(rowsum, -0.5).flatten() 74 | d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0. 75 | d_mat_inv_sqrt = sp.diags(d_inv_sqrt) 76 | return adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt).astype(np.float32).todense() 77 | 78 | 79 | def asym_adj(adj): 80 | adj = adj.astype(float) 81 | adj = sp.coo_matrix(adj) 82 | rowsum = np.array(adj.sum(1)).flatten() 83 | d_inv = np.power(rowsum, -1).flatten() 84 | d_inv[np.isinf(d_inv)] = 0. 85 | d_mat= sp.diags(d_inv) 86 | return d_mat.dot(adj).astype(np.float32).todense() 87 | 88 | 89 | def calculate_normalized_laplacian(adj): 90 | """ 91 | # L = D^-1/2 (D-A) D^-1/2 = I - D^-1/2 A D^-1/2 92 | # D = diag(A 1) 93 | :param adj: 94 | :return: 95 | """ 96 | adj = sp.coo_matrix(adj) 97 | d = np.array(adj.sum(1)) 98 | d_inv_sqrt = np.power(d, -0.5).flatten() 99 | d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0. 100 | d_mat_inv_sqrt = sp.diags(d_inv_sqrt) 101 | normalized_laplacian = sp.eye(adj.shape[0]) - adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt).tocoo() 102 | return normalized_laplacian 103 | 104 | 105 | def calculate_scaled_laplacian(adj_mx, lambda_max=2, undirected=True): 106 | if undirected: 107 | adj_mx = np.maximum.reduce([adj_mx, adj_mx.T]) 108 | L = calculate_normalized_laplacian(adj_mx) 109 | if lambda_max is None: 110 | lambda_max, _ = linalg.eigsh(L, 1, which='LM') 111 | lambda_max = lambda_max[0] 112 | L = sp.csr_matrix(L) 113 | M, _ = L.shape 114 | I = sp.identity(M, format='csr', dtype=L.dtype) 115 | L = (2 / lambda_max * L) - I 116 | return L.astype(np.float32).todense() 117 | 118 | 119 | def load_pickle(pickle_file): 120 | try: 121 | with open(pickle_file, 'rb') as f: 122 | pickle_data = pickle.load(f) 123 | except UnicodeDecodeError as e: 124 | with open(pickle_file, 'rb') as f: 125 | pickle_data = pickle.load(f, encoding='latin1') 126 | except Exception as e: 127 | print('Unable to load data ', pickle_file, ':', e) 128 | raise 129 | return pickle_data 130 | 131 | 132 | def load_adj(pkl_filename, adjtype): 133 | # sensor_ids, sensor_id_to_ind, adj_mx = load_pickle(pkl_filename) 134 | adj_mx = load_pickle(pkl_filename) 135 | if adjtype == "scalap": 136 | adj = [calculate_scaled_laplacian(adj_mx)] 137 | elif adjtype == "normlap": 138 | adj = [calculate_normalized_laplacian(adj_mx).astype(np.float32).todense()] 139 | elif adjtype == "symnadj": 140 | adj = [sym_adj(adj_mx)] 141 | elif adjtype == "transition": 142 | adj = [asym_adj(adj_mx)] 143 | elif adjtype == "doubletransition": 144 | adj = [asym_adj(adj_mx), asym_adj(np.transpose(adj_mx))] 145 | elif adjtype == "identity": 146 | adj = [np.diag(np.ones(adj_mx.shape[0])).astype(np.float32)] 147 | else: 148 | error = 0 149 | assert error, "adj type not defined" 150 | # return sensor_ids, sensor_id_to_ind, adj 151 | return adj 152 | 153 | 154 | def load_dataset(dataset_dir, batch_size, valid_batch_size= None, test_batch_size=None): 155 | data = {} 156 | for category in ['train', 'val', 'test']: 157 | # cat_data = np.load(os.path.join(dataset_dir, category + '.npz')) 158 | # data['x_' + category] = cat_data['x'] 159 | # data['y_' + category] = cat_data['y'] 160 | cat_data = load_pickle(os.path.join(dataset_dir, category + '.pkl')) 161 | data['x_' + category] = cat_data['finished'] / 1.0 162 | data['y_' + category] = cat_data['y'] / 1.0 163 | scaler = StandardScaler(mean=data['x_train'][..., 0].mean(), std=data['x_train'][..., 0].std()) 164 | # Data format 165 | for category in ['train', 'val', 'test']: 166 | data['x_' + category][..., 0] = scaler.transform(data['x_' + category][..., 0]) 167 | data['train_loader'] = DataLoader(data['x_train'], data['y_train'], batch_size) 168 | data['val_loader'] = DataLoader(data['x_val'], data['y_val'], valid_batch_size) 169 | data['test_loader'] = DataLoader(data['x_test'], data['y_test'], test_batch_size) 170 | data['scaler'] = scaler 171 | return data 172 | 173 | def load_dataset_do(dataset_dir, batch_size, valid_batch_size= None, test_batch_size=None): 174 | data = {} 175 | for category in ['train', 'val', 'test']: 176 | # cat_data = np.load(os.path.join(dataset_dir, category + '.npz')) 177 | # data['x_' + category] = cat_data['x'] 178 | # data['y_' + category] = cat_data['y'] 179 | cat_data = load_pickle(os.path.join(dataset_dir, category + '_do.pkl')) 180 | data['x_' + category] = cat_data['do_x'] / 1.0 181 | data['y_' + category] = cat_data['do_y'] / 1.0 182 | scaler = StandardScaler(mean=data['x_train'][..., 0].mean(), std=data['x_train'][..., 0].std()) 183 | # Data format 184 | for category in ['train', 'val', 'test']: 185 | data['x_' + category][..., 0] = scaler.transform(data['x_' + category][..., 0]) 186 | data['train_loader'] = DataLoader(data['x_train'], data['y_train'], batch_size) 187 | data['val_loader'] = DataLoader(data['x_val'], data['y_val'], valid_batch_size) 188 | data['test_loader'] = DataLoader(data['x_test'], data['y_test'], test_batch_size) 189 | data['scaler'] = scaler 190 | return data 191 | 192 | def get_logger(log_dir, 193 | name, 194 | log_filename='info.log', 195 | level=logging.INFO, 196 | write_to_file=True): 197 | logger = logging.getLogger(name) 198 | logger.setLevel(level) 199 | logger.propagate = False 200 | # Add file handler and stdout handler 201 | formatter = logging.Formatter( 202 | '%(asctime)s - %(name)s - %(levelname)s - %(message)s') 203 | file_handler = logging.FileHandler(os.path.join(log_dir, log_filename)) 204 | file_handler.setFormatter(formatter) 205 | # Add console handler. 206 | console_formatter = logging.Formatter( 207 | '%(asctime)s - %(levelname)s - %(message)s') 208 | console_handler = logging.StreamHandler(sys.stdout) 209 | console_handler.setFormatter(console_formatter) 210 | if write_to_file is True: 211 | logger.addHandler(file_handler) 212 | logger.addHandler(console_handler) 213 | # Add google cloud log handler 214 | logger.info('Log directory: %s', log_dir) 215 | return logger 216 | 217 | 218 | def setup_tensorboard_writer(base_dir, comment, **kargs): 219 | from torch.utils.tensorboard import SummaryWriter 220 | tf_file_name = '_'.join(k + str(v) for k, v in kargs.items()) 221 | tf_file_path = os.path.join(base_dir, tf_file_name) + comment 222 | return SummaryWriter( 223 | log_dir=tf_file_path 224 | ), tf_file_path 225 | -------------------------------------------------------------------------------- /SOTA/Graph-WaveNet-master/model1.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.autograd import Variable 5 | import sys 6 | 7 | 8 | class nconv(nn.Module): 9 | def __init__(self): 10 | super(nconv,self).__init__() 11 | 12 | def forward(self,x, A): 13 | x = torch.einsum('ncvl,vw->ncwl',(x,A)) 14 | return x.contiguous() 15 | 16 | class linear(nn.Module): 17 | def __init__(self,c_in,c_out): 18 | super(linear,self).__init__() 19 | self.mlp = torch.nn.Conv2d(c_in, c_out, kernel_size=(1, 1), padding=(0,0), stride=(1,1), bias=True) 20 | 21 | def forward(self,x): 22 | return self.mlp(x) 23 | 24 | class gcn(nn.Module): 25 | def __init__(self,c_in,c_out,dropout,support_len=3,order=2): 26 | super(gcn,self).__init__() 27 | self.nconv = nconv() 28 | c_in = (order*support_len+1)*c_in 29 | self.mlp = linear(c_in,c_out) 30 | self.dropout = dropout 31 | self.order = order 32 | 33 | def forward(self,x,support): 34 | out = [x] 35 | for a in support: 36 | x1 = self.nconv(x,a) 37 | out.append(x1) 38 | for k in range(2, self.order + 1): 39 | x2 = self.nconv(x1,a) 40 | out.append(x2) 41 | x1 = x2 42 | 43 | h = torch.cat(out,dim=1) 44 | h = self.mlp(h) 45 | h = F.dropout(h, self.dropout, training=self.training) 46 | return h 47 | 48 | 49 | class gwnet(nn.Module): 50 | def __init__(self, device, num_nodes, dropout=0.3, supports=None, gcn_bool=True, addaptadj=True, aptinit=None, 51 | in_dim=2,out_dim=12,residual_channels=32,dilation_channels=32,skip_channels=256,end_channels=512, 52 | kernel_size=2,blocks=4,layers=2): 53 | """ 54 | args: 55 | - out_dim: output node feature 56 | - in_dim: input node feature 57 | - skip_channels: hidden node feature 58 | - residual_channels: 59 | - end_channels: 60 | - layers: 61 | - kernel_size: 62 | returns: 63 | - prediction: [bn, feat_dim, node_num, seq] 64 | """ 65 | super(gwnet, self).__init__() 66 | self.device = device 67 | self.dropout = dropout 68 | self.blocks = blocks 69 | self.layers = layers 70 | self.gcn_bool = gcn_bool 71 | self.addaptadj = addaptadj 72 | 73 | self.filter_convs = nn.ModuleList() 74 | self.gate_convs = nn.ModuleList() 75 | self.residual_convs = nn.ModuleList() 76 | self.skip_convs = nn.ModuleList() 77 | self.pads = nn.ModuleList() 78 | self.bn = nn.ModuleList() 79 | self.gconv = nn.ModuleList() 80 | 81 | self.start_conv = nn.Conv2d(in_channels=in_dim, 82 | out_channels=residual_channels, 83 | kernel_size=(1,1)) 84 | self.supports = supports 85 | 86 | receptive_field = 1 87 | 88 | self.supports_len = 0 89 | if supports is not None: 90 | self.supports_len += len(supports) 91 | 92 | if gcn_bool and addaptadj: 93 | if aptinit is None: 94 | if supports is None: 95 | self.supports = [] 96 | self.nodevec1 = nn.Parameter(torch.randn(num_nodes, 10).to(device), requires_grad=True).to(device) 97 | self.nodevec2 = nn.Parameter(torch.randn(10, num_nodes).to(device), requires_grad=True).to(device) 98 | self.supports_len +=1 99 | else: 100 | if supports is None: 101 | self.supports = [] 102 | m, p, n = torch.svd(aptinit) 103 | initemb1 = torch.mm(m[:, :10], torch.diag(p[:10] ** 0.5)) 104 | initemb2 = torch.mm(torch.diag(p[:10] ** 0.5), n[:, :10].t()) 105 | self.nodevec1 = nn.Parameter(initemb1, requires_grad=True).to(device) 106 | self.nodevec2 = nn.Parameter(initemb2, requires_grad=True).to(device) 107 | self.supports_len += 1 108 | 109 | 110 | for b in range(blocks): 111 | additional_scope = kernel_size - 1 112 | new_dilation = 1 113 | for i in range(layers): 114 | # dilated convolutions 115 | pad_size = new_dilation * (kernel_size - 1) 116 | self.pads.append( 117 | nn.ZeroPad2d((pad_size, 0, 0, 0)) 118 | ) 119 | 120 | self.filter_convs.append(nn.Conv2d(in_channels=residual_channels, 121 | out_channels=dilation_channels, 122 | kernel_size=(1,kernel_size),dilation=new_dilation)) 123 | 124 | self.gate_convs.append(nn.Conv1d(in_channels=residual_channels, 125 | out_channels=dilation_channels, 126 | kernel_size=(1, kernel_size), dilation=new_dilation)) 127 | 128 | # 1x1 convolution for residual connection 129 | self.residual_convs.append(nn.Conv1d(in_channels=dilation_channels, 130 | out_channels=residual_channels, 131 | kernel_size=(1, 1))) 132 | 133 | # 1x1 convolution for skip connection 134 | self.skip_convs.append(nn.Conv1d(in_channels=dilation_channels, 135 | out_channels=skip_channels, 136 | kernel_size=(1, 1))) 137 | self.bn.append(nn.BatchNorm2d(residual_channels)) 138 | new_dilation *=2 139 | receptive_field += additional_scope 140 | additional_scope *= 2 141 | if self.gcn_bool: 142 | self.gconv.append(gcn(dilation_channels,residual_channels,dropout,support_len=self.supports_len)) 143 | 144 | 145 | 146 | self.end_conv_1 = nn.Conv2d(in_channels=skip_channels, 147 | out_channels=end_channels, 148 | kernel_size=(1,1), 149 | bias=True) 150 | 151 | self.end_conv_2 = nn.Conv2d(in_channels=end_channels, 152 | out_channels=out_dim, 153 | kernel_size=(1,1), 154 | bias=True) 155 | 156 | self.receptive_field = receptive_field 157 | 158 | 159 | 160 | def forward(self, input): 161 | in_len = input.size(3) 162 | input = input.transpose(1, 3) # [bn, in_dim, node_num, seq] 163 | x = input # x.shape (32,4,80,31) 164 | x = self.start_conv(x) # x.shape 32, 32, 80, 31 165 | skip = 0 166 | 167 | # calculate the current adaptive adj matrix once per iteration 168 | new_supports = None 169 | if self.gcn_bool and self.addaptadj and self.supports is not None: 170 | adp = F.softmax(F.relu(torch.mm(self.nodevec1, self.nodevec2)), dim=1) 171 | new_supports = self.supports + [adp] 172 | 173 | # WaveNet layers 174 | for i in range(self.blocks * self.layers): 175 | 176 | # |----------------------------------------| *residual* 177 | # | | 178 | # | |-- conv -- tanh --| | 179 | # -> dilate -|----| * ----|-- 1x1 -- + --> *input* 180 | # |-- conv -- sigm --| | 181 | # 1x1 182 | # | 183 | # ---------------------------------------> + -------------> *skip* 184 | 185 | #(dilation, init_dilation) = self.dilations[i] 186 | 187 | #residual = dilation_func(x, dilation, init_dilation, i) 188 | residual = x 189 | # dilated convolution 190 | # train with paddings 191 | filter = self.filter_convs[i](self.pads[i](residual)) 192 | filter = torch.tanh(filter) 193 | gate = self.gate_convs[i](self.pads[i](residual)) 194 | gate = torch.sigmoid(gate) 195 | 196 | x = filter * gate 197 | # parametrized skip connection 198 | 199 | s = x 200 | s = self.skip_convs[i](s) 201 | try: 202 | skip = skip[:, :, :, -s.size(3):] 203 | except: 204 | skip = 0 205 | skip = s + skip 206 | 207 | 208 | if self.gcn_bool and self.supports is not None: 209 | if self.addaptadj: 210 | x = self.gconv[i](x, new_supports) 211 | else: 212 | x = self.gconv[i](x,self.supports) 213 | else: 214 | x = self.residual_convs[i](x) 215 | 216 | x = x + residual[:, :, :, -x.size(3):] 217 | 218 | x = self.bn[i](x) 219 | x = F.relu(skip) 220 | # x = x.transpose(1, 3) 221 | x = F.relu(self.end_conv_1(x)) 222 | x = self.end_conv_2(x) 223 | return x.transpose(1, 3) # [bn, seq, node_num, feat] 224 | 225 | 226 | 227 | 228 | 229 | -------------------------------------------------------------------------------- /SOTA/Graph-WaveNet-master/train1.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import argparse 4 | import time 5 | import util1 as util 6 | import metrics 7 | # import matplotlib.pyplot as plt 8 | from engine1 import trainer 9 | import os 10 | 11 | parser = argparse.ArgumentParser() 12 | parser.add_argument('--data',type=str,default='data/hangzhou/OD/OD_26',help='data path') 13 | parser.add_argument('--adjdata',type=str,default='data/hangzhou/graph_hz_conn.pkl',help='adj data path') 14 | parser.add_argument('--adjtype',type=str,default='symnadj',help='adj type') 15 | parser.add_argument('--gcn_bool',action='store_true',help='whether to add graph convolution layer') 16 | parser.add_argument('--aptonly',action='store_true',help='whether only adaptive adj') 17 | parser.add_argument('--addaptadj',action='store_true',help='whether add adaptive adj') 18 | parser.add_argument('--randomadj',action='store_true',help='whether random initialize adaptive adj') 19 | parser.add_argument('--seq_length',type=int,default=4,help='') 20 | parser.add_argument('--nhid',type=int,default=32,help='') 21 | parser.add_argument('--in_dim',type=int,default=26,help='inputs dimension') 22 | parser.add_argument('--num_nodes',type=int,default=80,help='number of nodes') 23 | parser.add_argument('--out_dim',type=int,default=26,help='') 24 | parser.add_argument('--batch_size',type=int,default=32,help='batch size') 25 | parser.add_argument('--learning_rate',type=float,default=0.0005,help='learning rate') 26 | parser.add_argument('--dropout',type=float,default=0.1,help='dropout rate') 27 | parser.add_argument('--weight_decay',type=float,default=1e-5,help='weight decay rate') 28 | parser.add_argument('--epochs',type=int,default=400,help='') 29 | parser.add_argument('--print_every',type=int,default=1,help='') 30 | parser.add_argument('--save',type=str,default='data/checkpoint',help='save path') 31 | parser.add_argument('--expid',type=int,default=1,help='experiment id') 32 | parser.add_argument('--seed',type=int,default=777,help='random seed') 33 | 34 | parser.add_argument('--exp_base',type=str,default="data",help='log_base_dir') 35 | parser.add_argument('--runs',type=str,default="debug",help='log runs name') 36 | 37 | parser.add_argument('--train_type',type=str,default="od",help='od training or do training') 38 | 39 | args = parser.parse_args() 40 | 41 | def evaluate(scaler, dataloader, device, engine, logger,type, cat_list): 42 | results = {} 43 | y_preds = [] 44 | gt = [] 45 | for iter, (x, y) in enumerate(dataloader[type + '_loader'].get_iterator()): 46 | test_x = torch.tensor(x, dtype=torch.float, device=device) 47 | # testx = testx.transpose(1,3) 48 | with torch.no_grad(): 49 | preds = engine.model(test_x) # .transpose(1,3) 50 | # preds = preds.transpose(1,3) 51 | test_y = torch.tensor(y, dtype=torch.float, device=device) 52 | y_preds.append(preds.detach().cpu().numpy()) 53 | gt.append(test_y.detach().cpu().numpy()) 54 | 55 | res = [] 56 | for category in cat_list: 57 | y_preds = np.concatenate(y_preds, axis=0) # concat in batch_size dim. 58 | gt = np.concatenate(gt, axis=0) 59 | mae_list = [] 60 | mape_net_list = [] 61 | rmse_list = [] 62 | mae_sum = 0 63 | 64 | mape_net_sum = 0 65 | rmse_sum = 0 66 | horizon = 4 67 | for horizon_i in range(horizon): 68 | y_pred = scaler.inverse_transform( 69 | y_preds[:, horizon_i, :, :]) 70 | y_pred[y_pred < 0] = 0 71 | y_truth = gt[:, horizon_i, :, :] 72 | mae = metrics.masked_mae_np(y_pred, y_truth) 73 | mape_net = metrics.masked_mape_np(y_pred, y_truth) 74 | rmse = metrics.masked_rmse_np(y_pred, y_truth) 75 | mae_sum += mae 76 | mape_net_sum += mape_net 77 | rmse_sum += rmse 78 | mae_list.append(mae) 79 | 80 | mape_net_list.append(mape_net) 81 | rmse_list.append(rmse) 82 | 83 | msg = "Horizon {:02d}, MAE: {:.2f}, RMSE: {:.2f}, MAPE_net: {:.4f}" 84 | if type=='test': 85 | logger.info(msg.format(horizon_i + 1, mae, rmse, mape_net)) 86 | res.append( 87 | { 88 | "MAE": mae, 89 | "RMSE": rmse, 90 | "MAPE_net": mape_net, 91 | } 92 | ) 93 | results['MAE_' + category] = mae_sum / horizon 94 | results['RMSE_' + category] = rmse_sum / horizon 95 | results['MAPE_net_' + category] = mape_net_sum / horizon 96 | 97 | if type=='test': 98 | logger.info('Evaluation_{}_End:'.format(type)) 99 | return results, res 100 | else: 101 | return results 102 | 103 | def main(): 104 | # set seed 105 | torch.manual_seed(args.seed) 106 | np.random.seed(args.seed) 107 | 108 | tb_writer, log_dir = util.setup_tensorboard_writer(os.path.join(args.exp_base, args.runs), comment="", 109 | epochs=args.epochs, lr = args.learning_rate, bn=args.batch_size, nhid=args.nhid, gcn="t" if args.gcn_bool else 'f') 110 | 111 | # load data 112 | logger = util.get_logger(log_dir, __name__, 'info.log', level='INFO') 113 | args.log_dir = log_dir 114 | # device = torch.device(args.device) 115 | device = torch.device( 116 | 'cuda') if torch.cuda.is_available() else torch.device('cpu') 117 | # sensor_ids, sensor_id_to_ind, adj_mx = util.load_adj(args.adjdata,args.adjtype) 118 | adj_mx = util.load_adj(args.adjdata, args.adjtype) 119 | if args.train_type == "od": 120 | dataloader = util.load_dataset(args.data, args.batch_size, args.batch_size, args.batch_size) 121 | else: 122 | dataloader = util.load_dataset_do(args.data, args.batch_size, args.batch_size, args.batch_size) 123 | scaler = dataloader['scaler'] 124 | supports = [torch.tensor(i).to(device) for i in adj_mx] 125 | 126 | print(args) 127 | 128 | if args.randomadj: 129 | adjinit = None 130 | else: 131 | adjinit = supports[0] 132 | 133 | if args.aptonly: 134 | supports = None 135 | 136 | engine = trainer(scaler, args.in_dim, args.seq_length, args.num_nodes, args.out_dim, args.nhid, 137 | args.dropout, args.learning_rate, args.weight_decay, device, 138 | supports, args.gcn_bool, args.addaptadj,adjinit) 139 | 140 | 141 | print("start training...",flush=True) 142 | his_loss =[] 143 | val_time = [] 144 | train_time = [] 145 | 146 | 147 | best_epoch = 0 148 | best_val_net_mape = 1e6 149 | update = {} 150 | for category in ['od', 'do']: 151 | update['val_steady_count_'+category] = 0 152 | update['last_val_mae_'+category] = 1e6 153 | update['last_val_mape_net_'+category] = 1e6 154 | 155 | global_step = 0 156 | for i in range(1,args.epochs+1): 157 | #if i % 10 == 0: 158 | #lr = max(0.000002,args.learning_rate * (0.1 ** (i // 10))) 159 | #for g in engine.optimizer.param_groups: 160 | #g['lr'] = lr 161 | train_loss = [] 162 | t1 = time.time() 163 | dataloader['train_loader'].shuffle() 164 | for iter, (x, y) in enumerate(dataloader['train_loader'].get_iterator()): 165 | trainx = torch.tensor(x, dtype=torch.float, device=device) 166 | trainy = torch.tensor(y, dtype=torch.float, device=device) 167 | 168 | loss, predicts, gt = engine.train(trainx, trainy) 169 | train_loss.append(loss) 170 | tb_writer.add_scalar('train/loss', loss, global_step) 171 | 172 | global_step += 1 173 | 174 | if i % args.print_every == 0: 175 | logger.info(('Epoch:{}').format(i)) 176 | val_result = evaluate(scaler, dataloader, device, engine, logger, 'val', [args.train_type]) 177 | 178 | test_res, res = evaluate(scaler, dataloader, device, engine, logger, 'test', [args.train_type]) 179 | for k, v in test_res.items(): 180 | tb_writer.add_scalar(f'metric_test/{k}', v, i) 181 | 182 | val_category = [args.train_type] 183 | for k, v in val_result.items(): 184 | tb_writer.add_scalar(f'metric_val/{k}', v, i) 185 | for category in val_category: 186 | logger.info('{}:'.format(category)) 187 | logger.info(('val_mae:{}, val_mape_net:{}').format( 188 | val_result['MAE_' + category], 189 | val_result['MAPE_net_' + category])) 190 | if val_result['MAPE_net_' + category] < best_val_net_mape: 191 | best_val_net_mape = val_result['MAPE_net_' + category] 192 | metrics_strs = ['MAE', 'RMSE', 'MAPE_net'] 193 | best_epoch = i 194 | with open(os.path.join(args.log_dir, 'a_res.csv'), 'w') as f: 195 | f.write(f"{args.log_dir}_{best_epoch}\n") 196 | for met in metrics_strs: 197 | f.write(',\n'.join([str(e[met]) for e in res] + [',\n'])) 198 | torch.save(engine.model.state_dict(), os.path.join(args.log_dir, "best_model.pth")) 199 | if update['last_val_mae_' + category] > val_result['MAE_' + category]: 200 | logger.info('val_mae decreased from {} to {}'.format( 201 | update['last_val_mae_' + category], 202 | val_result['MAE_' + category])) 203 | update['last_val_mae_' + category] = val_result['MAE_' + category] 204 | 205 | if update['last_val_mape_net_' + category] > val_result['MAPE_net_' + category]: 206 | logger.info('val_mape_net decreased from {} to {}'.format( 207 | update['last_val_mape_net_' + category], 208 | val_result['MAPE_net_' + category])) 209 | update['last_val_mape_net_' + category] = val_result['MAPE_net_' + category] 210 | 211 | torch.save(engine.model.state_dict(), os.path.join(args.log_dir, "epoch_"+str(i)+".pth")) 212 | print("Average Training Time: {:.4f} secs/epoch".format(np.mean(train_time))) 213 | print("Average Inference Time: {:.4f} secs".format(np.mean(val_time))) 214 | 215 | print("Training finished") 216 | 217 | 218 | if __name__ == "__main__": 219 | t1 = time.time() 220 | main() 221 | t2 = time.time() 222 | print("Total time spent: {:.4f}".format(t2-t1)) 223 | -------------------------------------------------------------------------------- /models/OD_Net.py: -------------------------------------------------------------------------------- 1 | from torch_geometric import nn as gnn 2 | from torch import nn 3 | from torch.nn import functional as F 4 | from torch.nn import init, Parameter 5 | import torch 6 | import random 7 | import math 8 | import sys 9 | import os 10 | import copy 11 | # from lib.utils_unfinished import softmax 12 | from torch.nn import functional as F 13 | sys.path.insert(0, os.path.abspath('..')) 14 | 15 | from models.GGRUCell import GGRUCell 16 | 17 | class ODNet(torch.nn.Module): 18 | 19 | def __init__(self, cfg, logger): 20 | super(ODNet, self).__init__() 21 | self.logger = logger 22 | self.cfg = cfg 23 | 24 | self.num_nodes = cfg['model']['num_nodes'] 25 | self.num_output_dim = cfg['model']['output_dim'] 26 | self.num_units = cfg['model']['rnn_units'] 27 | self.num_finished_input_dim = cfg['model']['input_dim'] 28 | self.num_unfinished_input_dim = cfg['model']['input_dim'] 29 | self.num_rnn_layers = cfg['model']['num_rnn_layers'] 30 | self.seq_len = cfg['model']['seq_len'] 31 | self.horizon = cfg['model']['horizon'] 32 | self.num_relations = cfg['model'].get('num_relations', 1) 33 | self.K = cfg['model'].get('K', 2) 34 | self.num_bases = cfg['model'].get('num_bases', 1) 35 | 36 | self.dropout_type = cfg['model'].get('dropout_type', None) 37 | self.dropout_prob = cfg['model'].get('dropout_prob', 0.0) 38 | 39 | self.global_fusion = cfg['model'].get('global_fusion', False) 40 | 41 | self.encoder_first_finished_cells = GGRUCell(self.num_finished_input_dim, 42 | self.num_units, 43 | self.dropout_type, 44 | self.dropout_prob, 45 | self.num_relations, 46 | num_bases=self.num_bases, 47 | K=self.K, 48 | num_nodes=self.num_nodes, 49 | global_fusion=self.global_fusion) 50 | self.encoder_first_unfinished_cells = GGRUCell(self.num_unfinished_input_dim, 51 | self.num_units, 52 | self.dropout_type, 53 | self.dropout_prob, 54 | self.num_relations, 55 | num_bases=self.num_bases, 56 | K=self.K, 57 | num_nodes=self.num_nodes, 58 | global_fusion=self.global_fusion) 59 | self.encoder_first_short_his_cells = GGRUCell(self.num_unfinished_input_dim, 60 | self.num_units, 61 | self.dropout_type, 62 | self.dropout_prob, 63 | self.num_relations, 64 | num_bases=self.num_bases, 65 | K=self.K, 66 | num_nodes=self.num_nodes, 67 | global_fusion=self.global_fusion) 68 | 69 | self.unfinished_output_layer = nn.Conv1d(in_channels=self.num_units*2, 70 | out_channels=self.num_units, 71 | kernel_size=1) 72 | self.unfinished_hidden_layer = nn.Conv1d(in_channels=self.num_units*2, 73 | out_channels=self.num_units, 74 | kernel_size=1) 75 | 76 | self.encoder_second_cells = nn.ModuleList([GGRUCell(self.num_units, 77 | self.num_units, 78 | self.dropout_type, 79 | self.dropout_prob, 80 | self.num_relations, 81 | num_bases=self.num_bases, 82 | K=self.K, 83 | num_nodes=self.num_nodes, 84 | global_fusion=self.global_fusion) 85 | for _ in range(self.num_rnn_layers - 1)]) 86 | self.decoder_first_cells = GGRUCell(self.num_finished_input_dim, 87 | self.num_units, 88 | self.dropout_type, 89 | self.dropout_prob, 90 | self.num_relations, 91 | num_bases=self.num_bases, 92 | K=self.K, 93 | num_nodes=self.num_nodes, 94 | global_fusion=self.global_fusion) 95 | self.decoder_second_cells = nn.ModuleList([GGRUCell(self.num_units, 96 | self.num_units, 97 | self.dropout_type, 98 | self.dropout_prob, 99 | self.num_relations, 100 | self.K, 101 | num_nodes=self.num_nodes, 102 | global_fusion=self.global_fusion) 103 | for _ in range(self.num_rnn_layers - 1)]) 104 | 105 | self.output_type = cfg['model'].get('output_type', 'fc') 106 | if self.output_type == 'fc': 107 | self.output_layer = nn.Linear(self.num_units, self.num_output_dim) 108 | 109 | def encoder_first_layer(self, 110 | batch, 111 | finished_hidden, 112 | long_his_hidden, 113 | short_his_hidden, 114 | edge_index, 115 | edge_attr=None): 116 | finished_out, finished_hidden = self.encoder_first_finished_cells(inputs=batch.x_od, 117 | edge_index=edge_index, 118 | edge_attr=edge_attr, 119 | hidden=finished_hidden) 120 | enc_first_hidden = finished_hidden 121 | enc_first_out = finished_out 122 | 123 | long_his_out, long_his_hidden = self.encoder_first_unfinished_cells(inputs=batch.history, 124 | edge_index=edge_index, 125 | edge_attr=edge_attr, 126 | hidden=long_his_hidden) 127 | 128 | short_his_out, short_his_hidden = self.encoder_first_short_his_cells(inputs=batch.yesterday, 129 | edge_index=edge_index, 130 | edge_attr=edge_attr, 131 | hidden=short_his_hidden) 132 | 133 | hidden_fusion = torch.cat([long_his_hidden, short_his_hidden], -1).view(-1, self.num_units * 2, self.num_nodes) 134 | 135 | long_his_weight = torch.sigmoid(self.unfinished_hidden_layer(hidden_fusion)).view(-1, self.num_units) 136 | short_his_weight = torch.sigmoid(self.unfinished_output_layer(hidden_fusion)).view(-1, self.num_units) 137 | 138 | unfinished_hidden = long_his_weight * long_his_hidden + short_his_weight * short_his_hidden 139 | unfinished_out = long_his_weight * long_his_out + short_his_weight * short_his_out 140 | 141 | enc_first_out = finished_out + unfinished_out 142 | enc_first_hidden = enc_first_hidden + unfinished_hidden 143 | 144 | return enc_first_out, finished_hidden, long_his_hidden, short_his_hidden, enc_first_hidden 145 | 146 | def encoder_second_layer(self, 147 | index, 148 | first_out, 149 | enc_second_hidden, 150 | edge_index, 151 | edge_attr): 152 | enc_second_out, enc_second_hidden = self.encoder_second_cells[index](inputs=first_out, 153 | hidden=enc_second_hidden, 154 | edge_index=edge_index, 155 | edge_attr=edge_attr) 156 | return enc_second_out, enc_second_hidden 157 | 158 | def decoder_first_layer(self, 159 | decoder_input, 160 | dec_first_hidden, 161 | edge_index, 162 | edge_attr=None): 163 | dec_first_out, dec_first_hidden = self.decoder_first_cells(inputs=decoder_input, 164 | hidden=dec_first_hidden, 165 | edge_index=edge_index, 166 | edge_attr=edge_attr) 167 | return dec_first_out, dec_first_hidden 168 | 169 | def decoder_second_layer(self, 170 | index, 171 | decoder_first_out, 172 | dec_second_hidden, 173 | edge_index, 174 | edge_attr=None): 175 | dec_second_out, dec_second_hidden = self.decoder_second_cells[index](inputs=decoder_first_out, 176 | hidden=dec_second_hidden, 177 | edge_index=edge_index, 178 | edge_attr=edge_attr) 179 | return dec_second_out, dec_second_hidden 180 | 181 | 182 | 183 | -------------------------------------------------------------------------------- /evaluate.py: -------------------------------------------------------------------------------- 1 | # encoding:utf-8 2 | import random 3 | import argparse 4 | import time 5 | import yaml 6 | import numpy as np 7 | import torch 8 | import os 9 | 10 | from torch import nn 11 | from torch.optim.lr_scheduler import MultiStepLR 12 | from torch.nn.init import xavier_uniform_ 13 | from lib import utils_HIAM as utils 14 | from lib.utils_HIAM import collate_wrapper 15 | from lib import metrics 16 | from models.Net import Net 17 | try: 18 | from yaml import CLoader as Loader, CDumper as Dumper 19 | except ImportError: 20 | from yaml import Loader, Dumper 21 | seed = 1234 22 | random.seed(seed) 23 | np.random.seed(seed) 24 | torch.manual_seed(seed) 25 | parser = argparse.ArgumentParser() 26 | parser.add_argument('--config_filename', 27 | default=None, 28 | type=str, 29 | help='Configuration filename for restoring the model.') 30 | args = parser.parse_args() 31 | 32 | def read_cfg_file(filename): 33 | with open(filename, 'r') as ymlfile: 34 | cfg = yaml.load(ymlfile, Loader=Loader) 35 | return cfg 36 | 37 | def run_model(model, data_iterator, edge_index, edge_attr, device, seq_len, horizon, output_dim): 38 | """ 39 | return a list of (horizon_i, batch_size, num_nodes, output_dim) 40 | """ 41 | # while evaluation, we need model.eval and torch.no_grad 42 | model.eval() 43 | y_od_pred_list = [] 44 | y_do_pred_list = [] 45 | for _, (x_od, y_od, x_do, y_do, unfinished, history, yesterday, xtime, ytime) in enumerate(data_iterator): 46 | y_od = y_od[..., :output_dim] 47 | y_do = y_do[..., :output_dim] 48 | sequences, sequences_y, y_od, y_do = collate_wrapper(x_od=x_od, y_od=y_od, x_do=x_do, y_do=y_do, unfinished=unfinished, history=history, yesterday=yesterday, 49 | edge_index=edge_index, 50 | edge_attr=edge_attr, 51 | device=device, 52 | seq_len=seq_len, 53 | horizon=horizon) 54 | # (T, N, num_nodes, num_out_channels) 55 | with torch.no_grad(): 56 | y_od_pred, y_do_pred = model(sequences, sequences_y) 57 | if y_od_pred is not None: 58 | y_od_pred_list.append(y_od_pred.cpu().numpy()) 59 | if y_do_pred is not None: 60 | y_do_pred_list.append(y_do_pred.cpu().numpy()) 61 | return y_od_pred_list, y_do_pred_list 62 | 63 | 64 | def evaluate(model, 65 | dataset, 66 | dataset_type, 67 | edge_index, 68 | edge_attr, 69 | device, 70 | seq_Len, 71 | horizon, 72 | output_dim, 73 | logger, 74 | detail=True, 75 | cfg=None, 76 | format_result=False): 77 | if detail: 78 | logger.info('Evaluation_{}_Begin:'.format(dataset_type)) 79 | 80 | y_od_preds, y_do_preds = run_model( 81 | model, 82 | data_iterator=dataset['{}_loader'.format(dataset_type)].get_iterator(), 83 | edge_index=edge_index, 84 | edge_attr=edge_attr, 85 | device=device, 86 | seq_len=seq_Len, 87 | horizon=horizon, 88 | output_dim=output_dim) 89 | 90 | evaluate_category = [] 91 | if len(y_od_preds) > 0: 92 | evaluate_category.append("od") 93 | if len(y_do_preds) > 0: 94 | evaluate_category.append("do") 95 | results = {} 96 | for category in evaluate_category: 97 | if category == 'od': 98 | y_preds = y_od_preds 99 | scaler = dataset['scaler'] 100 | gt = dataset['y_{}'.format(dataset_type)] 101 | else: 102 | y_preds = y_do_preds 103 | scaler = dataset['do_scaler'] 104 | # scaler = dataset['scaler'] 105 | gt = dataset['do_y_{}'.format(dataset_type)] 106 | y_preds = np.concatenate(y_preds, axis=0) # concat in batch_size dim. 107 | mae_list = [] 108 | mape_net_list = [] 109 | rmse_list = [] 110 | mae_sum = 0 111 | 112 | mape_net_sum = 0 113 | rmse_sum = 0 114 | # horizon = dataset['y_{}'.format(dataset_type)].shape[1] 115 | logger.info("{}:".format(category)) 116 | horizon = cfg['model']['horizon'] 117 | for horizon_i in range(horizon): 118 | y_truth = scaler.inverse_transform( 119 | gt[:, horizon_i, :, :output_dim]) 120 | 121 | y_pred = scaler.inverse_transform( 122 | y_preds[:y_truth.shape[0], horizon_i, :, :output_dim]) 123 | y_pred[y_pred < 0] = 0 124 | mae = metrics.masked_mae_np(y_pred, y_truth) 125 | mape_net = metrics.masked_mape_np(y_pred, y_truth) 126 | rmse = metrics.masked_rmse_np(y_pred, y_truth) 127 | mae_sum += mae 128 | mape_net_sum += mape_net 129 | rmse_sum += rmse 130 | mae_list.append(mae) 131 | 132 | mape_net_list.append(mape_net) 133 | rmse_list.append(rmse) 134 | 135 | msg = "Horizon {:02d}, MAE: {:.2f}, RMSE: {:.2f}, MAPE_net: {:.4f}" 136 | if detail: 137 | logger.info(msg.format(horizon_i + 1, mae, rmse, mape_net)) 138 | results['MAE_' + category] = mae_sum / horizon 139 | results['RMSE_' + category] = rmse_sum / horizon 140 | results['MAPE_net_' + category] = mape_net_sum / horizon 141 | if detail: 142 | logger.info('Evaluation_{}_End:'.format(dataset_type)) 143 | if format_result: 144 | for i in range(len(mae_list)): 145 | print('{:.2f}'.format(mae_list[i])) 146 | print('{:.2f}'.format(rmse_list[i])) 147 | print('{:.2f}%'.format(mape_net_list[i] * 100)) 148 | print() 149 | else: 150 | # return mae_sum / horizon, rmse_sum / horizon, mape_sta_sum / horizon, mape_pair_sum / horizon, mape_net_sum/ horizon, mape_distribution_sum / horizon 151 | return results 152 | 153 | class StepLR2(MultiStepLR): 154 | """StepLR with min_lr""" 155 | def __init__(self, 156 | optimizer, 157 | milestones, 158 | gamma=0.1, 159 | last_epoch=-1, 160 | min_lr=2.0e-6): 161 | self.optimizer = optimizer 162 | self.milestones = milestones 163 | self.gamma = gamma 164 | self.last_epoch = last_epoch 165 | self.min_lr = min_lr 166 | super(StepLR2, self).__init__(optimizer, milestones, gamma) 167 | 168 | def get_lr(self): 169 | lr_candidate = super(StepLR2, self).get_lr() 170 | if isinstance(lr_candidate, list): 171 | for i in range(len(lr_candidate)): 172 | lr_candidate[i] = max(self.min_lr, lr_candidate[i]) 173 | 174 | else: 175 | lr_candidate = max(self.min_lr, lr_candidate) 176 | 177 | return lr_candidate 178 | 179 | def _get_log_dir(kwargs): 180 | log_dir = kwargs['train'].get('log_dir') 181 | if log_dir is None: 182 | batch_size = kwargs['data'].get('batch_size') 183 | learning_rate = kwargs['train'].get('base_lr') 184 | num_rnn_layers = kwargs['model'].get('num_rnn_layers') 185 | rnn_units = kwargs['model'].get('rnn_units') 186 | structure = '-'.join(['%d' % rnn_units for _ in range(num_rnn_layers)]) 187 | 188 | # name of dir for saving log 189 | run_id = 'HIAM_%s_lr%g_bs%d_%s/' % ( 190 | structure, 191 | learning_rate, 192 | batch_size, 193 | time.strftime('%m%d%H%M%S')) 194 | base_dir = kwargs.get('base_dir') 195 | log_dir = os.path.join(base_dir, run_id) 196 | if not os.path.exists(log_dir): 197 | os.makedirs(log_dir) 198 | return log_dir 199 | 200 | 201 | def init_weights(m): 202 | classname = m.__class__.__name__ # 2 203 | if classname.find('Conv') != -1 and classname.find('RGCN') == -1: 204 | xavier_uniform_(m.weight.data) 205 | if type(m) == nn.Linear: 206 | xavier_uniform_(m.weight.data) 207 | #xavier_uniform_(m.bias.data) 208 | 209 | def toDevice(datalist, device): 210 | for i in range(len(datalist)): 211 | datalist[i] = datalist[i].to(device) 212 | return datalist 213 | 214 | def main(args): 215 | cfg = read_cfg_file(args.config_filename) 216 | log_dir = _get_log_dir(cfg) 217 | log_level = cfg.get('log_level', 'INFO') 218 | 219 | logger = utils.get_logger(log_dir, __name__, 'info.log', level=log_level) 220 | 221 | device = torch.device( 222 | 'cuda') if torch.cuda.is_available() else torch.device('cpu') 223 | # all edge_index in same dataset is same 224 | # edge_index = adjacency_to_edge_index(adj_mx) # alreay added self-loop 225 | logger.info(cfg) 226 | batch_size = cfg['data']['batch_size'] 227 | seq_len = cfg['model']['seq_len'] 228 | horizon = cfg['model']['horizon'] 229 | # edge_index = utils.load_pickle(cfg['data']['edge_index_pkl_filename']) 230 | 231 | adj_mx_list = [] 232 | graph_pkl_filename = cfg['data']['graph_pkl_filename'] 233 | 234 | if not isinstance(graph_pkl_filename, list): 235 | graph_pkl_filename = [graph_pkl_filename] 236 | 237 | src = [] 238 | dst = [] 239 | for g in graph_pkl_filename: 240 | adj_mx = utils.load_graph_data(g) 241 | for i in range(len(adj_mx)): # 构建邻接矩阵 242 | adj_mx[i, i] = 0 243 | adj_mx_list.append(adj_mx) 244 | 245 | adj_mx = np.stack(adj_mx_list, axis=-1) 246 | print("adj_mx:", adj_mx.shape) 247 | if cfg['model'].get('norm', False): 248 | print('row normalization') 249 | adj_mx = adj_mx / (adj_mx.sum(axis=0) + 1e-18) # 归一化 250 | src, dst = adj_mx.sum(axis=-1).nonzero() 251 | print("src, dst:", src.shape, dst.shape) 252 | edge_index = torch.tensor([src, dst], dtype=torch.long, device=device) 253 | edge_attr = torch.tensor(adj_mx[adj_mx.sum(axis=-1) != 0], 254 | dtype=torch.float, 255 | device=device) 256 | print("train, edge:", edge_index.shape, edge_attr.shape) 257 | output_dim = cfg['model']['output_dim'] 258 | for i in range(adj_mx.shape[-1]): 259 | logger.info(adj_mx[..., i]) 260 | 261 | dataset = utils.load_dataset(**cfg['data'], scaler_axis=(0, 1, 2, 3)) 262 | for k, v in dataset.items(): 263 | if hasattr(v, 'shape'): 264 | logger.info((k, v.shape)) 265 | 266 | model = Net(cfg, logger).to(device) 267 | state = torch.load(cfg['model']['save_path']) 268 | model.load_state_dict(state, strict=False) 269 | evaluate(model=model, 270 | dataset=dataset, 271 | dataset_type='test', 272 | edge_index=edge_index, 273 | edge_attr=edge_attr, 274 | device=device, 275 | seq_Len=seq_len, 276 | horizon=horizon, 277 | output_dim=output_dim, 278 | logger=logger, 279 | cfg=cfg) 280 | 281 | if __name__ == "__main__": 282 | main(args) 283 | -------------------------------------------------------------------------------- /SOTA/DCRNN/model/dcrnn_supervisor.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | 4 | import numpy as np 5 | import torch 6 | # from torch.utils.tensorboard import SummaryWriter 7 | 8 | from lib import utils_HIAM, metrics 9 | from .dcrnn_model import DCRNNModel 10 | 11 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 12 | 13 | 14 | class DCRNNSupervisor: 15 | def __init__(self, adj_mx, **kwargs): 16 | self._kwargs = kwargs 17 | self._data_kwargs = kwargs.get('data') 18 | self._model_kwargs = kwargs.get('model') 19 | self._train_kwargs = kwargs.get('train') 20 | 21 | self.max_grad_norm = self._train_kwargs.get('max_grad_norm', 1.) 22 | 23 | # logging. 24 | self._log_dir = self._get_log_dir(kwargs) 25 | # self._writer = SummaryWriter('runs/' + self._log_dir) 26 | 27 | log_level = self._kwargs.get('log_level', 'INFO') 28 | self._logger = utils_HIAM.get_logger(self._log_dir, __name__, 'info.log', level=log_level) 29 | 30 | # data set 31 | self._data = utils_HIAM.load_dataset(**self._data_kwargs) 32 | self.standard_scaler = self._data['scaler'] 33 | 34 | self.num_nodes = int(self._model_kwargs.get('num_nodes', 1)) 35 | self.input_dim = int(self._model_kwargs.get('input_dim', 1)) 36 | self.seq_len = int(self._model_kwargs.get('seq_len')) # for the encoder 37 | self.output_dim = int(self._model_kwargs.get('output_dim', 1)) 38 | self.use_curriculum_learning = bool( 39 | self._model_kwargs.get('use_curriculum_learning', False)) 40 | self.horizon = int(self._model_kwargs.get('horizon', 1)) # for the decoder 41 | 42 | # setup model 43 | dcrnn_model = DCRNNModel(adj_mx, self._logger, **self._model_kwargs) 44 | self.dcrnn_model = dcrnn_model.cuda() if torch.cuda.is_available() else dcrnn_model 45 | self._logger.info("Model created") 46 | 47 | self._epoch_num = self._train_kwargs.get('epoch', 0) 48 | if self._epoch_num > 0: 49 | self.load_model() 50 | 51 | @staticmethod 52 | def _get_log_dir(kwargs): 53 | log_dir = kwargs['train'].get('log_dir') 54 | if log_dir is None: 55 | batch_size = kwargs['data'].get('batch_size') 56 | learning_rate = kwargs['train'].get('base_lr') 57 | max_diffusion_step = kwargs['model'].get('max_diffusion_step') 58 | num_rnn_layers = kwargs['model'].get('num_rnn_layers') 59 | rnn_units = kwargs['model'].get('rnn_units') 60 | structure = '-'.join( 61 | ['%d' % rnn_units for _ in range(num_rnn_layers)]) 62 | horizon = kwargs['model'].get('horizon') 63 | filter_type = kwargs['model'].get('filter_type') 64 | filter_type_abbr = 'L' 65 | if filter_type == 'random_walk': 66 | filter_type_abbr = 'R' 67 | elif filter_type == 'dual_random_walk': 68 | filter_type_abbr = 'DR' 69 | run_id = 'dcrnn_%s_%d_h_%d_%s_lr_%g_bs_%d_%s/' % ( 70 | filter_type_abbr, max_diffusion_step, horizon, 71 | structure, learning_rate, batch_size, 72 | time.strftime('%m%d%H%M%S')) 73 | base_dir = kwargs.get('base_dir') 74 | log_dir = os.path.join(base_dir, run_id) 75 | if not os.path.exists(log_dir): 76 | os.makedirs(log_dir) 77 | return log_dir 78 | 79 | def save_model(self, epoch): 80 | if not os.path.exists('models/'): 81 | os.makedirs('models/') 82 | 83 | config = dict(self._kwargs) 84 | config['model_state_dict'] = self.dcrnn_model.state_dict() 85 | config['epoch'] = epoch 86 | torch.save(config, 'models/epo%d.pt' % epoch) 87 | self._logger.info("Saved model at {}".format(epoch)) 88 | return 'models/epo%d.pt' % epoch 89 | 90 | def load_model(self): 91 | self._setup_graph() 92 | assert os.path.exists('models/epo%d.pt' % self._epoch_num), 'Weights at epoch %d not found' % self._epoch_num 93 | checkpoint = torch.load('models/epo%d.pt' % self._epoch_num, map_location='cpu') 94 | self.dcrnn_model.load_state_dict(checkpoint['model_state_dict']) 95 | self._logger.info("Loaded model at {}".format(self._epoch_num)) 96 | 97 | def _setup_graph(self): 98 | with torch.no_grad(): 99 | self.dcrnn_model = self.dcrnn_model.eval() 100 | 101 | val_iterator = self._data['val_loader'].get_iterator() 102 | 103 | for _, (x, y) in enumerate(val_iterator): 104 | x, y = self._prepare_data(x, y) 105 | output = self.dcrnn_model(x) 106 | 107 | output = output.view(self.horizon, -1, self.num_nodes, self.output_dim).transpose(0, 1) 108 | y = y.view(self.horizon, -1, self.num_nodes, self.output_dim).transpose(0, 1) 109 | break 110 | 111 | def train(self, **kwargs): 112 | kwargs.update(self._train_kwargs) 113 | return self._train(**kwargs) 114 | 115 | def evaluate(self, dataset='val', batches_seen=0): 116 | """ 117 | Computes mean L1Loss 118 | :return: mean L1Loss 119 | """ 120 | with torch.no_grad(): 121 | self.dcrnn_model = self.dcrnn_model.eval() 122 | 123 | val_iterator = self._data['{}_loader'.format(dataset)].get_iterator() 124 | # losses = [] 125 | mae_sum = 0 126 | mape_sum = 0 127 | 128 | y_preds = [] 129 | y_truths = [] 130 | 131 | for _, (x, y) in enumerate(val_iterator): 132 | x, y = self._prepare_data(x, y) 133 | output = self.dcrnn_model(x) 134 | 135 | output = output.view(self.horizon, -1, self.num_nodes, self.output_dim).transpose(0, 1) 136 | y = y.view(self.horizon, -1, self.num_nodes, self.output_dim).transpose(0, 1) 137 | y_preds.append(output.cpu()) 138 | y_truths.append(y.cpu()) 139 | 140 | y_preds = np.concatenate(y_preds, axis=0) 141 | y_truths = np.concatenate(y_truths, axis=0) 142 | 143 | scaler = self._data['scaler'] 144 | for horizon_i in range(self.horizon): 145 | y_truth = scaler.inverse_transform(y_truths[:, horizon_i, :, :]) 146 | y_pred = scaler.inverse_transform(y_preds[:, horizon_i, :, :]) 147 | 148 | mae = metrics.masked_mae_np(y_pred, y_truth) 149 | rmse = metrics.masked_rmse_np(y_pred, y_truth) 150 | mape_net = metrics.masked_mape_np(y_pred, y_truth) 151 | 152 | mae_sum += mae 153 | mape_sum += mape_net 154 | 155 | if dataset == 'test': 156 | self._logger.info( 157 | "Horizon {:02d}, MAE: {:.2f}, RMSE: {:.2f}, MAPE_net: {:.4f}".format( 158 | horizon_i + 1, mae, rmse, mape_net) 159 | ) 160 | mean_mae = mae_sum / self.horizon * 1.0 161 | mean_mape = mape_sum / self.horizon * 1.0 162 | return mean_mae, mean_mape 163 | 164 | def _train(self, base_lr, 165 | steps, patience=50, epochs=100, lr_decay_ratio=0.1, log_every=1, save_model=1, 166 | test_every_n_epochs=10, epsilon=1e-8, **kwargs): 167 | # steps is used in learning rate - will see if need to use it? 168 | min_val_loss = float('inf') 169 | min_val_mape_loss = float('inf') 170 | wait = 0 171 | optimizer = torch.optim.Adam(self.dcrnn_model.parameters(), lr=base_lr, eps=epsilon) 172 | 173 | lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=steps, 174 | gamma=lr_decay_ratio) 175 | 176 | self._logger.info('Start training ...') 177 | 178 | # this will fail if model is loaded with a changed batch_size 179 | num_batches = self._data['train_loader'].num_batch 180 | self._logger.info("num_batches:{}".format(num_batches)) 181 | 182 | batches_seen = num_batches * self._epoch_num 183 | 184 | for epoch_num in range(self._epoch_num, epochs): 185 | 186 | self.dcrnn_model = self.dcrnn_model.train() 187 | 188 | train_iterator = self._data['train_loader'].get_iterator() 189 | losses = [] 190 | 191 | start_time = time.time() 192 | 193 | for _, (x, y) in enumerate(train_iterator): 194 | optimizer.zero_grad() 195 | 196 | x, y = self._prepare_data(x, y) 197 | 198 | output = self.dcrnn_model(x, y, batches_seen) 199 | 200 | if batches_seen == 0: 201 | # this is a workaround to accommodate dynamically registered parameters in DCGRUCell 202 | optimizer = torch.optim.Adam(self.dcrnn_model.parameters(), lr=base_lr, eps=epsilon) 203 | 204 | output = output.view(self.horizon, -1, self.num_nodes, self.output_dim).transpose(0, 1) 205 | y = y.view(self.horizon, -1, self.num_nodes, self.output_dim).transpose(0, 1) 206 | loss = self._compute_loss(y, output) 207 | 208 | self._logger.debug(loss.item()) 209 | 210 | losses.append(loss.item()) 211 | 212 | batches_seen += 1 213 | loss.backward() 214 | 215 | # gradient clipping - this does it in place 216 | torch.nn.utils.clip_grad_norm_(self.dcrnn_model.parameters(), self.max_grad_norm) 217 | 218 | optimizer.step() 219 | self._logger.info("epoch complete") 220 | lr_scheduler.step() 221 | self._logger.info("evaluating now!") 222 | 223 | val_loss, val_mape_loss = self.evaluate(dataset='val', batches_seen=batches_seen) 224 | 225 | end_time = time.time() 226 | 227 | message = 'Epoch [{}/{}] ({}) train_mae: {:.4f}, val_mae: {:.4f}, val_mape: {:.4f}, lr: {:.6f}, ' \ 228 | '{:.1f}s'.format(epoch_num, epochs, batches_seen, 229 | np.mean(losses), val_loss, val_mape_loss, lr_scheduler.get_lr()[0], 230 | (end_time - start_time)) 231 | self._logger.info(message) 232 | 233 | if val_loss < min_val_loss: 234 | wait = 0 235 | if save_model: 236 | model_file_name = self.save_model(epoch_num) 237 | self._logger.info( 238 | 'Val MAE decrease from {:.4f} to {:.4f}, ' 239 | 'saving to {}'.format(min_val_loss, val_loss, model_file_name)) 240 | min_val_loss = val_loss 241 | 242 | if val_mape_loss < min_val_mape_loss: 243 | wait = 0 244 | if save_model: 245 | model_file_name = self.save_model(epoch_num) 246 | self._logger.info( 247 | 'Val MAPE decrease from {:.4f} to {:.4f}, ' 248 | 'saving to {}'.format(min_val_mape_loss, val_mape_loss, model_file_name)) 249 | min_val_mape_loss = val_mape_loss 250 | 251 | _, _ = self.evaluate(dataset='test', batches_seen=batches_seen) 252 | 253 | if val_loss >= min_val_loss: 254 | wait += 1 255 | if wait == patience: 256 | self._logger.warning('Early stopping at epoch: %d' % epoch_num) 257 | break 258 | 259 | def _prepare_data(self, x, y): 260 | x, y = self._get_x_y(x, y) 261 | x, y = self._get_x_y_in_correct_dims(x, y) 262 | return x.to(device), y.to(device) 263 | 264 | def _get_x_y(self, x, y): 265 | """ 266 | :param x: shape (batch_size, seq_len, num_sensor, input_dim) 267 | :param y: shape (batch_size, horizon, num_sensor, input_dim) 268 | :returns x shape (seq_len, batch_size, num_sensor, input_dim) 269 | y shape (horizon, batch_size, num_sensor, input_dim) 270 | """ 271 | x = torch.from_numpy(x).float() 272 | y = torch.from_numpy(y).float() 273 | self._logger.debug("X: {}".format(x.size())) 274 | self._logger.debug("y: {}".format(y.size())) 275 | x = x.permute(1, 0, 2, 3) 276 | y = y.permute(1, 0, 2, 3) 277 | return x, y 278 | 279 | def _get_x_y_in_correct_dims(self, x, y): 280 | """ 281 | :param x: shape (seq_len, batch_size, num_sensor, input_dim) 282 | :param y: shape (horizon, batch_size, num_sensor, input_dim) 283 | :return: x: shape (seq_len, batch_size, num_sensor * input_dim) 284 | y: shape (horizon, batch_size, num_sensor * output_dim) 285 | """ 286 | batch_size = x.size(1) 287 | x = x.view(self.seq_len, batch_size, self.num_nodes * self.input_dim) 288 | y = y[..., :self.output_dim].view(self.horizon, batch_size, 289 | self.num_nodes * self.output_dim) 290 | return x, y 291 | 292 | def _compute_loss(self, y_true, y_predicted): 293 | y_true = self.standard_scaler.inverse_transform(y_true) 294 | y_predicted = self.standard_scaler.inverse_transform(y_predicted) 295 | loss = torch.nn.L1Loss(reduction='mean') 296 | return loss(y_predicted, y_true) 297 | -------------------------------------------------------------------------------- /models/Net.py: -------------------------------------------------------------------------------- 1 | from torch_geometric import nn as gnn 2 | from torch import nn 3 | import torch 4 | import random 5 | import math 6 | import sys 7 | import os 8 | sys.path.insert(0, os.path.abspath('..')) 9 | 10 | from models.OD_Net import ODNet 11 | from models.DO_Net import DONet 12 | from models.DualInfoTransformer import DualInfoTransformer 13 | 14 | 15 | class Net(torch.nn.Module): 16 | 17 | def __init__(self, cfg, logger): 18 | super(Net, self).__init__() 19 | self.logger = logger 20 | self.cfg = cfg 21 | 22 | self.OD = ODNet(cfg, logger) 23 | self.DO = DONet(cfg, logger) 24 | 25 | self.num_nodes = cfg['model']['num_nodes'] 26 | self.num_output_dim = cfg['model']['output_dim'] 27 | self.num_units = cfg['model']['rnn_units'] 28 | self.num_finished_input_dim = cfg['model']['input_dim'] 29 | self.num_unfinished_input_dim = cfg['model']['input_dim'] 30 | self.num_rnn_layers = cfg['model']['num_rnn_layers'] 31 | 32 | self.seq_len = cfg['model']['seq_len'] 33 | self.horizon = cfg['model']['horizon'] 34 | self.head = cfg['model'].get('head', 4) 35 | self.d_channel = cfg['model'].get('channel', 512) 36 | 37 | self.use_curriculum_learning = self.cfg['model']['use_curriculum_learning'] 38 | self.cl_decay_steps = torch.FloatTensor(data=[self.cfg['model']['cl_decay_steps']]) 39 | self.use_input = cfg['model'].get('use_input', True) 40 | 41 | self.mediate_activation = nn.PReLU(self.num_units) 42 | 43 | self.global_step = 0 44 | 45 | self.encoder_first_interact = DualInfoTransformer( 46 | h=self.head, 47 | d_nodes=self.num_nodes, 48 | d_model=self.num_units, 49 | d_channel=self.d_channel) 50 | self.decoder_first_interact = DualInfoTransformer( 51 | h=self.head, 52 | d_nodes=self.num_nodes, 53 | d_model=self.num_units, 54 | d_channel=self.d_channel) 55 | self.encoder_second_interact = nn.ModuleList([DualInfoTransformer( 56 | h=self.head, 57 | d_nodes=self.num_nodes, 58 | d_model=self.num_units, 59 | d_channel=self.d_channel) 60 | for _ in range(self.num_rnn_layers - 1)]) 61 | self.decoder_second_interact = nn.ModuleList([DualInfoTransformer( 62 | h=self.head, 63 | d_nodes=self.num_nodes, 64 | d_model=self.num_units, 65 | d_channel=self.d_channel) 66 | for _ in range(self.num_rnn_layers - 1)]) 67 | 68 | 69 | @staticmethod 70 | def inverse_sigmoid_scheduler_sampling(step, k): 71 | """TODO: Docstring for linear_scheduler_sampling. 72 | :returns: TODO 73 | 74 | """ 75 | try: 76 | return k / (k + math.exp(step / k)) 77 | except OverflowError: 78 | return float('inf') 79 | 80 | def encoder_od_do(self, 81 | sequences, 82 | edge_index, 83 | edge_attr=None): 84 | """ 85 | Encodes input into hidden state on one branch for T steps. 86 | 87 | Return: hidden state on one branch.""" 88 | enc_hiddens_od = [None] * self.num_rnn_layers 89 | enc_hiddens_do = [None] * self.num_rnn_layers 90 | 91 | finished_hidden_od = None 92 | long_his_hidden_od = None 93 | short_his_hidden_od = None 94 | 95 | for t, batch in enumerate(sequences): 96 | encoder_first_out_od, finished_hidden_od, \ 97 | long_his_hidden_od, short_his_hidden_od, \ 98 | enc_first_hidden_od = self.OD.encoder_first_layer(batch, 99 | finished_hidden_od, 100 | long_his_hidden_od, 101 | short_his_hidden_od, 102 | edge_index, 103 | edge_attr) 104 | 105 | enc_first_out_do, enc_first_hidden_do = self.DO.encoder_first_layer(batch, 106 | enc_hiddens_do[0], 107 | edge_index, 108 | edge_attr) 109 | 110 | enc_first_interact_info_od, enc_first_interact_info_do = self.encoder_first_interact( 111 | enc_first_hidden_od, 112 | enc_first_hidden_do) 113 | 114 | enc_hiddens_od[0] = enc_first_hidden_od + enc_first_interact_info_od 115 | enc_hiddens_do[0] = enc_first_hidden_do + enc_first_interact_info_do 116 | 117 | enc_mid_out_od = encoder_first_out_od + enc_first_interact_info_od 118 | enc_mid_out_do = enc_first_out_do + enc_first_interact_info_do 119 | 120 | for index in range(self.num_rnn_layers - 1): 121 | enc_mid_out_od = self.mediate_activation(enc_mid_out_od) 122 | enc_mid_out_do = self.mediate_activation(enc_mid_out_do) 123 | 124 | enc_mid_out_od, enc_mid_hidden_od = self.OD.encoder_second_layer(index, 125 | enc_mid_out_od, 126 | enc_hiddens_od[index + 1], 127 | edge_index, 128 | edge_attr) 129 | enc_mid_out_do, enc_mid_hidden_do = self.DO.encoder_second_layer(index, 130 | enc_mid_out_do, 131 | enc_hiddens_do[index + 1], 132 | edge_index, 133 | edge_attr) 134 | 135 | enc_mid_interact_info_od, enc_mid_interact_info_do = self.encoder_second_interact[index]( 136 | enc_mid_hidden_od, 137 | enc_mid_hidden_do) 138 | 139 | enc_hiddens_od[index + 1] = enc_mid_hidden_od + enc_mid_interact_info_od 140 | enc_hiddens_do[index + 1] = enc_mid_hidden_do + enc_mid_interact_info_do 141 | 142 | return enc_hiddens_od, enc_hiddens_do 143 | 144 | def scheduled_sampling(self, 145 | out, 146 | label, 147 | GO): 148 | if self.training and self.use_curriculum_learning: 149 | c = random.uniform(0, 1) 150 | T = self.inverse_sigmoid_scheduler_sampling( 151 | self.global_step, 152 | self.cl_decay_steps) 153 | use_truth_sequence = True if c < T else False 154 | else: 155 | use_truth_sequence = False 156 | 157 | if use_truth_sequence: 158 | # Feed the prev label as the next input 159 | decoder_input = label 160 | else: 161 | # detach from history as input 162 | decoder_input = out.detach().view(-1, self.num_output_dim) 163 | if not self.use_input: 164 | decoder_input = GO.detach() 165 | 166 | return decoder_input 167 | 168 | def decoder_od_do(self, 169 | sequences, 170 | enc_hiddens_od, 171 | enc_hiddens_do, 172 | edge_index, 173 | edge_attr=None): 174 | predictions_od = [] 175 | predictions_do = [] 176 | 177 | GO_od = torch.zeros(enc_hiddens_od[0].size()[0], 178 | self.num_output_dim, 179 | dtype=enc_hiddens_od[0].dtype, 180 | device=enc_hiddens_od[0].device) 181 | GO_do = torch.zeros(enc_hiddens_do[0].size()[0], 182 | self.num_output_dim, 183 | dtype=enc_hiddens_do[0].dtype, 184 | device=enc_hiddens_do[0].device) 185 | 186 | dec_input_od = GO_od 187 | dec_hiddens_od = enc_hiddens_od 188 | 189 | dec_input_do = GO_do 190 | dec_hiddens_do = enc_hiddens_do 191 | 192 | for t in range(self.horizon): 193 | dec_first_out_od, dec_first_hidden_od = self.OD.decoder_first_layer(dec_input_od, 194 | dec_hiddens_od[0], 195 | edge_index, 196 | edge_attr) 197 | 198 | dec_first_out_do, dec_first_hidden_do = self.DO.decoder_first_layer(dec_input_do, 199 | dec_hiddens_do[0], 200 | edge_index, 201 | edge_attr) 202 | 203 | dec_first_interact_info_od, dec_first_interact_info_do = self.decoder_first_interact( 204 | dec_first_hidden_od, 205 | dec_first_hidden_do) 206 | 207 | dec_hiddens_od[0] = dec_first_hidden_od + dec_first_interact_info_od 208 | dec_hiddens_do[0] = dec_first_hidden_do + dec_first_interact_info_do 209 | dec_mid_out_od = dec_first_out_od + dec_first_interact_info_od 210 | dec_mid_out_do = dec_first_out_do + dec_first_interact_info_do 211 | 212 | for index in range(self.num_rnn_layers - 1): 213 | dec_mid_out_od = self.mediate_activation(dec_mid_out_od) 214 | dec_mid_out_do = self.mediate_activation(dec_mid_out_do) 215 | dec_mid_out_od, dec_mid_hidden_od = self.OD.decoder_second_layer(index, 216 | dec_mid_out_od, 217 | dec_hiddens_od[index + 1], 218 | edge_index, 219 | edge_attr) 220 | dec_mid_out_do, dec_mid_hidden_do = self.DO.decoder_second_layer(index, 221 | dec_mid_out_do, 222 | dec_hiddens_do[index + 1], 223 | edge_index, 224 | edge_attr) 225 | 226 | dec_second_interact_info_od, dec_second_interact_info_do = self.decoder_second_interact[index]( 227 | dec_mid_hidden_od, 228 | dec_mid_hidden_do) 229 | 230 | dec_hiddens_od[index + 1] = dec_mid_hidden_od + dec_second_interact_info_od 231 | dec_hiddens_do[index + 1] = dec_mid_hidden_do + dec_second_interact_info_do 232 | dec_mid_out_od = dec_mid_out_od + dec_second_interact_info_od 233 | dec_mid_out_do = dec_mid_out_do + dec_second_interact_info_do 234 | 235 | dec_mid_out_od = dec_mid_out_od.reshape(-1, self.num_units) 236 | dec_mid_out_do = dec_mid_out_do.reshape(-1, self.num_units) 237 | 238 | dec_mid_out_od = self.OD.output_layer(dec_mid_out_od).view(-1, self.num_nodes, self.num_output_dim) 239 | dec_mid_out_do = self.DO.output_layer(dec_mid_out_do).view(-1, self.num_nodes, self.num_output_dim) 240 | 241 | predictions_od.append(dec_mid_out_od) 242 | predictions_do.append(dec_mid_out_do) 243 | 244 | dec_input_od = self.scheduled_sampling(dec_mid_out_od, sequences[t].y_od, GO_od) 245 | dec_input_do = self.scheduled_sampling(dec_mid_out_do, sequences[t].y_do, GO_do) 246 | 247 | if self.training: 248 | self.global_step += 1 249 | 250 | return torch.stack(predictions_od).transpose(0, 1), torch.stack(predictions_do).transpose(0, 1) 251 | 252 | def forward(self, sequences, sequences_y): 253 | edge_index = sequences[0].edge_index.detach() 254 | edge_attr = sequences[0].edge_attr.detach() 255 | 256 | enc_hiddens_od, enc_hiddens_do = self.encoder_od_do(sequences, 257 | edge_index=edge_index, 258 | edge_attr=edge_attr) 259 | predictions_od, predictions_do = self.decoder_od_do(sequences_y, 260 | enc_hiddens_od, 261 | enc_hiddens_do, 262 | edge_index=edge_index, 263 | edge_attr=edge_attr) 264 | 265 | return predictions_od, predictions_do 266 | 267 | 268 | -------------------------------------------------------------------------------- /lib/utils_HIAM.py: -------------------------------------------------------------------------------- 1 | # part of this code are copied from DCRNN 2 | import logging 3 | import os 4 | import pickle 5 | import sys 6 | import numpy as np 7 | import torch 8 | from torch_geometric.data import Batch, Data 9 | # from sklearn.externals import joblib 10 | 11 | class DataLoader(object): 12 | 13 | def __init__(self, 14 | x_od, 15 | y_od, 16 | unfinished, 17 | history, 18 | yesterday, 19 | x_do, 20 | y_do, 21 | xtime, 22 | ytime, 23 | batch_size, 24 | pad_with_last_sample=True, 25 | shuffle=False): 26 | """ 27 | 28 | :param xs: 29 | :param ys: 30 | :param batch_size: 31 | :param pad_with_last_sample: pad with the last sample to make number of samples divisible to batch_size. 32 | """ 33 | self.batch_size = batch_size 34 | self.current_ind = 0 35 | if pad_with_last_sample: 36 | num_padding = (batch_size - (len(x_od) % batch_size)) % batch_size 37 | x_od_padding = np.repeat(x_od[-1:], num_padding, axis=0) 38 | x_do_padding = np.repeat(x_do[-1:], num_padding, axis=0) 39 | xtime_padding = np.repeat(xtime[-1:], num_padding, axis=0) 40 | y_od_padding = np.repeat(y_od[-1:], num_padding, axis=0) 41 | y_do_padding = np.repeat(y_do[-1:], num_padding, axis=0) 42 | ytime_padding = np.repeat(ytime[-1:], num_padding, axis=0) 43 | 44 | x_od = np.concatenate([x_od, x_od_padding], axis=0) 45 | x_do = np.concatenate([x_do, x_do_padding], axis=0) 46 | xtime = np.concatenate([xtime, xtime_padding], axis=0) 47 | y_od = np.concatenate([y_od, y_od_padding], axis=0) 48 | y_do = np.concatenate([y_do, y_do_padding], axis=0) 49 | ytime = np.concatenate([ytime, ytime_padding], axis=0) 50 | 51 | unfi_num_padding = (batch_size - (len(unfinished) % batch_size)) % batch_size 52 | unfinished_padding = np.repeat(unfinished[-1:], unfi_num_padding, axis=0) 53 | unfinished = np.concatenate([unfinished, unfinished_padding], axis=0) 54 | 55 | history_num_padding = (batch_size - (len(history) % batch_size)) % batch_size 56 | history_padding = np.repeat(history[-1:], history_num_padding, axis=0) 57 | history = np.concatenate([history, history_padding], axis=0) 58 | 59 | yesterday_num_padding = (batch_size - (len(yesterday) % batch_size)) % batch_size 60 | yesterday_padding = np.repeat(yesterday[-1:], yesterday_num_padding, axis=0) 61 | yesterday = np.concatenate([yesterday, yesterday_padding], axis=0) 62 | self.size = len(x_od) 63 | self.num_batch = int(self.size // self.batch_size) 64 | # if shuffle: 65 | # permutation = np.random.permutation(self.size) 66 | # x_do, y_do = x_do[permutation], y_do[permutation] 67 | # x_od, y_od = x_od[permutation], y_od[permutation] 68 | # xtime, ytime = xtime[permutation], ytime[permutation] 69 | # unfinished = unfinished[permutation] 70 | # history = history[permutation] 71 | # yesterday = yesterday[permutation] 72 | self.x_do = x_do 73 | self.y_do = y_do 74 | self.x_od = x_od 75 | self.y_od = y_od 76 | self.unfinished = unfinished 77 | self.history = history 78 | self.yesterday = yesterday 79 | self.xtime = xtime 80 | self.ytime = ytime 81 | 82 | def shuffle(self): 83 | permutation = np.random.permutation(self.size) 84 | x_do, y_do = self.x_do[permutation], self.y_do[permutation] 85 | x_od, y_od = self.x_od[permutation], self.y_od[permutation] 86 | xtime, ytime = self.xtime[permutation], self.ytime[permutation] 87 | unfinished = self.unfinished[permutation] 88 | history = self.history[permutation] 89 | yesterday = self.yesterday[permutation] 90 | self.x_do = x_do 91 | self.y_do = y_do 92 | self.x_od = x_od 93 | self.y_od = y_od 94 | self.unfinished = unfinished 95 | self.history = history 96 | self.yesterday = yesterday 97 | self.xtime = xtime 98 | self.ytime = ytime 99 | 100 | def get_iterator(self): 101 | self.current_ind = 0 102 | 103 | def _wrapper(): 104 | while self.current_ind < self.num_batch: 105 | start_ind = self.batch_size * self.current_ind 106 | end_ind = min(self.size, 107 | self.batch_size * (self.current_ind + 1)) 108 | x_do_i = self.x_do[start_ind:end_ind, ...] 109 | y_do_i = self.y_do[start_ind:end_ind, ...] 110 | x_od_i = self.x_od[start_ind:end_ind, ...] 111 | y_od_i = self.y_od[start_ind:end_ind, ...] 112 | unfinished_i = self.unfinished[start_ind:end_ind, ...] 113 | history_i = self.history[start_ind:end_ind, ...] 114 | yesterday_i = self.yesterday[start_ind:end_ind, ...] 115 | xtime_i = self.xtime[start_ind:end_ind, ...] 116 | ytime_i = self.ytime[start_ind:end_ind, ...] 117 | yield (x_od_i, y_od_i, x_do_i, y_do_i, unfinished_i, history_i, yesterday_i, xtime_i, ytime_i) # 注意顺序 118 | self.current_ind += 1 119 | return _wrapper() 120 | 121 | 122 | class StandardScaler_Torch: 123 | """ 124 | Standard the input 125 | """ 126 | 127 | def __init__(self, mean, std, device): 128 | self.mean = torch.tensor(data=mean, dtype=torch.float, device=device) 129 | self.std = torch.tensor(data=std, dtype=torch.float, device=device) 130 | 131 | def transform(self, data): 132 | return (data - self.mean) / self.std 133 | 134 | def inverse_transform(self, data): 135 | return (data * self.std) + self.mean 136 | 137 | 138 | class StandardScaler: 139 | """ 140 | Standard the input 141 | """ 142 | 143 | def __init__(self, mean, std): 144 | self.mean = mean 145 | self.std = std 146 | 147 | def transform(self, data): 148 | return (data - self.mean) / self.std 149 | 150 | def inverse_transform(self, data): 151 | return (data * self.std) + self.mean 152 | 153 | 154 | 155 | def config_logging(log_dir, log_filename='info.log', level=logging.INFO): 156 | # Add file handler and stdout handler 157 | formatter = logging.Formatter( 158 | '%(asctime)s - %(name)s - %(levelname)s - %(message)s') 159 | # Create the log directory if necessary. 160 | try: 161 | os.makedirs(log_dir) 162 | except OSError: 163 | pass 164 | file_handler = logging.FileHandler(os.path.join(log_dir, log_filename)) 165 | file_handler.setFormatter(formatter) 166 | file_handler.setLevel(level=level) 167 | # Add console handler. 168 | console_formatter = logging.Formatter( 169 | '%(asctime)s - %(levelname)s - %(message)s') 170 | console_handler = logging.StreamHandler(sys.stdout) 171 | console_handler.setFormatter(console_formatter) 172 | console_handler.setLevel(level=level) 173 | logging.basicConfig(handlers=[file_handler, console_handler], level=level) 174 | 175 | 176 | def get_logger(log_dir, 177 | name, 178 | log_filename='info.log', 179 | level=logging.INFO, 180 | write_to_file=True): 181 | logger = logging.getLogger(name) 182 | logger.setLevel(level) 183 | logger.propagate = False 184 | # Add file handler and stdout handler 185 | formatter = logging.Formatter( 186 | '%(asctime)s - %(name)s - %(levelname)s - %(message)s') 187 | file_handler = logging.FileHandler(os.path.join(log_dir, log_filename)) 188 | file_handler.setFormatter(formatter) 189 | # Add console handler. 190 | console_formatter = logging.Formatter( 191 | '%(asctime)s - %(levelname)s - %(message)s') 192 | console_handler = logging.StreamHandler(sys.stdout) 193 | console_handler.setFormatter(console_formatter) 194 | if write_to_file is True: 195 | logger.addHandler(file_handler) 196 | logger.addHandler(console_handler) 197 | # Add google cloud log handler 198 | logger.info('Log directory: %s', log_dir) 199 | return logger 200 | 201 | 202 | def load_dataset(dataset_dir, 203 | do_dataset_dir, 204 | batch_size, 205 | test_batch_size=None, 206 | scaler_axis=(0, 207 | 1, 208 | 2, 209 | 3), 210 | **kwargs): 211 | data = {} 212 | 213 | for category in ['train', 'val', 'test']: 214 | cat_data = load_pickle(os.path.join(dataset_dir, category + '.pkl')) 215 | history_data = load_pickle(os.path.join(dataset_dir, category + '_history_long.pkl')) 216 | yesterday_data = load_pickle(os.path.join(dataset_dir, category + '_history_short.pkl')) 217 | do_data = load_pickle(os.path.join(do_dataset_dir, category + '_do.pkl')) 218 | data['x_' + category] = cat_data['finished'] 219 | data['unfinished_' + category] = cat_data['unfinished'] 220 | data['xtime_' + category] = cat_data['xtime'] 221 | data['y_' + category] = cat_data['y'] 222 | data['ytime_' + category] = cat_data['ytime'] 223 | data['do_x_'+ category] = do_data['do_x'] 224 | data['do_y_' + category] = do_data['do_y'] 225 | 226 | his_sum = np.sum(history_data['history'], axis=-1) 227 | history_distribution = np.nan_to_num(np.divide(history_data['history'], np.expand_dims(his_sum, axis=-1))) 228 | data['history_' + category] = np.multiply(history_distribution, cat_data['unfinished']) 229 | 230 | yesterday_sum = np.sum(yesterday_data['history'], axis=-1) 231 | yesterday_distribution = np.nan_to_num(np.divide(yesterday_data['history'], np.expand_dims(yesterday_sum, axis=-1))) 232 | data['yesterday_' + category] = np.multiply(yesterday_distribution, cat_data['unfinished']) 233 | 234 | 235 | scaler = StandardScaler(mean=data['x_train'].mean(axis=scaler_axis), 236 | std=data['x_train'].std(axis=scaler_axis)) 237 | do_scaler = StandardScaler(mean=data['do_x_train'].mean(axis=scaler_axis), 238 | std=data['do_x_train'].std(axis=scaler_axis)) 239 | # Data format 240 | for category in ['train', 'val', 'test']: 241 | data['x_' + category] = scaler.transform(data['x_' + category]) 242 | data['y_' + category] = scaler.transform(data['y_' + category]) 243 | data['do_x_' + category] = do_scaler.transform(data['do_x_' + category]) 244 | data['do_y_' + category] = do_scaler.transform(data['do_y_' + category]) 245 | data['unfinished_' + category] = scaler.transform(data['unfinished_' + category]) 246 | data['history_' + category] = scaler.transform(data['history_' + category]) 247 | data['yesterday_' + category] = scaler.transform(data['yesterday_' + category]) 248 | 249 | data['train_loader'] = DataLoader(data['x_train'], 250 | data['y_train'], 251 | data['unfinished_train'], 252 | data['history_train'], 253 | data['yesterday_train'], 254 | data['do_x_train'], 255 | data['do_y_train'], 256 | data['xtime_train'], 257 | data['ytime_train'], 258 | batch_size, 259 | shuffle=True) 260 | data['val_loader'] = DataLoader(data['x_val'], 261 | data['y_val'], 262 | data['unfinished_val'], 263 | data['history_val'], 264 | data['yesterday_val'], 265 | data['do_x_val'], 266 | data['do_y_val'], 267 | data['xtime_val'], 268 | data['ytime_val'], 269 | test_batch_size, 270 | shuffle=False) 271 | 272 | data['test_loader'] = DataLoader(data['x_test'], 273 | data['y_test'], 274 | data['unfinished_test'], 275 | data['history_test'], 276 | data['yesterday_test'], 277 | data['do_x_test'], 278 | data['do_y_test'], 279 | data['xtime_test'], 280 | data['ytime_test'], 281 | test_batch_size, 282 | shuffle=False) 283 | data['scaler'] = scaler 284 | data['do_scaler'] = do_scaler 285 | 286 | return data 287 | 288 | def load_graph_data(pkl_filename): 289 | adj_mx = load_pickle(pkl_filename) 290 | return adj_mx.astype(np.float32) 291 | 292 | 293 | def load_pickle(pickle_file): 294 | try: 295 | with open(pickle_file, 'rb') as f: 296 | pickle_data = pickle.load(f) 297 | except UnicodeDecodeError as e: 298 | with open(pickle_file, 'rb') as f: 299 | pickle_data = pickle.load(f, encoding='latin1') 300 | except Exception as e: 301 | print('Unable to load data ', pickle_file, ':', e) 302 | raise 303 | return pickle_data 304 | 305 | 306 | class SimpleBatch(list): 307 | 308 | def to(self, device): 309 | for ele in self: 310 | ele.to(device) 311 | return self 312 | 313 | 314 | def collate_wrapper(x_od, y_od, x_do, y_do, unfinished, history, yesterday, edge_index, edge_attr, seq_len, horizon, device, return_y=True): 315 | x_od = torch.tensor(x_od, dtype=torch.float, device=device) 316 | y_od = torch.tensor(y_od, dtype=torch.float, device=device) 317 | x_od = x_od.transpose(dim0=1, dim1=0) # (T, N, num_nodes, num_features) 318 | y_od_T_first = y_od.transpose(dim0=1, dim1=0) # (T, N, num_nodes, num_features) 319 | 320 | x_do = torch.tensor(x_do, dtype=torch.float, device=device) 321 | y_do = torch.tensor(y_do, dtype=torch.float, device=device) 322 | x_do = x_do.transpose(dim0=1, dim1=0) # (T, N, num_nodes, num_features) 323 | y_do_T_first = y_do.transpose(dim0=1, dim1=0) # (T, N, num_nodes, num_features) 324 | 325 | unfinished = torch.tensor(unfinished, dtype=torch.float, device=device) 326 | unfinished = unfinished.transpose(dim0=1, dim1=0) 327 | 328 | history = torch.tensor(history, dtype=torch.float, device=device) 329 | history = history.transpose(dim0=1, dim1=0) 330 | 331 | yesterday = torch.tensor(yesterday, dtype=torch.float, device=device) 332 | yesterday = yesterday.transpose(dim0=1, dim1=0) 333 | # do not tranpose y_truth 334 | T = x_od.size()[0] 335 | H = y_od_T_first.size()[0] 336 | N = x_od.size()[1] 337 | # generate batched sequence. 338 | sequences = [] 339 | sequences_y = [] 340 | for t in range(T - seq_len, T): 341 | cur_batch_x_od = x_od[t] 342 | cur_batch_x_do = x_do[t] 343 | cur_batch_unfinished = unfinished[t] 344 | cur_batch_history = history[t] 345 | cur_batch_yesterday = yesterday[t] 346 | 347 | batch = Batch.from_data_list([ 348 | Data(x_do=cur_batch_x_do[i], 349 | x_od=cur_batch_x_od[i], 350 | unfinished=cur_batch_unfinished[i], 351 | history=cur_batch_history[i], 352 | yesterday=cur_batch_yesterday[i], 353 | edge_index=edge_index, 354 | edge_attr=edge_attr) for i in range(N) 355 | ]) 356 | sequences.append(batch) 357 | 358 | for t in range(H - horizon, H): 359 | cur_batch_y_od = y_od_T_first[t] 360 | cur_batch_y_do = y_do_T_first[t] 361 | 362 | batch_y = Batch.from_data_list([ 363 | Data(y_do=cur_batch_y_do[i], 364 | y_od=cur_batch_y_od[i]) for i in range(N) 365 | ]) 366 | sequences_y.append(batch_y) 367 | if return_y: 368 | return SimpleBatch(sequences), SimpleBatch(sequences_y), y_od, y_do 369 | else: 370 | return SimpleBatch(sequences), SimpleBatch(sequences_y) 371 | 372 | 373 | def collate_wrapper_multi_branches(x_numpy, y_numpy, edge_index_list, device): 374 | sequences_multi_branches = [] 375 | for edge_index in edge_index_list: 376 | sequences, y = collate_wrapper(x_numpy, y_numpy, edge_index, device, return_y=True) 377 | sequences_multi_branches.append(sequences) 378 | 379 | return sequences_multi_branches, y 380 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | # encoding:utf-8 2 | import random 3 | import argparse 4 | import time 5 | import yaml 6 | import numpy as np 7 | import torch 8 | import os 9 | 10 | from torch import nn 11 | from torch.nn.utils import clip_grad_norm_ 12 | from torch import optim 13 | from torch.optim.lr_scheduler import MultiStepLR 14 | from torch.nn.init import xavier_uniform_ 15 | from lib import utils_HIAM as utils 16 | from lib import metrics 17 | from lib.utils_HIAM import collate_wrapper 18 | from models.Net import Net 19 | try: 20 | from yaml import CLoader as Loader, CDumper as Dumper 21 | except ImportError: 22 | from yaml import Loader, Dumper 23 | 24 | seed = 1234 25 | random.seed(seed) 26 | np.random.seed(seed) 27 | torch.manual_seed(seed) 28 | parser = argparse.ArgumentParser() 29 | parser.add_argument('--config_filename', 30 | default=None, 31 | type=str, 32 | help='Configuration filename for restoring the model.') 33 | args = parser.parse_args() 34 | 35 | 36 | def read_cfg_file(filename): 37 | with open(filename, 'r') as ymlfile: 38 | cfg = yaml.load(ymlfile, Loader=Loader) 39 | return cfg 40 | 41 | 42 | def run_model(model, data_iterator, edge_index, edge_attr, device, seq_len, horizon, output_dim): 43 | """ 44 | return a list of (horizon_i, batch_size, num_nodes, output_dim) 45 | """ 46 | # while evaluation, we need model.eval and torch.no_grad 47 | model.eval() 48 | y_od_pred_list = [] 49 | y_do_pred_list = [] 50 | for _, (x_od, y_od, x_do, y_do, unfinished, history, yesterday, xtime, ytime) in enumerate(data_iterator): 51 | y_od = y_od[..., :output_dim] 52 | y_do = y_do[..., :output_dim] 53 | sequences, sequences_y, y_od, y_do = collate_wrapper(x_od=x_od, y_od=y_od, x_do=x_do, y_do=y_do, unfinished=unfinished, history=history, yesterday=yesterday, 54 | edge_index=edge_index, 55 | edge_attr=edge_attr, 56 | device=device, 57 | seq_len=seq_len, 58 | horizon=horizon) 59 | # (T, N, num_nodes, num_out_channels) 60 | with torch.no_grad(): 61 | y_od_pred, y_do_pred = model(sequences, sequences_y) 62 | if y_od_pred is not None: 63 | y_od_pred_list.append(y_od_pred.detach().cpu().numpy()) 64 | if y_do_pred is not None: 65 | y_do_pred_list.append(y_do_pred.detach().cpu().numpy()) 66 | return y_od_pred_list, y_do_pred_list 67 | 68 | 69 | def evaluate(model, 70 | dataset, 71 | dataset_type, 72 | edge_index, 73 | edge_attr, 74 | device, 75 | seq_Len, 76 | horizon, 77 | output_dim, 78 | logger, 79 | detail=True, 80 | cfg=None, 81 | format_result=False): 82 | if detail: 83 | logger.info('Evaluation_{}_Begin:'.format(dataset_type)) 84 | 85 | y_od_preds, y_do_preds = run_model( 86 | model, 87 | data_iterator=dataset['{}_loader'.format(dataset_type)].get_iterator(), 88 | edge_index=edge_index, 89 | edge_attr=edge_attr, 90 | device=device, 91 | seq_len=seq_Len, 92 | horizon=horizon, 93 | output_dim=output_dim) 94 | 95 | evaluate_category = [] 96 | if len(y_od_preds) > 0: 97 | evaluate_category.append("od") 98 | if len(y_do_preds) > 0: 99 | evaluate_category.append("do") 100 | results = {} 101 | for category in evaluate_category: 102 | if category == 'od': 103 | y_preds = y_od_preds 104 | scaler = dataset['scaler'] 105 | gt = dataset['y_{}'.format(dataset_type)] 106 | else: 107 | y_preds = y_do_preds 108 | scaler = dataset['do_scaler'] 109 | # scaler = dataset['scaler'] 110 | gt = dataset['do_y_{}'.format(dataset_type)] 111 | y_preds = np.concatenate(y_preds, axis=0) # concat in batch_size dim. 112 | mae_list = [] 113 | mape_net_list = [] 114 | rmse_list = [] 115 | mae_sum = 0 116 | 117 | mape_net_sum = 0 118 | rmse_sum = 0 119 | logger.info("{}:".format(category)) 120 | horizon = cfg['model']['horizon'] 121 | for horizon_i in range(horizon): 122 | y_truth = scaler.inverse_transform( 123 | gt[:, horizon_i, :, :output_dim]) 124 | 125 | y_pred = scaler.inverse_transform( 126 | y_preds[:y_truth.shape[0], horizon_i, :, :output_dim]) 127 | y_pred[y_pred < 0] = 0 128 | mae = metrics.masked_mae_np(y_pred, y_truth) 129 | mape_net = metrics.masked_mape_np(y_pred, y_truth) 130 | rmse = metrics.masked_rmse_np(y_pred, y_truth) 131 | mae_sum += mae 132 | mape_net_sum += mape_net 133 | rmse_sum += rmse 134 | mae_list.append(mae) 135 | 136 | mape_net_list.append(mape_net) 137 | rmse_list.append(rmse) 138 | 139 | msg = "Horizon {:02d}, MAE: {:.2f}, RMSE: {:.2f}, MAPE_net: {:.4f}" 140 | if detail: 141 | logger.info(msg.format(horizon_i + 1, mae, rmse, mape_net)) 142 | results['MAE_' + category] = mae_sum / horizon 143 | results['RMSE_' + category] = rmse_sum / horizon 144 | results['MAPE_net_' + category] = mape_net_sum / horizon 145 | if detail: 146 | logger.info('Evaluation_{}_End:'.format(dataset_type)) 147 | if format_result: 148 | for i in range(len(mae_list)): 149 | print('{:.2f}'.format(mae_list[i])) 150 | print('{:.2f}'.format(rmse_list[i])) 151 | print('{:.2f}%'.format(mape_net_list[i] * 100)) 152 | print() 153 | else: 154 | return results 155 | 156 | 157 | class StepLR2(MultiStepLR): 158 | """StepLR with min_lr""" 159 | def __init__(self, 160 | optimizer, 161 | milestones, 162 | gamma=0.1, 163 | last_epoch=-1, 164 | min_lr=2.0e-6): 165 | 166 | self.optimizer = optimizer 167 | self.milestones = milestones 168 | self.gamma = gamma 169 | self.last_epoch = last_epoch 170 | self.min_lr = min_lr 171 | super(StepLR2, self).__init__(optimizer, milestones, gamma) 172 | 173 | def get_lr(self): 174 | lr_candidate = super(StepLR2, self).get_lr() 175 | if isinstance(lr_candidate, list): 176 | for i in range(len(lr_candidate)): 177 | lr_candidate[i] = max(self.min_lr, lr_candidate[i]) 178 | 179 | else: 180 | lr_candidate = max(self.min_lr, lr_candidate) 181 | 182 | return lr_candidate 183 | 184 | 185 | def _get_log_dir(kwargs): 186 | log_dir = kwargs['train'].get('log_dir') 187 | if log_dir is None: 188 | batch_size = kwargs['data'].get('batch_size') 189 | learning_rate = kwargs['train'].get('base_lr') 190 | num_rnn_layers = kwargs['model'].get('num_rnn_layers') 191 | rnn_units = kwargs['model'].get('rnn_units') 192 | structure = '-'.join(['%d' % rnn_units for _ in range(num_rnn_layers)]) 193 | 194 | # 保存的文件夹名称 195 | run_id = 'HIAM_%s_lr%g_bs%d_%s/' % ( 196 | structure, 197 | learning_rate, 198 | batch_size, 199 | time.strftime('%m%d%H%M%S')) 200 | base_dir = kwargs.get('base_dir') 201 | log_dir = os.path.join(base_dir, run_id) 202 | if not os.path.exists(log_dir): 203 | os.makedirs(log_dir) 204 | return log_dir 205 | 206 | 207 | def init_weights(m): 208 | classname = m.__class__.__name__ # 2 209 | if classname.find('Conv') != -1 and classname.find('RGCN') == -1: 210 | xavier_uniform_(m.weight.data) 211 | if type(m) == nn.Linear: 212 | xavier_uniform_(m.weight.data) 213 | #xavier_uniform_(m.bias.data) 214 | 215 | 216 | def main(args): 217 | cfg = read_cfg_file(args.config_filename) 218 | log_dir = _get_log_dir(cfg) 219 | log_level = cfg.get('log_level', 'INFO') 220 | 221 | logger = utils.get_logger(log_dir, __name__, 'info.log', level=log_level) 222 | 223 | device = torch.device( 224 | 'cuda') if torch.cuda.is_available() else torch.device('cpu') 225 | # all edge_index in same dataset is same 226 | logger.info(cfg) 227 | batch_size = cfg['data']['batch_size'] 228 | seq_len = cfg['model']['seq_len'] 229 | horizon = cfg['model']['horizon'] 230 | 231 | adj_mx_list = [] 232 | graph_pkl_filename = cfg['data']['graph_pkl_filename'] 233 | 234 | if not isinstance(graph_pkl_filename, list): 235 | graph_pkl_filename = [graph_pkl_filename] 236 | 237 | src = [] 238 | dst = [] 239 | for g in graph_pkl_filename: 240 | adj_mx = utils.load_graph_data(g) 241 | for i in range(len(adj_mx)): # 构建邻接矩阵 242 | adj_mx[i, i] = 0 243 | adj_mx_list.append(adj_mx) 244 | 245 | adj_mx = np.stack(adj_mx_list, axis=-1) 246 | print("adj_mx:", adj_mx.shape) 247 | if cfg['model'].get('norm', False): 248 | print('row normalization') 249 | adj_mx = adj_mx / (adj_mx.sum(axis=0) + 1e-18) # 归一化 250 | src, dst = adj_mx.sum(axis=-1).nonzero() 251 | print("src, dst:", src.shape, dst.shape) 252 | edge_index = torch.tensor([src, dst], dtype=torch.long, device=device) 253 | edge_attr = torch.tensor(adj_mx[adj_mx.sum(axis=-1) != 0], 254 | dtype=torch.float, 255 | device=device) 256 | print("train, edge:", edge_index.shape, edge_attr.shape) 257 | output_dim = cfg['model']['output_dim'] 258 | for i in range(adj_mx.shape[-1]): 259 | logger.info(adj_mx[..., i]) 260 | 261 | ## load dataset 262 | dataset = utils.load_dataset(**cfg['data'], scaler_axis=(0, 1, 2, 3)) 263 | for k, v in dataset.items(): 264 | if hasattr(v, 'shape'): 265 | logger.info((k, v.shape)) 266 | 267 | scaler_od = dataset['scaler'] 268 | scaler_od_torch = utils.StandardScaler_Torch(scaler_od.mean, 269 | scaler_od.std, 270 | device=device) 271 | logger.info('scaler_od.mean:{}, scaler_od.std:{}'.format(scaler_od.mean, 272 | scaler_od.std)) 273 | scaler_do = dataset['do_scaler'] 274 | scaler_do_torch = utils.StandardScaler_Torch(scaler_do.mean, 275 | scaler_do.std, 276 | device=device) 277 | logger.info('scaler_do.mean:{}, scaler_do.std:{}'.format(scaler_do.mean, 278 | scaler_do.std)) 279 | ## define model 280 | model = Net(cfg, logger).to(device) 281 | model.apply(init_weights) 282 | criterion = nn.L1Loss(reduction='mean') 283 | optimizer = optim.Adam(model.parameters(), 284 | lr=cfg['train']['base_lr'], 285 | eps=cfg['train']['epsilon']) 286 | scheduler = StepLR2(optimizer=optimizer, 287 | milestones=cfg['train']['steps'], 288 | gamma=cfg['train']['lr_decay_ratio'], 289 | min_lr=cfg['train']['min_learning_rate']) 290 | 291 | max_grad_norm = cfg['train']['max_grad_norm'] 292 | train_patience = cfg['train']['patience'] 293 | 294 | update = {} 295 | for category in ['od', 'do']: 296 | update['val_steady_count_'+category] = 0 297 | update['last_val_mae_'+category] = 1e6 298 | update['last_val_mape_net_'+category] = 1e6 299 | 300 | horizon = cfg['model']['horizon'] 301 | 302 | for epoch in range(cfg['train']['epochs']): 303 | total_loss = 0 304 | i = 0 305 | begin_time = time.perf_counter() 306 | dataset['train_loader'].shuffle() 307 | train_iterator = dataset['train_loader'].get_iterator() 308 | model.train() 309 | for _, (x_od, y_od, x_do, y_do, unfinished, history, yesterday, xtime, ytime) in enumerate(train_iterator): 310 | optimizer.zero_grad() 311 | y_od = y_od[:, :horizon, :, :output_dim] 312 | y_do = y_do[:, :horizon, :, :output_dim] 313 | sequences, sequences_y, y_od, y_do = collate_wrapper( 314 | x_od=x_od, y_od=y_od, x_do=x_do, y_do=y_do, 315 | unfinished=unfinished, history=history, yesterday=yesterday, 316 | edge_index=edge_index, edge_attr=edge_attr, 317 | device=device, seq_len=seq_len, horizon=horizon) # generate batch data 318 | 319 | y_od_pred, y_do_pred = model(sequences, sequences_y) 320 | 321 | # OD loss 322 | y_od_pred = scaler_od_torch.inverse_transform(y_od_pred) # *std+mean 323 | y_od = scaler_od_torch.inverse_transform(y_od) 324 | loss_od = criterion(y_od_pred, y_od) 325 | 326 | # DO loss 327 | y_do_pred = scaler_do_torch.inverse_transform(y_do_pred) # *std+mean 328 | y_do = scaler_do_torch.inverse_transform(y_do) 329 | loss_do = criterion(y_do_pred, y_do) 330 | 331 | loss = loss_od + loss_do 332 | total_loss += loss.item() 333 | loss.backward() 334 | 335 | clip_grad_norm_(model.parameters(), max_grad_norm) 336 | optimizer.step() 337 | 338 | i += 1 339 | 340 | # evaluation on validation set 341 | val_result = evaluate(model=model, 342 | dataset=dataset, 343 | dataset_type='val', 344 | edge_index=edge_index, 345 | edge_attr=edge_attr, 346 | device=device, 347 | seq_Len=seq_len, 348 | horizon=horizon, 349 | output_dim=output_dim, 350 | logger=logger, 351 | detail=False, 352 | cfg=cfg) 353 | 354 | time_elapsed = time.perf_counter() - begin_time 355 | 356 | logger.info(('Epoch:{}, total_loss:{}').format(epoch, total_loss / i)) 357 | val_category = ['od', 'do'] 358 | for category in val_category: 359 | logger.info('{}:'.format(category)) 360 | logger.info(('val_mae:{}, val_mape_net:{}' 361 | 'r_loss={:.2f},lr={}, time_elapsed:{}').format( 362 | val_result['MAE_'+category], 363 | val_result['MAPE_net_' + category], 364 | 0, 365 | str(scheduler.get_lr()), 366 | time_elapsed)) 367 | if update['last_val_mae_'+category] > val_result['MAE_'+category]: 368 | logger.info('val_mae decreased from {} to {}'.format( 369 | update['last_val_mae_' + category], 370 | val_result['MAE_' + category])) 371 | update['last_val_mae_' + category] = val_result['MAE_' + category] 372 | update['val_steady_count_' + category] = 0 373 | else: 374 | update['val_steady_count_' + category] += 1 375 | 376 | if update['last_val_mape_net_'+category] > val_result['MAPE_net_'+category]: 377 | logger.info('val_mape_net decreased from {} to {}'.format( 378 | update['last_val_mape_net_' + category], 379 | val_result['MAPE_net_' + category])) 380 | update['last_val_mape_net_' + category] = val_result['MAPE_net_'+category] 381 | 382 | # after per epoch, run evaluation on test dataset 383 | if (epoch + 1) % cfg['train']['test_every_n_epochs'] == 0: 384 | evaluate(model=model, 385 | dataset=dataset, 386 | dataset_type='test', 387 | edge_index=edge_index, 388 | edge_attr=edge_attr, 389 | device=device, 390 | seq_Len=seq_len, 391 | horizon=horizon, 392 | output_dim=output_dim, 393 | logger=logger, 394 | cfg=cfg) 395 | 396 | if (epoch + 1) % cfg['train']['save_every_n_epochs'] == 0: 397 | save_dir = log_dir 398 | if not os.path.exists(save_dir): 399 | os.mkdir(save_dir) 400 | config_path = os.path.join(save_dir, 401 | 'config-{}.yaml'.format(epoch + 1)) 402 | epoch_path = os.path.join(save_dir, 403 | 'epoch-{}.pt'.format(epoch + 1)) 404 | torch.save(model.state_dict(), epoch_path) 405 | with open(config_path, 'w') as f: 406 | from copy import deepcopy 407 | save_cfg = deepcopy(cfg) 408 | save_cfg['model']['save_path'] = epoch_path 409 | f.write(yaml.dump(save_cfg, Dumper=Dumper)) 410 | 411 | if train_patience <= update['val_steady_count_od'] or train_patience <= update['val_steady_count_do']: 412 | logger.info('early stopping.') 413 | break 414 | scheduler.step() 415 | 416 | 417 | if __name__ == "__main__": 418 | main(args) 419 | --------------------------------------------------------------------------------