├── models ├── stgcn │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-38.pyc │ │ ├── layers.cpython-38.pyc │ │ └── models.cpython-38.pyc │ ├── models.py │ └── layers.py ├── __pycache__ │ ├── lstm.cpython-38.pyc │ ├── ma2gcn.cpython-38.pyc │ ├── ASTGCN_r.cpython-38.pyc │ └── graph_wavenet.cpython-38.pyc ├── lstm.py ├── graph_wavenet.py ├── ma2gcn.py └── ASTGCN_r.py ├── imgs ├── arch.png ├── adj_attention.png └── ttb_pipeline.png ├── data ├── io_matrix_grid.npy └── feature_matrix_grid.npy ├── trainer ├── __pycache__ │ ├── __init__.cpython-38.pyc │ ├── trainer_lstm.cpython-38.pyc │ ├── trainer_stgcn.cpython-38.pyc │ ├── trainer_astgcn.cpython-38.pyc │ ├── trainer_ma2gcn.cpython-38.pyc │ └── trainer_graph_wavenet.cpython-38.pyc ├── trainer_astgcn.py ├── trainer_graph_wavenet.py ├── trainer_stgcn.py ├── trainer_lstm.py └── trainer_ma2gcn.py ├── requirements.txt ├── .idea └── .gitignore ├── config ├── config_arima.yaml ├── config_svr.yaml ├── config_lstm.yaml ├── config_graph_wavenet.yaml ├── config_astgcn.yaml ├── config_stgcn.yaml └── config_ma2gcn.yaml ├── README.md ├── main ├── main_svr.py ├── main_arima.py ├── main_lstm.py ├── main_astgcn.py ├── main_graph_wavenet.py ├── main_stgcn.py └── main_ma2gcn.py └── utils.py /models/stgcn/__init__.py: -------------------------------------------------------------------------------- 1 | __all__ = ['layers', 'models'] -------------------------------------------------------------------------------- /imgs/arch.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/imgs/arch.png -------------------------------------------------------------------------------- /data/io_matrix_grid.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/data/io_matrix_grid.npy -------------------------------------------------------------------------------- /imgs/adj_attention.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/imgs/adj_attention.png -------------------------------------------------------------------------------- /imgs/ttb_pipeline.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/imgs/ttb_pipeline.png -------------------------------------------------------------------------------- /data/feature_matrix_grid.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/data/feature_matrix_grid.npy -------------------------------------------------------------------------------- /models/__pycache__/lstm.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/models/__pycache__/lstm.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/ma2gcn.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/models/__pycache__/ma2gcn.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/ASTGCN_r.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/models/__pycache__/ASTGCN_r.cpython-38.pyc -------------------------------------------------------------------------------- /trainer/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/trainer/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/graph_wavenet.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/models/__pycache__/graph_wavenet.cpython-38.pyc -------------------------------------------------------------------------------- /models/stgcn/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/models/stgcn/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /models/stgcn/__pycache__/layers.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/models/stgcn/__pycache__/layers.cpython-38.pyc -------------------------------------------------------------------------------- /models/stgcn/__pycache__/models.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/models/stgcn/__pycache__/models.cpython-38.pyc -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy==1.23.2 2 | pandas==2.2.1 3 | PyYAML==6.0.1 4 | scikit_learn==1.4.1.post1 5 | scipy==1.12.0 6 | statsmodels==0.14.1 7 | torch==1.12.1 8 | -------------------------------------------------------------------------------- /trainer/__pycache__/trainer_lstm.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/trainer/__pycache__/trainer_lstm.cpython-38.pyc -------------------------------------------------------------------------------- /trainer/__pycache__/trainer_stgcn.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/trainer/__pycache__/trainer_stgcn.cpython-38.pyc -------------------------------------------------------------------------------- /trainer/__pycache__/trainer_astgcn.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/trainer/__pycache__/trainer_astgcn.cpython-38.pyc -------------------------------------------------------------------------------- /trainer/__pycache__/trainer_ma2gcn.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/trainer/__pycache__/trainer_ma2gcn.cpython-38.pyc -------------------------------------------------------------------------------- /trainer/__pycache__/trainer_graph_wavenet.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zachysun/Taxi_Traffic_Benchmark/HEAD/trainer/__pycache__/trainer_graph_wavenet.cpython-38.pyc -------------------------------------------------------------------------------- /.idea/.gitignore: -------------------------------------------------------------------------------- 1 | # Default ignored files 2 | /shelf/ 3 | /workspace.xml 4 | # Editor-based HTTP Client requests 5 | /httpRequests/ 6 | # Datasource local storage ignored files 7 | /dataSources/ 8 | /dataSources.local.xml 9 | -------------------------------------------------------------------------------- /config/config_arima.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | adj_data_path: '../data/adjacency_matrix_grid.csv' 3 | feature_matrix_data_path: '../data/feature_matrix_grid.npy' 4 | input_seq: 12 5 | output_seq: 1 6 | num_of_vertices: 225 7 | batch_size: 8 8 | split_rate: [7, 2, 1] 9 | device: None 10 | -------------------------------------------------------------------------------- /config/config_svr.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | adj_data_path: '../data/adjacency_matrix_grid.csv' 3 | feature_matrix_data_path: '../data/feature_matrix_grid.npy' 4 | input_seq: 12 5 | output_seq: 1 6 | num_of_vertices: 225 7 | batch_size: 8 8 | split_rate: [7, 2, 1] 9 | device: cuda:0 10 | -------------------------------------------------------------------------------- /models/lstm.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | class LSTM(nn.Module): 5 | def __init__(self, input_size, hidden_size, output_size): 6 | super(LSTM, self).__init__() 7 | self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True) 8 | self.fc = nn.Linear(hidden_size, output_size) 9 | 10 | def forward(self, x): 11 | lstm_out, _ = self.lstm(x) 12 | out = self.fc(lstm_out[:, -1, :]) 13 | return out 14 | -------------------------------------------------------------------------------- /config/config_lstm.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | adj_data_path: '../data/adjacency_matrix_grid.csv' 3 | feature_matrix_data_path: '../data/feature_matrix_grid.npy' 4 | input_seq: 12 5 | output_seq: 1 6 | num_of_vertices: 225 7 | batch_size: 256 8 | device: cuda:0 9 | split_rate: [7, 2, 1] 10 | 11 | optimize: 12 | learning_rate: 0.001 13 | weight_decay: 0.0001 14 | name: Adam 15 | 16 | trainer: 17 | in_dim: 2 18 | hidden_dim: 16 19 | out_dim: 2 20 | theta: 0.5 21 | epoch: 300 22 | model_path: '../saved_models' 23 | results_path: '../results' 24 | -------------------------------------------------------------------------------- /config/config_graph_wavenet.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | adj_data_path: '../data/adjacency_matrix_grid.csv' 3 | feature_matrix_data_path: '../data/feature_matrix_grid.npy' 4 | input_seq: 12 5 | output_seq: 1 6 | num_of_vertices: 225 7 | batch_size: 8 8 | device: cuda:0 9 | split_rate: [7, 2, 1] 10 | 11 | optimize: 12 | learning_rate: 0.001 13 | weight_decay: 0.0001 14 | name: Adam 15 | 16 | trainer: 17 | K: 3 18 | theta: 0.5 19 | epoch: 300 20 | dropout: 0.5 21 | gcn_bool: True 22 | aptonly: False 23 | addaptadj: True 24 | randomadj: True 25 | nhid: 32 26 | in_dim: 2 27 | out_dim: 1 28 | model_path: '../saved_models' 29 | results_path: '../results' 30 | -------------------------------------------------------------------------------- /config/config_astgcn.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | adj_data_path: '../data/adjacency_matrix_grid.csv' 3 | feature_matrix_data_path: '../data/feature_matrix_grid.npy' 4 | input_seq: 12 5 | output_seq: 1 6 | num_of_vertices: 225 7 | batch_size: 8 8 | device: cuda:0 9 | split_rate: [7, 2, 1] 10 | 11 | optimize: 12 | learning_rate: 0.001 13 | weight_decay: 0.0001 14 | name: Adam 15 | 16 | trainer: 17 | K: 3 18 | theta: 0.5 19 | epoch: 300 20 | ctx: 0 21 | in_channels: 2 22 | nb_block: 2 23 | nb_chev_filter: 64 24 | nb_time_filter: 64 25 | start_epoch: 0 26 | num_of_hours: 1 27 | num_of_days: 0 28 | num_of_weeks: 0 29 | model_path: '../saved_models' 30 | results_path: '../results' 31 | -------------------------------------------------------------------------------- /config/config_stgcn.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | adj_data_path: '../data/adjacency_matrix_grid.csv' 3 | feature_matrix_data_path: '../data/feature_matrix_grid.npy' 4 | input_seq: 12 5 | output_seq: 1 6 | num_of_vertices: 225 7 | batch_size: 8 8 | device: cuda:0 9 | split_rate: [7, 2, 1] 10 | 11 | optimize: 12 | learning_rate: 0.001 13 | weight_decay: 0.0001 14 | name: Adam 15 | 16 | trainer: 17 | Kt: 3 18 | Ks: 3 # [2, 3] 19 | theta: 0.5 20 | epoch: 300 21 | n_his: 12 22 | n_pred: 1 23 | channel: 2 24 | act_func: glu # [glu, gtu] 25 | graph_conv_type: cheb_graph_conv # [cheb_graph_conv, graph_conv] 26 | stblock_num: 2 27 | enable_bias: True 28 | dropout: 0.5 29 | model_path: '../saved_models' 30 | results_path: '../results' -------------------------------------------------------------------------------- /config/config_ma2gcn.yaml: -------------------------------------------------------------------------------- 1 | dataset: 2 | origin_adj_data_path: '../data/adjacency_matrix_grid.csv' 3 | io_adj_data_path: '../data/io_matrix_grid.npy' 4 | feature_matrix_data_path: '../data/feature_matrix_grid.npy' 5 | input_seq: 12 6 | output_seq: 1 7 | batch_size: 8 8 | device: cuda:0 9 | basic_interval: 5 # minutes 10 | io_matrix_interval: 3 # hours 11 | split_rate: [7, 2, 1] 12 | 13 | # Adam, AdamW, RMSprop 14 | optimize: 15 | learning_rate: 0.001 16 | weight_decay: 0.0001 17 | name: Adam 18 | 19 | trainer: 20 | K: 3 21 | epoch: 300 22 | theta: 0.5 23 | block_nums: 3 24 | drop_rate: 0.00 25 | adj_hidden_dim: 16 26 | kernel_size: [1, 3] 27 | in_channel_list: [2, 16, 8] 28 | out_channel_list: [16, 8, 2] 29 | in_tcn_dim_list: [2, 16, 8] 30 | out_tcn_dim_list: [2, 16, 8] 31 | model_path: '../saved_models' 32 | results_path: '../results' 33 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## Taxi Traffic Benchmark 2 | 3 | ### Pipeline 4 | 5 | 6 | 7 | ### Dataset 8 | 9 | - [Shanghai Taxi](https://cse.hkust.edu.hk/scrg/) 10 | - **Note:** This link has expired. The processed data is located in the "/data" directory. 11 | 12 | ### Baseline 13 | 14 | - ARIMA 15 | - SVR 16 | - LSTM 17 | - STGCN[[Paper]](https://arxiv.org/abs/1709.04875) [[Code]](./models/stgcn) [[Official Code]](https://github.com/hazdzz/STGCN) 18 | - ASTGCN[[Paper]](https://ojs.aaai.org/index.php/AAAI/article/view/3881) [[Code]](./models/ASTGCN_r.py) [[Official Code]](https://github.com/guoshnBJTU/ASTGCN-r-pytorch) 19 | - Graph WaveNet[[Paper]](https://arxiv.org/pdf/1906.00121) [[Code]](./models/graph_wavenet.py) [[Official Code]](https://github.com/nnzhan/Graph-WaveNet) 20 | 21 | ### Proposed 22 | 23 | [[Paper]](https://arxiv.org/abs/2401.08727)**MA2GCN: Multi Adjacency relationship Attention Graph Convolutional 24 | Networks for Traffic Prediction using Trajectory data** 25 | 26 | | | | 27 | | :----------------------------------------------------------: | :----------------------------------------------------------: | 28 | | Architecture | Multi Adjacency relationship Attention mechanism | 29 | 30 | ### Related Links 31 | 32 | - GPS2Graph [Link](https://github.com/zachysun/Gps2graph) 33 | 34 | ### Important Notes 35 | 36 | - The Shanghai Taxi Dataset rarely appears in traffic prediction tasks, so the six baselines above can be used as a reference. 37 | - The 'io_matrix' in code corresponds to the 'vehicle entry and exit matrix ' in ma2gcn paper. The 'io_adj' in code corresponds to the $A_{mo}$ in ma2gcn paper. 38 | - The six baselines can't use the mobility matrix as it's specificly used for proposed ma2gcn model. 39 | - Due to the mobility matrix is time sensitive, the dataset needs to be divided in chronological order and cannot be shuffled. For the six baselines, there is no such limit. 40 | - The proposed model is on arxiv.org. Now it can be seen as a simple technical report exploring the use of taxi trajectories for traffic prediction tasks. There's still significant potential for enhancement. 41 | -------------------------------------------------------------------------------- /models/stgcn/models.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from models.stgcn import layers 3 | 4 | 5 | class STGCNChebGraphConv(nn.Module): 6 | # STGCNChebGraphConv contains 'TGTND TGTND TNFF' structure 7 | # ChebGraphConv is the graph convolution from ChebyNet. 8 | # Using the Chebyshev polynomials of the first kind as a graph filter. 9 | 10 | # T: Gated Temporal Convolution Layer (GLU or GTU) 11 | # G: Graph Convolution Layer (ChebGraphConv) 12 | # T: Gated Temporal Convolution Layer (GLU or GTU) 13 | # N: Layer Normolization 14 | # D: Dropout 15 | 16 | # T: Gated Temporal Convolution Layer (GLU or GTU) 17 | # G: Graph Convolution Layer (ChebGraphConv) 18 | # T: Gated Temporal Convolution Layer (GLU or GTU) 19 | # N: Layer Normolization 20 | # D: Dropout 21 | 22 | # T: Gated Temporal Convolution Layer (GLU or GTU) 23 | # N: Layer Normalization 24 | # F: Fully-Connected Layer 25 | # F: Fully-Connected Layer 26 | 27 | def __init__(self, blocks, n_vertex, Kt, Ks, act_func, graph_conv_type, gso, enable_bias, droprate, n_his): 28 | super(STGCNChebGraphConv, self).__init__() 29 | modules = [] 30 | for l in range(len(blocks) - 3): 31 | modules.append(layers.STConvBlock(Kt, Ks, n_vertex, blocks[l][-1], blocks[l + 1], act_func, 32 | graph_conv_type, gso, enable_bias, droprate)) 33 | self.st_blocks = nn.Sequential(*modules) 34 | Ko = n_his - (len(blocks) - 3) * 2 * (Kt - 1) 35 | self.Ko = Ko 36 | if self.Ko > 1: 37 | self.output = layers.OutputBlock(Ko, blocks[-3][-1], blocks[-2], blocks[-1][0], n_vertex, act_func, 38 | enable_bias, droprate) 39 | elif self.Ko == 0: 40 | self.fc1 = nn.Linear(in_features=blocks[-3][-1], out_features=blocks[-2][0], bias=enable_bias) 41 | self.fc2 = nn.Linear(in_features=blocks[-2][0], out_features=blocks[-1][0], bias=enable_bias) 42 | self.relu = nn.ReLU() 43 | self.leaky_relu = nn.LeakyReLU() 44 | self.silu = nn.SiLU() 45 | self.dropout = nn.Dropout(p=droprate) 46 | 47 | def forward(self, x): 48 | x = self.st_blocks(x) 49 | if self.Ko > 1: 50 | x = self.output(x) 51 | elif self.Ko == 0: 52 | x = self.fc1(x.permute(0, 2, 3, 1)) 53 | x = self.relu(x) 54 | x = self.fc2(x).permute(0, 3, 1, 2) 55 | 56 | return x 57 | -------------------------------------------------------------------------------- /main/main_svr.py: -------------------------------------------------------------------------------- 1 | from utils import ProcessData, masked_mae_np, masked_mape_np, masked_rmse_np, StandardScalar_np 2 | import numpy as np 3 | from sklearn.svm import SVR 4 | import yaml 5 | import argparse 6 | 7 | 8 | def load_config(config_path): 9 | with open(config_path, 'r', encoding='utf-8') as f: 10 | config = yaml.safe_load(f) 11 | return config 12 | 13 | 14 | def main(config): 15 | data_config = config['dataset'] 16 | feature_matrix = np.load(data_config['feature_matrix_data_path']) 17 | # shape -> (number of samples, sequence length, number of vertexes, features) 18 | train_data, train_labels, val_data, val_labels, test_data, test_labels = ProcessData(feature_matrix, 19 | data_config['input_seq'], 20 | data_config['output_seq'], 21 | data_config['split_rate']) 22 | 23 | scalar = StandardScalar_np(train_data, axis=(0, 1, 2)) 24 | data_x_sets = [train_data, val_data, test_data] 25 | data_y_sets = [train_labels, val_labels, test_labels] 26 | data_x_normed_sets = [] 27 | data_y_real_sets = [] 28 | data_y_normed_sets = [] 29 | 30 | for data in data_x_sets: 31 | data_normed = scalar.transform(data) 32 | data_x_normed_sets.append(data_normed) 33 | train_x, val_x, test_x = data_x_normed_sets 34 | 35 | for data in data_y_sets: 36 | data_y_real_sets.append(data) 37 | train_real_y, val_real_y, test_real_y = data_y_real_sets 38 | 39 | for data in data_y_sets: 40 | data_normed = scalar.transform(data) 41 | data_y_normed_sets.append(data_normed) 42 | train_normed_y, val_normed_y, test_normed_y = data_y_normed_sets 43 | 44 | X_train = np.einsum('btnf->bnft', train_x) 45 | X_train = X_train.reshape(-1, X_train.shape[-1]) 46 | y_train = np.reshape(train_normed_y, (-1)) 47 | X_test = np.einsum('btnf->bnft', test_x) 48 | X_test = X_test.reshape(-1, X_test.shape[-1]) 49 | y_test = np.reshape(test_real_y, (-1)) 50 | b, t, n, f = test_real_y.shape 51 | 52 | svr = SVR(kernel='rbf') 53 | svr.fit(X_train, y_train) 54 | y_pred = svr.predict(X_test) 55 | y_pred = scalar.inverse_transform(y_pred.reshape(b, t, n, f)).reshape(-1) 56 | 57 | loss = masked_mae_np(y_pred, y_test, 0.0) 58 | mape = masked_mape_np(y_pred, y_test, 0.0) 59 | rmse = masked_rmse_np(y_pred, y_test, 0.0) 60 | 61 | print(loss) 62 | print(mape) 63 | print(rmse) 64 | 65 | 66 | if __name__ == '__main__': 67 | parser = argparse.ArgumentParser() 68 | parser.add_argument('--config_path', type=str, default='../config/config_svr.yaml', 69 | help='Config path') 70 | args = parser.parse_args() 71 | configs = load_config(args.config_path) 72 | main(configs) 73 | -------------------------------------------------------------------------------- /main/main_arima.py: -------------------------------------------------------------------------------- 1 | from utils import ProcessData, masked_mae_np, masked_mape_np, masked_rmse_np 2 | from statsmodels.tsa.arima.model import ARIMA 3 | import numpy as np 4 | import yaml 5 | import argparse 6 | 7 | 8 | def load_config(config_path): 9 | with open(config_path, 'r', encoding='utf-8') as f: 10 | config = yaml.safe_load(f) 11 | return config 12 | 13 | 14 | def main(config): 15 | data_config = config['dataset'] 16 | feature_matrix = np.load(data_config['feature_matrix_data_path']) 17 | # shape -> (number of samples, sequence length, number of vertexes, features) 18 | train_data, train_labels, val_data, val_labels, test_data, test_labels = ProcessData(feature_matrix, 19 | data_config['input_seq'], 20 | data_config['output_seq'], 21 | data_config['split_rate']) 22 | 23 | data = test_data 24 | labels = test_labels 25 | reshaped_data = data.transpose(0, 2, 1, 3).reshape(-1, data.shape[1], data.shape[3]) 26 | reshaped_data_1 = reshaped_data[..., 0] 27 | reshaped_data_2 = reshaped_data[..., 1] 28 | reshaped_labels = labels.transpose(0, 2, 1, 3).reshape(-1, labels.shape[1], labels.shape[3]) 29 | reshaped_labels_1 = reshaped_labels[..., 0] 30 | reshaped_labels_2 = reshaped_labels[..., 1] 31 | 32 | preds = [] 33 | for i in range(reshaped_data_1.shape[0]): 34 | node_feature_data = reshaped_data_1[i] 35 | model = ARIMA(node_feature_data, order=(1, 1, 1)) 36 | model_fit = model.fit() 37 | y_pred = model_fit.forecast(steps=1)[0] 38 | preds.append(y_pred) 39 | 40 | preds = np.array(preds) 41 | preds = preds.reshape(len(preds), 1) 42 | loss1 = masked_mae_np(preds, reshaped_labels_1, 0.0) 43 | mape1 = masked_mape_np(preds, reshaped_labels_1, 0.0) 44 | rmse1 = masked_rmse_np(preds, reshaped_labels_1, 0.0) 45 | 46 | preds2 = [] 47 | for i in range(reshaped_data_2.shape[0]): 48 | node_feature_data = reshaped_data_2[i] 49 | model = ARIMA(node_feature_data, order=(1, 1, 1)) 50 | model_fit = model.fit() 51 | y_pred = model_fit.forecast(steps=1)[0] 52 | preds2.append(y_pred) 53 | 54 | preds2 = np.array(preds) 55 | preds2 = preds.reshape(len(preds), 1) 56 | loss2 = masked_mae_np(preds, reshaped_labels_2, 0.0) 57 | mape2 = masked_mape_np(preds, reshaped_labels_2, 0.0) 58 | rmse2 = masked_rmse_np(preds, reshaped_labels_2, 0.0) 59 | 60 | print((loss1 + loss2) / 2) 61 | print((mape1 + mape2) / 2) 62 | print((rmse1 + rmse2) / 2) 63 | 64 | 65 | if __name__ == '__main__': 66 | parser = argparse.ArgumentParser() 67 | parser.add_argument('--config_path', type=str, default='../config/config_arima.yaml', 68 | help='Config path') 69 | args = parser.parse_args() 70 | configs = load_config(args.config_path) 71 | main(configs) 72 | -------------------------------------------------------------------------------- /trainer/trainer_astgcn.py: -------------------------------------------------------------------------------- 1 | from utils import masked_mae, masked_mape, masked_rmse 2 | import numpy as np 3 | import torch.nn as nn 4 | import torch 5 | 6 | 7 | class trainer: 8 | def __init__(self, model, optimizer, scaler, theta): 9 | self.model = model 10 | self.criterion = masked_mae 11 | self.optimizer = optimizer 12 | self.scaler = scaler 13 | self.theta = theta 14 | for p in model.parameters(): 15 | if p.dim() > 1: 16 | nn.init.xavier_uniform_(p) 17 | else: 18 | nn.init.uniform_(p) 19 | 20 | def train(self, x, y_real): 21 | self.optimizer.zero_grad() 22 | 23 | output = self.model(x) 24 | y_pred = self.scaler.inverse_transform(output) 25 | loss = self.theta * self.criterion(y_pred[..., 0], y_real[..., 0], 0.0) + \ 26 | (1 - self.theta) * self.criterion(y_pred[..., 1], y_real[..., 1], 0.0) 27 | 28 | loss.backward() 29 | self.optimizer.step() 30 | 31 | mape = masked_mape(y_pred, y_real, 0.0).item() 32 | rmse = masked_rmse(y_pred, y_real, 0.0).item() 33 | return loss.item(), mape, rmse 34 | 35 | def train_one_epoch(self, train_loader): 36 | self.model.train() 37 | loss_inner_epoch = [] 38 | mape_inner_epoch = [] 39 | rmse_inner_epoch = [] 40 | for batch_idx, (x, y_real) in enumerate(train_loader): 41 | 42 | loss, mape, rmse = self.train(x, y_real) 43 | loss_inner_epoch.append(loss) 44 | mape_inner_epoch.append(mape) 45 | rmse_inner_epoch.append(rmse) 46 | 47 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 48 | np.array(rmse_inner_epoch).mean() 49 | 50 | def eval(self, x, y_real): 51 | output = self.model(x) 52 | y_pred = self.scaler.inverse_transform(output) 53 | 54 | loss = self.theta * self.criterion(y_pred[..., 0], y_real[..., 0], 0.0) + \ 55 | (1 - self.theta) * self.criterion(y_pred[..., 1], y_real[..., 1], 0.0) 56 | 57 | mape = masked_mape(y_pred, y_real, 0.0).item() 58 | rmse = masked_rmse(y_pred, y_real, 0.0).item() 59 | return loss.item(), mape, rmse 60 | 61 | def eval_one_epoch(self, val_loader): 62 | self.model.eval() 63 | loss_inner_epoch = [] 64 | mape_inner_epoch = [] 65 | rmse_inner_epoch = [] 66 | for batch_idx, (x, y_real) in enumerate(val_loader): 67 | loss, mape, rmse = self.eval(x, y_real) 68 | loss_inner_epoch.append(loss) 69 | mape_inner_epoch.append(mape) 70 | rmse_inner_epoch.append(rmse) 71 | 72 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 73 | np.array(rmse_inner_epoch).mean() 74 | 75 | def test(self, test_loader): 76 | loss_inner_epoch = [] 77 | mape_inner_epoch = [] 78 | rmse_inner_epoch = [] 79 | with torch.no_grad(): 80 | for batch_idx, (x, y_real) in enumerate(test_loader): 81 | loss, mape, rmse = self.eval(x, y_real) 82 | loss_inner_epoch.append(loss) 83 | mape_inner_epoch.append(mape) 84 | rmse_inner_epoch.append(rmse) 85 | 86 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 87 | np.array(rmse_inner_epoch).mean() 88 | -------------------------------------------------------------------------------- /trainer/trainer_graph_wavenet.py: -------------------------------------------------------------------------------- 1 | from utils import masked_mae, masked_mape, masked_rmse 2 | import numpy as np 3 | import torch.nn as nn 4 | import torch 5 | 6 | 7 | class trainer: 8 | def __init__(self, model, optimizer, scaler, theta): 9 | self.model = model 10 | self.criterion = masked_mae 11 | self.optimizer = optimizer 12 | self.scaler = scaler 13 | self.theta = theta 14 | for p in model.parameters(): 15 | if p.dim() > 1: 16 | nn.init.xavier_uniform_(p) 17 | else: 18 | nn.init.uniform_(p) 19 | 20 | def train(self, x, y_real): 21 | self.optimizer.zero_grad() 22 | 23 | output = self.model(x) 24 | y_pred = self.scaler.inverse_transform(output) 25 | loss = self.theta * self.criterion(y_pred[..., 0], y_real[..., 0], 0.0) + \ 26 | (1 - self.theta) * self.criterion(y_pred[..., 1], y_real[..., 1], 0.0) 27 | 28 | loss.backward() 29 | self.optimizer.step() 30 | 31 | mape = masked_mape(y_pred, y_real, 0.0).item() 32 | rmse = masked_rmse(y_pred, y_real, 0.0).item() 33 | return loss.item(), mape, rmse 34 | 35 | def train_one_epoch(self, train_loader): 36 | self.model.train() 37 | loss_inner_epoch = [] 38 | mape_inner_epoch = [] 39 | rmse_inner_epoch = [] 40 | for batch_idx, (x, y_real) in enumerate(train_loader): 41 | 42 | loss, mape, rmse = self.train(x, y_real) 43 | loss_inner_epoch.append(loss) 44 | mape_inner_epoch.append(mape) 45 | rmse_inner_epoch.append(rmse) 46 | 47 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 48 | np.array(rmse_inner_epoch).mean() 49 | 50 | def eval(self, x, y_real): 51 | output = self.model(x) 52 | y_pred = self.scaler.inverse_transform(output) 53 | 54 | loss = self.theta * self.criterion(y_pred[..., 0], y_real[..., 0], 0.0) + \ 55 | (1 - self.theta) * self.criterion(y_pred[..., 1], y_real[..., 1], 0.0) 56 | 57 | mape = masked_mape(y_pred, y_real, 0.0).item() 58 | rmse = masked_rmse(y_pred, y_real, 0.0).item() 59 | return loss.item(), mape, rmse 60 | 61 | def eval_one_epoch(self, val_loader): 62 | self.model.eval() 63 | loss_inner_epoch = [] 64 | mape_inner_epoch = [] 65 | rmse_inner_epoch = [] 66 | for batch_idx, (x, y_real) in enumerate(val_loader): 67 | loss, mape, rmse = self.eval(x, y_real) 68 | loss_inner_epoch.append(loss) 69 | mape_inner_epoch.append(mape) 70 | rmse_inner_epoch.append(rmse) 71 | 72 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 73 | np.array(rmse_inner_epoch).mean() 74 | 75 | def test(self, test_loader): 76 | loss_inner_epoch = [] 77 | mape_inner_epoch = [] 78 | rmse_inner_epoch = [] 79 | with torch.no_grad(): 80 | for batch_idx, (x, y_real) in enumerate(test_loader): 81 | loss, mape, rmse = self.eval(x, y_real) 82 | loss_inner_epoch.append(loss) 83 | mape_inner_epoch.append(mape) 84 | rmse_inner_epoch.append(rmse) 85 | 86 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 87 | np.array(rmse_inner_epoch).mean() 88 | 89 | -------------------------------------------------------------------------------- /trainer/trainer_stgcn.py: -------------------------------------------------------------------------------- 1 | from utils import masked_mae, masked_mape, masked_rmse 2 | import numpy as np 3 | import torch.nn as nn 4 | import torch 5 | 6 | 7 | class trainer: 8 | def __init__(self, model, optimizer, scaler, theta): 9 | self.model = model 10 | self.criterion = masked_mae 11 | self.optimizer = optimizer 12 | self.scaler = scaler 13 | self.theta = theta 14 | for p in model.parameters(): 15 | if p.dim() > 1: 16 | nn.init.xavier_uniform_(p) 17 | else: 18 | nn.init.uniform_(p) 19 | 20 | def train(self, x, y_real): 21 | self.optimizer.zero_grad() 22 | 23 | output = self.model(x) 24 | y_pred = self.scaler.inverse_transform(output.permute(0, 2, 3, 1)) 25 | loss = self.theta * self.criterion(y_pred[..., 0], y_real[..., 0], 0.0) + \ 26 | (1 - self.theta) * self.criterion(y_pred[..., 1], y_real[..., 1], 0.0) 27 | 28 | loss.backward() 29 | self.optimizer.step() 30 | 31 | mape = masked_mape(y_pred, y_real, 0.0).item() 32 | rmse = masked_rmse(y_pred, y_real, 0.0).item() 33 | return loss.item(), mape, rmse 34 | 35 | def train_one_epoch(self, train_loader): 36 | self.model.train() 37 | loss_inner_epoch = [] 38 | mape_inner_epoch = [] 39 | rmse_inner_epoch = [] 40 | for batch_idx, (x, y_real) in enumerate(train_loader): 41 | 42 | loss, mape, rmse = self.train(x, y_real) 43 | loss_inner_epoch.append(loss) 44 | mape_inner_epoch.append(mape) 45 | rmse_inner_epoch.append(rmse) 46 | 47 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 48 | np.array(rmse_inner_epoch).mean() 49 | 50 | def eval(self, x, y_real): 51 | output = self.model(x) 52 | y_pred = self.scaler.inverse_transform(output.permute(0, 2, 3, 1)) 53 | 54 | loss = self.theta * self.criterion(y_pred[..., 0], y_real[..., 0], 0.0) + \ 55 | (1 - self.theta) * self.criterion(y_pred[..., 1], y_real[..., 1], 0.0) 56 | 57 | mape = masked_mape(y_pred, y_real, 0.0).item() 58 | rmse = masked_rmse(y_pred, y_real, 0.0).item() 59 | return loss.item(), mape, rmse 60 | 61 | def eval_one_epoch(self, val_loader): 62 | self.model.eval() 63 | loss_inner_epoch = [] 64 | mape_inner_epoch = [] 65 | rmse_inner_epoch = [] 66 | for batch_idx, (x, y_real) in enumerate(val_loader): 67 | loss, mape, rmse = self.eval(x, y_real) 68 | loss_inner_epoch.append(loss) 69 | mape_inner_epoch.append(mape) 70 | rmse_inner_epoch.append(rmse) 71 | 72 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 73 | np.array(rmse_inner_epoch).mean() 74 | 75 | def test(self, test_loader): 76 | loss_inner_epoch = [] 77 | mape_inner_epoch = [] 78 | rmse_inner_epoch = [] 79 | with torch.no_grad(): 80 | for batch_idx, (x, y_real) in enumerate(test_loader): 81 | loss, mape, rmse = self.eval(x, y_real) 82 | loss_inner_epoch.append(loss) 83 | mape_inner_epoch.append(mape) 84 | rmse_inner_epoch.append(rmse) 85 | 86 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 87 | np.array(rmse_inner_epoch).mean() 88 | -------------------------------------------------------------------------------- /trainer/trainer_lstm.py: -------------------------------------------------------------------------------- 1 | from utils import masked_mae, masked_mape, masked_rmse 2 | import numpy as np 3 | import torch.nn as nn 4 | import torch 5 | 6 | 7 | class trainer: 8 | def __init__(self, model, optimizer, scaler, theta): 9 | self.model = model 10 | self.criterion = masked_mae 11 | self.optimizer = optimizer 12 | self.scaler = scaler 13 | self.theta = theta 14 | for p in model.parameters(): 15 | if p.dim() > 1: 16 | nn.init.xavier_uniform_(p) 17 | else: 18 | nn.init.uniform_(p) 19 | 20 | def train(self, x, y_real): 21 | self.optimizer.zero_grad() 22 | 23 | output = self.model(x) 24 | y_pred = self.scaler.inverse_transform(output) 25 | y_pred = y_pred.unsqueeze(1) 26 | loss = self.theta * self.criterion(y_pred[..., 0], y_real[..., 0], 0.0) + \ 27 | (1 - self.theta) * self.criterion(y_pred[..., 1], y_real[..., 1], 0.0) 28 | 29 | loss.backward() 30 | self.optimizer.step() 31 | 32 | mape = masked_mape(y_pred, y_real, 0.0).item() 33 | rmse = masked_rmse(y_pred, y_real, 0.0).item() 34 | return loss.item(), mape, rmse 35 | 36 | def train_one_epoch(self, train_loader): 37 | self.model.train() 38 | loss_inner_epoch = [] 39 | mape_inner_epoch = [] 40 | rmse_inner_epoch = [] 41 | for batch_idx, (x, y_real) in enumerate(train_loader): 42 | 43 | loss, mape, rmse = self.train(x, y_real) 44 | loss_inner_epoch.append(loss) 45 | mape_inner_epoch.append(mape) 46 | rmse_inner_epoch.append(rmse) 47 | 48 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 49 | np.array(rmse_inner_epoch).mean() 50 | 51 | def eval(self, x, y_real): 52 | output = self.model(x) 53 | y_pred = self.scaler.inverse_transform(output) 54 | y_pred = y_pred.unsqueeze(1) 55 | loss = self.theta * self.criterion(y_pred[..., 0], y_real[..., 0], 0.0) + \ 56 | (1 - self.theta) * self.criterion(y_pred[..., 1], y_real[..., 1], 0.0) 57 | 58 | mape = masked_mape(y_pred, y_real, 0.0).item() 59 | rmse = masked_rmse(y_pred, y_real, 0.0).item() 60 | return loss.item(), mape, rmse 61 | 62 | def eval_one_epoch(self, val_loader): 63 | self.model.eval() 64 | loss_inner_epoch = [] 65 | mape_inner_epoch = [] 66 | rmse_inner_epoch = [] 67 | for batch_idx, (x, y_real) in enumerate(val_loader): 68 | loss, mape, rmse = self.eval(x, y_real) 69 | loss_inner_epoch.append(loss) 70 | mape_inner_epoch.append(mape) 71 | rmse_inner_epoch.append(rmse) 72 | 73 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 74 | np.array(rmse_inner_epoch).mean() 75 | 76 | def test(self, test_loader): 77 | loss_inner_epoch = [] 78 | mape_inner_epoch = [] 79 | rmse_inner_epoch = [] 80 | with torch.no_grad(): 81 | for batch_idx, (x, y_real) in enumerate(test_loader): 82 | loss, mape, rmse = self.eval(x, y_real) 83 | loss_inner_epoch.append(loss) 84 | mape_inner_epoch.append(mape) 85 | rmse_inner_epoch.append(rmse) 86 | 87 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 88 | np.array(rmse_inner_epoch).mean() -------------------------------------------------------------------------------- /trainer/trainer_ma2gcn.py: -------------------------------------------------------------------------------- 1 | from utils import masked_mae, masked_mape, masked_rmse 2 | import numpy as np 3 | import torch.nn as nn 4 | import torch 5 | 6 | 7 | class trainer: 8 | def __init__(self, model, optimizer, scaler, theta, io_range, io_adj): 9 | self.model = model 10 | self.criterion = masked_mae 11 | self.optimizer = optimizer 12 | self.scaler = scaler 13 | self.theta = theta 14 | self.io_range = io_range 15 | self.io_adj = io_adj 16 | for p in model.parameters(): 17 | if p.dim() > 1: 18 | nn.init.xavier_uniform_(p) 19 | else: 20 | nn.init.uniform_(p) 21 | 22 | def train(self, x, y_real, cur_io_adj): 23 | self.optimizer.zero_grad() 24 | 25 | output, cur_io_adj = self.model((x, cur_io_adj)) 26 | y_pred = self.scaler.inverse_transform(output) 27 | loss = self.theta * self.criterion(y_pred[..., 0], y_real[..., 0], 0.0) + \ 28 | (1 - self.theta) * self.criterion(y_pred[..., 1], y_real[..., 1], 0.0) 29 | 30 | loss.backward() 31 | self.optimizer.step() 32 | 33 | mape = masked_mape(y_pred, y_real, 0.0).item() 34 | rmse = masked_rmse(y_pred, y_real, 0.0).item() 35 | return loss.item(), mape, rmse 36 | 37 | def train_one_epoch(self, train_loader): 38 | self.model.train() 39 | loss_inner_epoch = [] 40 | mape_inner_epoch = [] 41 | rmse_inner_epoch = [] 42 | for batch_idx, (x, y_real) in enumerate(train_loader): 43 | batch_num, _, _, _ = x.shape 44 | batch_start = batch_idx * batch_num 45 | batch_end = (batch_idx + 1) * batch_num 46 | io_start = np.digitize(batch_start, self.io_range) - 1 47 | io_end = np.digitize(batch_end, self.io_range) - 1 48 | cur_io_adj = (self.io_adj[io_start] + self.io_adj[io_end]) / 2 49 | loss, mape, rmse = self.train(x, y_real, cur_io_adj) 50 | loss_inner_epoch.append(loss) 51 | mape_inner_epoch.append(mape) 52 | rmse_inner_epoch.append(rmse) 53 | 54 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 55 | np.array(rmse_inner_epoch).mean() 56 | 57 | def eval(self, x, y_real, cur_io_adj): 58 | output, cur_io_adj = self.model((x, cur_io_adj)) 59 | y_pred = self.scaler.inverse_transform(output) 60 | 61 | loss = self.theta * self.criterion(y_pred[..., 0], y_real[..., 0], 0.0) + \ 62 | (1 - self.theta) * self.criterion(y_pred[..., 1], y_real[..., 1], 0.0) 63 | 64 | mape = masked_mape(y_pred, y_real, 0.0).item() 65 | rmse = masked_rmse(y_pred, y_real, 0.0).item() 66 | return loss.item(), mape, rmse 67 | 68 | def eval_one_epoch(self, val_loader, cur_io_adj): 69 | self.model.eval() 70 | loss_inner_epoch = [] 71 | mape_inner_epoch = [] 72 | rmse_inner_epoch = [] 73 | for batch_idx, (x, y_real) in enumerate(val_loader): 74 | loss, mape, rmse = self.eval(x, y_real, cur_io_adj) 75 | loss_inner_epoch.append(loss) 76 | mape_inner_epoch.append(mape) 77 | rmse_inner_epoch.append(rmse) 78 | 79 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 80 | np.array(rmse_inner_epoch).mean() 81 | 82 | def test(self, test_loader, cur_io_adj): 83 | loss_inner_epoch = [] 84 | mape_inner_epoch = [] 85 | rmse_inner_epoch = [] 86 | with torch.no_grad(): 87 | for batch_idx, (x, y_real) in enumerate(test_loader): 88 | loss, mape, rmse = self.eval(x, y_real, cur_io_adj) 89 | loss_inner_epoch.append(loss) 90 | mape_inner_epoch.append(mape) 91 | rmse_inner_epoch.append(rmse) 92 | 93 | return np.array(loss_inner_epoch).mean(), np.array(mape_inner_epoch).mean(), \ 94 | np.array(rmse_inner_epoch).mean() 95 | -------------------------------------------------------------------------------- /main/main_lstm.py: -------------------------------------------------------------------------------- 1 | from utils import ProcessData, StandardScalar 2 | from models.lstm import LSTM 3 | import yaml 4 | import argparse 5 | import numpy as np 6 | import torch 7 | import torch.optim as optim 8 | from trainer.trainer_lstm import trainer 9 | from torch.utils.data import DataLoader, TensorDataset 10 | 11 | 12 | def load_config(config_path): 13 | with open(config_path, 'r', encoding='utf-8') as f: 14 | config = yaml.safe_load(f) 15 | return config 16 | 17 | 18 | def main(config): 19 | # ------------------------ Data loading ------------------------ 20 | data_config = config['dataset'] 21 | optim_config = config['optimize'] 22 | trainer_config = config['trainer'] 23 | feature_matrix = np.load(data_config['feature_matrix_data_path']) 24 | # shape -> (number of samples, sequence length, number of vertexes, features) 25 | train_data, train_labels, val_data, val_labels, test_data, test_labels = ProcessData(feature_matrix, 26 | data_config['input_seq'], 27 | data_config['output_seq'], 28 | data_config['split_rate']) 29 | # data scalar 30 | device = data_config['device'] 31 | scalar = StandardScalar(train_data, device, axis=(0, 1, 2)) 32 | data_x_sets = [train_data, val_data, test_data] 33 | data_y_sets = [train_labels, val_labels, test_labels] 34 | data_x_normed_sets = [] 35 | data_y_real_sets = [] 36 | 37 | for data in data_x_sets: 38 | data_normed = scalar.transform(data) 39 | data_normed = data_normed.transpose(0, 2, 1, 3) 40 | data_normed = data_normed.reshape(-1, data_normed.shape[2], data_normed.shape[3]) 41 | data_normed = torch.Tensor(data_normed).to(device) 42 | data_x_normed_sets.append(data_normed) 43 | train_x, val_x, test_x = data_x_normed_sets 44 | 45 | for data in data_y_sets: 46 | data = data.transpose(0, 2, 1, 3) 47 | data = data.reshape(-1, data.shape[2], data.shape[3]) 48 | data_real = torch.Tensor(data).to(device) 49 | data_y_real_sets.append(data_real) 50 | train_real_y, val_real_y, test_real_y = data_y_real_sets 51 | 52 | data_final_list = [train_x, train_real_y, val_x, val_real_y, test_x, test_real_y] 53 | 54 | # get dataloader 55 | data_loaders = [] 56 | for i in range(0, len(data_final_list), 2): 57 | dataset = TensorDataset(data_final_list[i], data_final_list[i + 1]) 58 | data_loader = DataLoader(dataset, batch_size=data_config['batch_size']) 59 | data_loaders.append(data_loader) 60 | 61 | # ------------------------ Model setting ------------------------ 62 | 63 | model = LSTM(trainer_config['in_dim'], trainer_config['hidden_dim'], trainer_config['out_dim']) 64 | model.to(device) 65 | 66 | if optim_config['name'] == 'Adam': 67 | optimizer = optim.Adam(model.parameters(), 68 | lr=optim_config['learning_rate'], 69 | weight_decay=optim_config['weight_decay']) 70 | 71 | elif optim_config['name'] == 'AdamW': 72 | optimizer = optim.AdamW(model.parameters(), 73 | lr=optim_config['learning_rate'], 74 | weight_decay=optim_config['weight_decay']) 75 | 76 | elif optim_config['name'] == 'RMSprop': 77 | optimizer = optim.RMSprop(model.parameters(), 78 | lr=optim_config['learning_rate'], 79 | weight_decay=optim_config['weight_decay']) 80 | else: 81 | raise NotImplementedError(f'ERROR: The optimizer {optim_config["name"]} is not implemented.') 82 | 83 | engine = trainer(model, optimizer, scalar, trainer_config['theta']) 84 | # --------------------------- Train/Validation ------------------------- 85 | print('Training...') 86 | train_loss_epochs = [] 87 | train_mape_epochs = [] 88 | train_rmse_epochs = [] 89 | val_loss_epochs = [] 90 | val_mape_epochs = [] 91 | val_rmse_epochs = [] 92 | for epoch in range(trainer_config['epoch']): 93 | loss, mape, rmse = engine.train_one_epoch(data_loaders[0]) 94 | train_loss_epochs.append(loss) 95 | train_mape_epochs.append(mape) 96 | train_rmse_epochs.append(rmse) 97 | 98 | print(f'epoch:{epoch}, train_loss(MAE):{train_loss_epochs[epoch]}, ' 99 | f'train_MAPE:{train_mape_epochs[epoch]}, train_RMSE:{train_rmse_epochs[epoch]}') 100 | 101 | loss, mape, rmse = engine.eval_one_epoch(data_loaders[1]) 102 | val_loss_epochs.append(loss) 103 | val_mape_epochs.append(mape) 104 | val_rmse_epochs.append(rmse) 105 | 106 | print(f'val_loss(MAE):{val_loss_epochs[epoch]}, ' 107 | f'val_MAPE:{val_mape_epochs[epoch]}, val_RMSE:{val_rmse_epochs[epoch]}') 108 | 109 | torch.save(model.state_dict(), trainer_config['model_path'] + '/' + 'lstm_pretrained' + '.pth') 110 | print('Model saved!') 111 | 112 | # --------------------------- Test ------------------------- 113 | print('Testing...') 114 | test_loss, test_mape, test_rmse = engine.test(data_loaders[2]) 115 | print(f'test_loss(MAE):{test_loss}, test_mape:{test_mape}, test_rmse:{test_rmse}') 116 | 117 | 118 | if __name__ == '__main__': 119 | parser = argparse.ArgumentParser() 120 | parser.add_argument('--config_path', type=str, default='../config/config_lstm.yaml', 121 | help='Config path') 122 | args = parser.parse_args() 123 | configs = load_config(args.config_path) 124 | main(configs) 125 | -------------------------------------------------------------------------------- /main/main_astgcn.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | import argparse 3 | import numpy as np 4 | import pandas as pd 5 | import torch 6 | import torch.optim as optim 7 | from models.ASTGCN_r import make_model 8 | from utils import ProcessData, StandardScalar 9 | from torch.utils.data import DataLoader, TensorDataset 10 | from trainer.trainer_astgcn import trainer 11 | 12 | 13 | def load_config(config_path): 14 | with open(config_path, 'r', encoding='utf-8') as f: 15 | config = yaml.safe_load(f) 16 | return config 17 | 18 | 19 | def main(config): 20 | # ------------------------ Data loading ------------------------ 21 | data_config = config['dataset'] 22 | optim_config = config['optimize'] 23 | trainer_config = config['trainer'] 24 | feature_matrix = np.load(data_config['feature_matrix_data_path']) 25 | samples = feature_matrix.shape[0] 26 | # shape -> (number of samples, sequence length, number of vertexes, features) 27 | train_data, train_labels, val_data, val_labels, test_data, test_labels = ProcessData(feature_matrix, 28 | data_config['input_seq'], 29 | data_config['output_seq'], 30 | data_config['split_rate']) 31 | # data scalar 32 | device = data_config['device'] 33 | scalar = StandardScalar(train_data, device, axis=(0, 1, 2)) 34 | data_x_sets = [train_data, val_data, test_data] 35 | data_y_sets = [train_labels, val_labels, test_labels] 36 | data_x_normed_sets = [] 37 | data_y_real_sets = [] 38 | for data in data_x_sets: 39 | data_normed = scalar.transform(data) 40 | data_normed = np.einsum('btnf->bnft', data_normed) 41 | data_normed = torch.Tensor(data_normed).to(device) 42 | data_x_normed_sets.append(data_normed) 43 | train_x, val_x, test_x = data_x_normed_sets 44 | for data in data_y_sets: 45 | data_real = torch.Tensor(data).to(device) 46 | data_y_real_sets.append(data_real) 47 | train_real_y, val_real_y, test_real_y = data_y_real_sets 48 | data_final_list = [train_x, train_real_y, val_x, val_real_y, test_x, test_real_y] 49 | 50 | # get dataloader 51 | data_loaders = [] 52 | for i in range(0, len(data_final_list), 2): 53 | dataset = TensorDataset(data_final_list[i], data_final_list[i + 1]) 54 | data_loader = DataLoader(dataset, batch_size=data_config['batch_size']) 55 | data_loaders.append(data_loader) 56 | 57 | # get adjacency matrix 58 | origin_adj = pd.read_csv(data_config['adj_data_path'], header=None) 59 | origin_adj = torch.Tensor(origin_adj.values).to(device) 60 | 61 | # ------------------------ Model setting ------------------------ 62 | 63 | model = make_model(data_config['device'], trainer_config['nb_block'], trainer_config['in_channels'], 64 | trainer_config['K'], trainer_config['nb_chev_filter'], trainer_config['nb_time_filter'], 65 | trainer_config['num_of_hours'], origin_adj, data_config['output_seq'], 66 | data_config['input_seq'], data_config['num_of_vertices']) 67 | model.to(device) 68 | 69 | if optim_config['name'] == 'Adam': 70 | optimizer = optim.Adam(model.parameters(), 71 | lr=optim_config['learning_rate'], 72 | weight_decay=optim_config['weight_decay']) 73 | 74 | elif optim_config['name'] == 'AdamW': 75 | optimizer = optim.AdamW(model.parameters(), 76 | lr=optim_config['learning_rate'], 77 | weight_decay=optim_config['weight_decay']) 78 | 79 | elif optim_config['name'] == 'RMSprop': 80 | optimizer = optim.RMSprop(model.parameters(), 81 | lr=optim_config['learning_rate'], 82 | weight_decay=optim_config['weight_decay']) 83 | else: 84 | raise NotImplementedError(f'ERROR: The optimizer {optim_config["name"]} is not implemented.') 85 | 86 | engine = trainer(model, optimizer, scalar, trainer_config['theta']) 87 | # --------------------------- Train/Validation ------------------------- 88 | print('Training...') 89 | train_loss_epochs = [] 90 | train_mape_epochs = [] 91 | train_rmse_epochs = [] 92 | val_loss_epochs = [] 93 | val_mape_epochs = [] 94 | val_rmse_epochs = [] 95 | for epoch in range(trainer_config['epoch']): 96 | loss, mape, rmse = engine.train_one_epoch(data_loaders[0]) 97 | train_loss_epochs.append(loss) 98 | train_mape_epochs.append(mape) 99 | train_rmse_epochs.append(rmse) 100 | 101 | print(f'epoch:{epoch}, train_loss(MAE):{train_loss_epochs[epoch]}, ' 102 | f'train_MAPE:{train_mape_epochs[epoch]}, train_RMSE:{train_rmse_epochs[epoch]}') 103 | 104 | loss, mape, rmse = engine.eval_one_epoch(data_loaders[1]) 105 | val_loss_epochs.append(loss) 106 | val_mape_epochs.append(mape) 107 | val_rmse_epochs.append(rmse) 108 | 109 | print(f'val_loss(MAE):{val_loss_epochs[epoch]}, ' 110 | f'val_MAPE:{val_mape_epochs[epoch]}, val_RMSE:{val_rmse_epochs[epoch]}') 111 | 112 | torch.save(model.state_dict(), trainer_config['model_path'] + '/' + 'astgcn_pretrained' + '.pth') 113 | print('Model saved!') 114 | 115 | # --------------------------- Test ------------------------- 116 | print('Testing...') 117 | test_loss, test_mape, test_rmse = engine.test(data_loaders[2]) 118 | print(f'test_loss(MAE):{test_loss}, test_mape:{test_mape}, test_rmse:{test_rmse}') 119 | 120 | 121 | if __name__ == '__main__': 122 | parser = argparse.ArgumentParser() 123 | parser.add_argument('--config_path', type=str, default='../config/config_astgcn.yaml', 124 | help='Config path') 125 | args = parser.parse_args() 126 | configs = load_config(args.config_path) 127 | main(configs) 128 | -------------------------------------------------------------------------------- /main/main_graph_wavenet.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | import argparse 3 | import numpy as np 4 | import pandas as pd 5 | import torch 6 | import torch.optim as optim 7 | from models.graph_wavenet import gwnet 8 | from utils import ProcessData, StandardScalar 9 | from torch.utils.data import DataLoader, TensorDataset 10 | from utils import masked_mae 11 | from trainer.trainer_graph_wavenet import trainer 12 | 13 | 14 | def load_config(config_path): 15 | with open(config_path, 'r', encoding='utf-8') as f: 16 | config = yaml.safe_load(f) 17 | return config 18 | 19 | 20 | def main(config): 21 | # ------------------------ Data loading ------------------------ 22 | data_config = config['dataset'] 23 | optim_config = config['optimize'] 24 | trainer_config = config['trainer'] 25 | feature_matrix = np.load(data_config['feature_matrix_data_path']) 26 | samples = feature_matrix.shape[0] 27 | # shape -> (number of samples, sequence length, number of vertexes, features) 28 | train_data, train_labels, val_data, val_labels, test_data, test_labels = ProcessData(feature_matrix, 29 | data_config['input_seq'], 30 | data_config['output_seq'], 31 | data_config['split_rate']) 32 | # data scalar 33 | device = data_config['device'] 34 | scalar = StandardScalar(train_data, device, axis=(0, 1, 2)) 35 | data_x_sets = [train_data, val_data, test_data] 36 | data_y_sets = [train_labels, val_labels, test_labels] 37 | data_x_normed_sets = [] 38 | data_y_real_sets = [] 39 | for data in data_x_sets: 40 | data_normed = scalar.transform(data) 41 | data_normed = np.einsum('btnf->bfnt', data_normed) 42 | data_normed = torch.Tensor(data_normed).to(device) 43 | data_x_normed_sets.append(data_normed) 44 | train_x, val_x, test_x = data_x_normed_sets 45 | for data in data_y_sets: 46 | data_real = torch.Tensor(data).to(device) 47 | data_y_real_sets.append(data_real) 48 | train_real_y, val_real_y, test_real_y = data_y_real_sets 49 | data_final_list = [train_x, train_real_y, val_x, val_real_y, test_x, test_real_y] 50 | 51 | # get dataloader 52 | data_loaders = [] 53 | for i in range(0, len(data_final_list), 2): 54 | dataset = TensorDataset(data_final_list[i], data_final_list[i + 1]) 55 | data_loader = DataLoader(dataset, batch_size=data_config['batch_size']) 56 | data_loaders.append(data_loader) 57 | 58 | # get adjacency matrix 59 | origin_adj = pd.read_csv(data_config['adj_data_path'], header=None) 60 | origin_adj = torch.Tensor(origin_adj.values).to(device) 61 | 62 | # ------------------------ Model setting ------------------------ 63 | supports = [origin_adj] 64 | if trainer_config['randomadj']: 65 | adjinit = None 66 | else: 67 | adjinit = supports[0] 68 | if trainer_config['aptonly']: 69 | supports = None 70 | 71 | model = gwnet(device, data_config['num_of_vertices'], trainer_config['dropout'], supports, 72 | trainer_config['gcn_bool'], trainer_config['addaptadj'], adjinit, 73 | trainer_config['in_dim'], trainer_config['out_dim'], trainer_config['nhid'], 74 | trainer_config['nhid'], trainer_config['nhid'] * 8, trainer_config['nhid'] * 16) 75 | model.to(device) 76 | 77 | if optim_config['name'] == 'Adam': 78 | optimizer = optim.Adam(model.parameters(), 79 | lr=optim_config['learning_rate'], 80 | weight_decay=optim_config['weight_decay']) 81 | 82 | elif optim_config['name'] == 'AdamW': 83 | optimizer = optim.AdamW(model.parameters(), 84 | lr=optim_config['learning_rate'], 85 | weight_decay=optim_config['weight_decay']) 86 | 87 | elif optim_config['name'] == 'RMSprop': 88 | optimizer = optim.RMSprop(model.parameters(), 89 | lr=optim_config['learning_rate'], 90 | weight_decay=optim_config['weight_decay']) 91 | else: 92 | raise NotImplementedError(f'ERROR: The optimizer {optim_config["name"]} is not implemented.') 93 | 94 | engine = trainer(model, optimizer, scalar, trainer_config['theta']) 95 | # --------------------------- Train/Validation ------------------------- 96 | print('Training...') 97 | train_loss_epochs = [] 98 | train_mape_epochs = [] 99 | train_rmse_epochs = [] 100 | val_loss_epochs = [] 101 | val_mape_epochs = [] 102 | val_rmse_epochs = [] 103 | for epoch in range(trainer_config['epoch']): 104 | loss, mape, rmse = engine.train_one_epoch(data_loaders[0]) 105 | train_loss_epochs.append(loss) 106 | train_mape_epochs.append(mape) 107 | train_rmse_epochs.append(rmse) 108 | 109 | print(f'epoch:{epoch}, train_loss(MAE):{train_loss_epochs[epoch]}, ' 110 | f'train_MAPE:{train_mape_epochs[epoch]}, train_RMSE:{train_rmse_epochs[epoch]}') 111 | 112 | loss, mape, rmse = engine.eval_one_epoch(data_loaders[1]) 113 | val_loss_epochs.append(loss) 114 | val_mape_epochs.append(mape) 115 | val_rmse_epochs.append(rmse) 116 | 117 | print(f'val_loss(MAE):{val_loss_epochs[epoch]}, ' 118 | f'val_MAPE:{val_mape_epochs[epoch]}, val_RMSE:{val_rmse_epochs[epoch]}') 119 | 120 | torch.save(model.state_dict(), trainer_config['model_path'] + '/' + 'gwnet_pretrained' + '.pth') 121 | print('Model saved!') 122 | 123 | # --------------------------- Test ------------------------- 124 | print('Testing...') 125 | test_loss, test_mape, test_rmse = engine.test(data_loaders[2]) 126 | print(f'test_loss(MAE):{test_loss}, test_mape:{test_mape}, test_rmse:{test_rmse}') 127 | 128 | 129 | if __name__ == '__main__': 130 | parser = argparse.ArgumentParser() 131 | parser.add_argument('--config_path', type=str, default='../config/config_graph_wavenet.yaml', 132 | help='Config path') 133 | args = parser.parse_args() 134 | configs = load_config(args.config_path) 135 | main(configs) 136 | -------------------------------------------------------------------------------- /main/main_stgcn.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | import argparse 3 | import numpy as np 4 | import pandas as pd 5 | import torch 6 | import torch.optim as optim 7 | from models.stgcn.models import STGCNChebGraphConv 8 | from utils import ProcessData, StandardScalar 9 | from torch.utils.data import DataLoader, TensorDataset 10 | from utils import masked_mae, calculate_scaled_laplacian_torch 11 | from trainer.trainer_stgcn import trainer 12 | 13 | 14 | def load_config(config_path): 15 | with open(config_path, 'r', encoding='utf-8') as f: 16 | config = yaml.safe_load(f) 17 | return config 18 | 19 | 20 | def main(config): 21 | # ------------------------ Data loading ------------------------ 22 | data_config = config['dataset'] 23 | optim_config = config['optimize'] 24 | trainer_config = config['trainer'] 25 | feature_matrix = np.load(data_config['feature_matrix_data_path']) 26 | samples = feature_matrix.shape[0] 27 | # shape -> (number of samples, sequence length, number of vertexes, features) 28 | train_data, train_labels, val_data, val_labels, test_data, test_labels = ProcessData(feature_matrix, 29 | data_config['input_seq'], 30 | data_config['output_seq'], 31 | data_config['split_rate']) 32 | # data scalar 33 | device = data_config['device'] 34 | scalar = StandardScalar(train_data, device, axis=(0, 1, 2)) 35 | data_x_sets = [train_data, val_data, test_data] 36 | data_y_sets = [train_labels, val_labels, test_labels] 37 | data_x_normed_sets = [] 38 | data_y_real_sets = [] 39 | for data in data_x_sets: 40 | data_normed = scalar.transform(data) 41 | data_normed = np.einsum('btnf->bftn', data_normed) 42 | data_normed = torch.Tensor(data_normed).to(device) 43 | data_x_normed_sets.append(data_normed) 44 | train_x, val_x, test_x = data_x_normed_sets 45 | for data in data_y_sets: 46 | data_real = torch.Tensor(data).to(device) 47 | data_y_real_sets.append(data_real) 48 | train_real_y, val_real_y, test_real_y = data_y_real_sets 49 | data_final_list = [train_x, train_real_y, val_x, val_real_y, test_x, test_real_y] 50 | 51 | # get dataloader 52 | data_loaders = [] 53 | for i in range(0, len(data_final_list), 2): 54 | dataset = TensorDataset(data_final_list[i], data_final_list[i + 1]) 55 | data_loader = DataLoader(dataset, batch_size=data_config['batch_size']) 56 | data_loaders.append(data_loader) 57 | 58 | # get adjacency matrix 59 | origin_adj = pd.read_csv(data_config['adj_data_path'], header=None) 60 | origin_adj = torch.Tensor(origin_adj.values).to(device) 61 | # ------------------------ Model setting ------------------------ 62 | 63 | Ko = trainer_config['n_his'] - (trainer_config['Kt'] - 1) * 2 * trainer_config['stblock_num'] 64 | blocks = [[trainer_config['channel']]] 65 | for l in range(trainer_config['stblock_num']): 66 | blocks.append([64, 16, 64]) 67 | if Ko == 0: 68 | blocks.append([128]) 69 | elif Ko > 0: 70 | blocks.append([128, 128]) 71 | blocks.append([trainer_config['channel']]) 72 | 73 | gso = calculate_scaled_laplacian_torch(origin_adj) 74 | 75 | model = STGCNChebGraphConv(blocks, data_config['num_of_vertices'], trainer_config['Kt'], trainer_config['Ks'], 76 | trainer_config['act_func'], trainer_config['graph_conv_type'], gso 77 | ,trainer_config['enable_bias'], trainer_config['dropout'], trainer_config['n_his']) 78 | model.to(device) 79 | 80 | if optim_config['name'] == 'Adam': 81 | optimizer = optim.Adam(model.parameters(), 82 | lr=optim_config['learning_rate'], 83 | weight_decay=optim_config['weight_decay']) 84 | 85 | elif optim_config['name'] == 'AdamW': 86 | optimizer = optim.AdamW(model.parameters(), 87 | lr=optim_config['learning_rate'], 88 | weight_decay=optim_config['weight_decay']) 89 | 90 | elif optim_config['name'] == 'RMSprop': 91 | optimizer = optim.RMSprop(model.parameters(), 92 | lr=optim_config['learning_rate'], 93 | weight_decay=optim_config['weight_decay']) 94 | else: 95 | raise NotImplementedError(f'ERROR: The optimizer {optim_config["name"]} is not implemented.') 96 | 97 | engine = trainer(model, optimizer, scalar, trainer_config['theta']) 98 | # --------------------------- Train/Validation ------------------------- 99 | print('Training...') 100 | train_loss_epochs = [] 101 | train_mape_epochs = [] 102 | train_rmse_epochs = [] 103 | val_loss_epochs = [] 104 | val_mape_epochs = [] 105 | val_rmse_epochs = [] 106 | for epoch in range(trainer_config['epoch']): 107 | loss, mape, rmse = engine.train_one_epoch(data_loaders[0]) 108 | train_loss_epochs.append(loss) 109 | train_mape_epochs.append(mape) 110 | train_rmse_epochs.append(rmse) 111 | 112 | print(f'epoch:{epoch}, train_loss(MAE):{train_loss_epochs[epoch]}, ' 113 | f'train_MAPE:{train_mape_epochs[epoch]}, train_RMSE:{train_rmse_epochs[epoch]}') 114 | 115 | loss, mape, rmse = engine.eval_one_epoch(data_loaders[1]) 116 | val_loss_epochs.append(loss) 117 | val_mape_epochs.append(mape) 118 | val_rmse_epochs.append(rmse) 119 | 120 | print(f'val_loss(MAE):{val_loss_epochs[epoch]}, ' 121 | f'val_MAPE:{val_mape_epochs[epoch]}, val_RMSE:{val_rmse_epochs[epoch]}') 122 | 123 | torch.save(model.state_dict(), trainer_config['model_path'] + '/' + 'stgcn_pretrained' + '.pth') 124 | print('Model saved!') 125 | 126 | # --------------------------- Test ------------------------- 127 | print('Testing...') 128 | test_loss, test_mape, test_rmse = engine.test(data_loaders[2]) 129 | print(f'test_loss(MAE):{test_loss}, test_mape:{test_mape}, test_rmse:{test_rmse}') 130 | 131 | if __name__ == '__main__': 132 | parser = argparse.ArgumentParser() 133 | parser.add_argument('--config_path', type=str, default='../config/config_stgcn.yaml', 134 | help='Config path') 135 | args = parser.parse_args() 136 | configs = load_config(args.config_path) 137 | main(configs) 138 | -------------------------------------------------------------------------------- /main/main_ma2gcn.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | import argparse 3 | import numpy as np 4 | import pandas as pd 5 | import torch 6 | import torch.optim as optim 7 | from torch.utils.data import DataLoader, TensorDataset 8 | from models.ma2gcn import MA2GCN 9 | from trainer.trainer_ma2gcn import trainer 10 | from utils import ProcessData, TransAdj, StandardScalar 11 | 12 | 13 | def load_config(config_path): 14 | with open(config_path, 'r', encoding='utf-8') as f: 15 | config = yaml.safe_load(f) 16 | return config 17 | 18 | 19 | def main(config): 20 | # ------------------------ Data loading ------------------------ 21 | data_config = config['dataset'] 22 | optim_config = config['optimize'] 23 | trainer_config = config['trainer'] 24 | feature_matrix = np.load(data_config['feature_matrix_data_path']) 25 | samples = feature_matrix.shape[0] 26 | # shape -> (number of samples, sequence length, number of vertexes, features) 27 | train_data, train_labels, val_data, val_labels, test_data, test_labels = ProcessData(feature_matrix, 28 | data_config['input_seq'], 29 | data_config['output_seq'], 30 | data_config['split_rate']) 31 | # data scalar 32 | device = data_config['device'] 33 | scalar = StandardScalar(train_data, device, axis=(0, 1, 2)) 34 | data_x_sets = [train_data, val_data, test_data] 35 | data_y_sets = [train_labels, val_labels, test_labels] 36 | data_x_normed_sets = [] 37 | data_y_real_sets = [] 38 | for data in data_x_sets: 39 | data_normed = scalar.transform(data) 40 | data_normed = torch.Tensor(data_normed).to(device) 41 | data_x_normed_sets.append(data_normed) 42 | train_x, val_x, test_x = data_x_normed_sets 43 | for data in data_y_sets: 44 | data_real = torch.Tensor(data).to(device) 45 | data_y_real_sets.append(data_real) 46 | train_real_y, val_real_y, test_real_y = data_y_real_sets 47 | data_final_list = [train_x, train_real_y, val_x, val_real_y, test_x, test_real_y] 48 | 49 | # get dataloader 50 | data_loaders = [] 51 | for i in range(0, len(data_final_list), 2): 52 | dataset = TensorDataset(data_final_list[i], data_final_list[i + 1]) 53 | data_loader = DataLoader(dataset, batch_size=data_config['batch_size']) 54 | data_loaders.append(data_loader) 55 | 56 | # get adjacency matrix 57 | origin_adj = pd.read_csv(data_config['origin_adj_data_path'], header=None) 58 | origin_adj = torch.Tensor(origin_adj.values).to(device) 59 | io_matrix = np.load(data_config['io_adj_data_path']) 60 | io_adj = torch.Tensor(TransAdj(io_matrix)).to(device) 61 | 62 | # process io adj 63 | io_bound = data_config['io_matrix_interval'] * 60 / data_config['basic_interval'] 64 | io_range = np.array(range(0, samples + 1, int(io_bound))) 65 | 66 | # get cur_io_adj for val and test 67 | val_start = train_data.shape[0] 68 | val_cur_io_adj = io_adj[np.digitize(val_start, io_range) - 2] 69 | test_start = val_start + val_data.shape[0] 70 | test_cur_io_adj = io_adj[np.digitize(test_start, io_range) - 2] 71 | 72 | # ------------------------ Model parameters getting ------------------------ 73 | blocks_num = trainer_config['block_nums'] 74 | kernel_size = trainer_config['kernel_size'] 75 | in_tcn_dim_list = trainer_config['in_tcn_dim_list'] 76 | adj_input_dim = origin_adj.shape[0] 77 | _, input_seq, nodes_num, features = train_x.shape 78 | # [Start]--get receptive field, dilation list, all_feature_dim, input_seq(padded) 79 | cur_dilation_list = [0] * blocks_num 80 | all_feature_dim = [0] * blocks_num 81 | cur_dilation_list[0] = 1 82 | receptive_field = 1 83 | additional_scope = kernel_size[1] - 1 84 | 85 | for i in range(blocks_num): 86 | receptive_field = receptive_field + additional_scope 87 | additional_scope *= 2 88 | 89 | new_input_seq = receptive_field 90 | all_feature_dim[0] = new_input_seq * features 91 | 92 | for j in range(1, blocks_num): 93 | cur_dilation_list[j] = 2 * cur_dilation_list[j - 1] 94 | all_feature_dim[j] = in_tcn_dim_list[j] * ( 95 | new_input_seq - cur_dilation_list[j - 1] * (kernel_size[1] - 1)) 96 | new_input_seq = new_input_seq - cur_dilation_list[j - 1] * (kernel_size[1] - 1) 97 | # [End]--get receptive field, dilation list, all_feature_dim, new_input_seq(padded) 98 | 99 | # ------------------------ Model setting ------------------------ 100 | model = MA2GCN(trainer_config['block_nums'], trainer_config['K'], adj_input_dim, trainer_config['adj_hidden_dim'], 101 | origin_adj, trainer_config['in_channel_list'], trainer_config['out_channel_list'], 102 | trainer_config['in_tcn_dim_list'], trainer_config['out_tcn_dim_list'], 103 | cur_dilation_list, trainer_config['kernel_size'], nodes_num, all_feature_dim, 104 | receptive_field, trainer_config['drop_rate'], device) 105 | model.to(device) 106 | 107 | if optim_config['name'] == 'Adam': 108 | optimizer = optim.Adam(model.parameters(), 109 | lr=optim_config['learning_rate'], 110 | weight_decay=optim_config['weight_decay']) 111 | 112 | elif optim_config['name'] == 'AdamW': 113 | optimizer = optim.AdamW(model.parameters(), 114 | lr=optim_config['learning_rate'], 115 | weight_decay=optim_config['weight_decay']) 116 | 117 | elif optim_config['name'] == 'RMSprop': 118 | optimizer = optim.RMSprop(model.parameters(), 119 | lr=optim_config['learning_rate'], 120 | weight_decay=optim_config['weight_decay']) 121 | else: 122 | raise NotImplementedError(f'ERROR: The optimizer {optim_config["name"]} is not implemented.') 123 | 124 | engine = trainer(model, optimizer, scalar, trainer_config['theta'], io_range, io_adj) 125 | # --------------------------- Train/Validation ------------------------- 126 | print('Training...') 127 | train_loss_epochs = [] 128 | train_mape_epochs = [] 129 | train_rmse_epochs = [] 130 | val_loss_epochs = [] 131 | val_mape_epochs = [] 132 | val_rmse_epochs = [] 133 | for epoch in range(trainer_config['epoch']): 134 | loss, mape, rmse = engine.train_one_epoch(data_loaders[0]) 135 | train_loss_epochs.append(loss) 136 | train_mape_epochs.append(mape) 137 | train_rmse_epochs.append(rmse) 138 | 139 | print(f'epoch:{epoch}, train_loss(MAE):{train_loss_epochs[epoch]}, ' 140 | f'train_MAPE:{train_mape_epochs[epoch]}, train_RMSE:{train_rmse_epochs[epoch]}') 141 | 142 | loss, mape, rmse = engine.eval_one_epoch(data_loaders[1], val_cur_io_adj) 143 | val_loss_epochs.append(loss) 144 | val_mape_epochs.append(mape) 145 | val_rmse_epochs.append(rmse) 146 | 147 | print(f'val_loss(MAE):{val_loss_epochs[epoch]}, ' 148 | f'val_MAPE:{val_mape_epochs[epoch]}, val_RMSE:{val_rmse_epochs[epoch]}') 149 | 150 | torch.save(model.state_dict(), trainer_config['model_path'] + '/' + 'ma2gcn_pretrained' + '.pth') 151 | print('Model saved!') 152 | 153 | # --------------------------- Test ------------------------- 154 | print('Testing...') 155 | test_loss, test_mape, test_rmse = engine.test(data_loaders[2], test_cur_io_adj) 156 | print(f'test_loss(MAE):{test_loss}, test_mape:{test_mape}, test_rmse:{test_rmse}') 157 | 158 | 159 | if __name__ == '__main__': 160 | parser = argparse.ArgumentParser() 161 | parser.add_argument('--config_path', type=str, default='../config/config_ma2gcn.yaml', 162 | help='Config path') 163 | args = parser.parse_args() 164 | configs = load_config(args.config_path) 165 | main(configs) 166 | -------------------------------------------------------------------------------- /models/graph_wavenet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class nconv(nn.Module): 7 | def __init__(self): 8 | super(nconv, self).__init__() 9 | 10 | def forward(self, x, A): 11 | x = torch.einsum('ncvl,vw->ncwl', (x, A)) 12 | return x.contiguous() 13 | 14 | 15 | class linear(nn.Module): 16 | def __init__(self, c_in, c_out): 17 | super(linear, self).__init__() 18 | self.mlp = torch.nn.Conv2d(c_in, c_out, kernel_size=(1, 1), padding=(0, 0), stride=(1, 1), bias=True) 19 | 20 | def forward(self, x): 21 | return self.mlp(x) 22 | 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, 51 | aptinit=None, in_dim=2, out_dim=12, residual_channels=32, dilation_channels=32, 52 | skip_channels=256, end_channels=512, kernel_size=2, blocks=4, layers=2): 53 | super(gwnet, self).__init__() 54 | self.dropout = dropout 55 | self.blocks = blocks 56 | self.layers = layers 57 | self.gcn_bool = gcn_bool 58 | self.addaptadj = addaptadj 59 | 60 | self.filter_convs = nn.ModuleList() 61 | self.gate_convs = nn.ModuleList() 62 | self.residual_convs = nn.ModuleList() 63 | self.skip_convs = nn.ModuleList() 64 | self.bn = nn.ModuleList() 65 | self.gconv = nn.ModuleList() 66 | 67 | self.start_conv = nn.Conv2d(in_channels=in_dim, 68 | out_channels=residual_channels, 69 | kernel_size=(1, 1)) 70 | self.supports = supports 71 | 72 | receptive_field = 1 73 | 74 | self.supports_len = 0 75 | if supports is not None: 76 | self.supports_len += len(supports) 77 | 78 | if gcn_bool and addaptadj: 79 | if aptinit is None: 80 | if supports is None: 81 | self.supports = [] 82 | self.nodevec1 = nn.Parameter(torch.randn(num_nodes, 10).to(device), requires_grad=True).to(device) 83 | self.nodevec2 = nn.Parameter(torch.randn(10, num_nodes).to(device), requires_grad=True).to(device) 84 | self.supports_len += 1 85 | else: 86 | if supports is None: 87 | self.supports = [] 88 | m, p, n = torch.svd(aptinit) 89 | initemb1 = torch.mm(m[:, :10], torch.diag(p[:10] ** 0.5)) 90 | initemb2 = torch.mm(torch.diag(p[:10] ** 0.5), n[:, :10].t()) 91 | self.nodevec1 = nn.Parameter(initemb1, requires_grad=True).to(device) 92 | self.nodevec2 = nn.Parameter(initemb2, requires_grad=True).to(device) 93 | self.supports_len += 1 94 | 95 | for b in range(blocks): 96 | additional_scope = kernel_size - 1 97 | new_dilation = 1 98 | for i in range(layers): 99 | # dilated convolutions 100 | self.filter_convs.append(nn.Conv2d(in_channels=residual_channels, 101 | out_channels=dilation_channels, 102 | kernel_size=(1, kernel_size), dilation=new_dilation)) 103 | 104 | self.gate_convs.append(nn.Conv2d(in_channels=residual_channels, 105 | out_channels=dilation_channels, 106 | kernel_size=(1, kernel_size), dilation=new_dilation)) 107 | 108 | # 1x1 convolution for residual connection 109 | self.residual_convs.append(nn.Conv2d(in_channels=dilation_channels, 110 | out_channels=residual_channels, 111 | kernel_size=(1, 1))) 112 | 113 | # 1x1 convolution for skip connection 114 | self.skip_convs.append(nn.Conv2d(in_channels=dilation_channels, 115 | out_channels=skip_channels, 116 | kernel_size=(1, 1))) 117 | 118 | self.bn.append(nn.BatchNorm2d(residual_channels)) 119 | new_dilation *= 2 120 | receptive_field += additional_scope 121 | additional_scope *= 2 122 | if self.gcn_bool: 123 | self.gconv.append(gcn(dilation_channels, residual_channels, dropout, support_len=self.supports_len)) 124 | 125 | self.end_conv_1 = nn.Conv2d(in_channels=skip_channels, 126 | out_channels=end_channels, 127 | kernel_size=(1, 1), 128 | bias=True) 129 | 130 | self.end_conv_2 = nn.Conv2d(in_channels=end_channels, 131 | out_channels=out_dim, 132 | kernel_size=(1, 1), 133 | bias=True) 134 | 135 | self.receptive_field = receptive_field + 1 136 | 137 | def forward(self, input): 138 | in_len = input.size(3) 139 | if in_len < self.receptive_field: 140 | x = nn.functional.pad(input, (self.receptive_field - in_len, 0, 0, 0)) 141 | else: 142 | x = input 143 | x = self.start_conv(x) 144 | skip = 0 145 | 146 | # calculate the current adaptive adj matrix once per iteration 147 | new_supports = None 148 | if self.gcn_bool and self.addaptadj and self.supports is not None: 149 | adp = F.softmax(F.relu(torch.mm(self.nodevec1, self.nodevec2)), dim=1) 150 | new_supports = self.supports + [adp] 151 | 152 | # WaveNet layers 153 | for i in range(self.blocks * self.layers): 154 | 155 | # |----------------------------------------| *residual* 156 | # | | 157 | # | |-- conv -- tanh --| | 158 | # -> dilate -|----| * ----|-- 1x1 -- + --> *input* 159 | # |-- conv -- sigm --| | 160 | # 1x1 161 | # | 162 | # ---------------------------------------> + -------------> *skip* 163 | 164 | residual = x 165 | # dilated convolution 166 | filter = self.filter_convs[i](residual) 167 | filter = torch.tanh(filter) 168 | gate = self.gate_convs[i](residual) 169 | gate = torch.sigmoid(gate) 170 | x = filter * gate 171 | 172 | # parametrized skip connection 173 | 174 | s = x 175 | s = self.skip_convs[i](s) 176 | try: 177 | skip = skip[:, :, :, -s.size(3):] 178 | except: 179 | skip = 0 180 | skip = s + skip 181 | 182 | if self.gcn_bool and self.supports is not None: 183 | if self.addaptadj: 184 | x = self.gconv[i](x, new_supports) 185 | else: 186 | x = self.gconv[i](x, self.supports) 187 | else: 188 | x = self.residual_convs[i](x) 189 | 190 | x = x + residual[:, :, :, -x.size(3):] 191 | 192 | x = self.bn[i](x) 193 | 194 | x = F.relu(skip) 195 | x = F.relu(self.end_conv_1(x)) 196 | x = self.end_conv_2(x) 197 | return x 198 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from scipy.sparse.linalg import eigs 4 | 5 | 6 | # ------------------------ Process data ------------------------ 7 | 8 | def ProcessData(data, in_seq_length, out_seq_length, rates): 9 | """ 10 | :param data: raw data 11 | :param in_seq_length: the number of timestamps for prediction 12 | :param out_seq_length: the number of timestamps as labels 13 | :param rates: train/val/test 14 | :return: shape -> (number of samples, sequence length, number of vertexes, features) 15 | """ 16 | # Set parameters 17 | features = data 18 | input_seq_length = in_seq_length 19 | output_seq_length = out_seq_length 20 | train_rate = rates[0] 21 | val_rate = rates[1] 22 | 23 | assert 10 - train_rate - val_rate == rates[2] 24 | 25 | # Get time sequence length 26 | data_length = features.shape[0] 27 | 28 | # Get index of train/val/test set 29 | train_length = int(train_rate / 10 * (data_length - input_seq_length - output_seq_length)) 30 | val_length = int(val_rate / 10 * (data_length - input_seq_length - output_seq_length)) 31 | 32 | train_start_idx = 0 33 | train_end_idx = train_length + 1 34 | 35 | val_start_idx = train_end_idx 36 | val_end_idx = val_start_idx + val_length + 1 37 | 38 | test_start_idx = val_end_idx 39 | test_end_idx = data_length - input_seq_length - output_seq_length + 1 40 | 41 | # Split dataset 42 | def SplitData(start_idx, end_idx): 43 | inputs = [] 44 | labels = [] 45 | for i in range(start_idx, end_idx): 46 | inputs.append(features[i:i + input_seq_length]) 47 | labels.append(features[i + input_seq_length:i + input_seq_length + output_seq_length]) 48 | return np.array(inputs), np.array(labels) 49 | 50 | train_data, train_labels = SplitData(train_start_idx, train_end_idx) 51 | val_data, val_labels = SplitData(val_start_idx, val_end_idx) 52 | test_data, test_labels = SplitData(test_start_idx, test_end_idx) 53 | 54 | return train_data, train_labels, val_data, val_labels, test_data, test_labels 55 | 56 | 57 | def TransAdj(io_matrix): 58 | # Set parameter 59 | io_matrix = io_matrix 60 | 61 | # Calculate weights 62 | num_trains = io_matrix[:, :, :, 0] 63 | travel_times = io_matrix[:, :, :, 1] 64 | weights = np.divide(num_trains, travel_times, out=np.zeros_like(num_trains), where=travel_times != 0) 65 | 66 | # Get adj_matrix 67 | adj_matrix = np.zeros((io_matrix.shape[0], io_matrix.shape[1], io_matrix.shape[1])) 68 | mask = np.sum(io_matrix, axis=-1) != 0 69 | adj_matrix[mask] = weights[mask] 70 | 71 | return adj_matrix 72 | 73 | 74 | class StandardScalar: 75 | def __init__(self, data, device, axis=None): 76 | # Get mean, std, axis 77 | self.axis = axis 78 | self.mean = data.mean(axis=self.axis) 79 | self.std = data.std(axis=self.axis) 80 | self.device = device 81 | 82 | def transform(self, data): 83 | return (data - self.mean) / self.std 84 | 85 | def inverse_transform(self, data): 86 | self.std = torch.Tensor(self.std).to(self.device) 87 | self.mean = torch.Tensor(self.mean).to(self.device) 88 | return (data * self.std) + self.mean 89 | 90 | 91 | class StandardScalar_np: 92 | def __init__(self, data, axis=None): 93 | # Get mean, std, axis 94 | self.axis = axis 95 | self.mean = data.mean(axis=self.axis) 96 | self.std = data.std(axis=self.axis) 97 | 98 | def transform(self, data): 99 | return (data - self.mean) / self.std 100 | 101 | def inverse_transform(self, data): 102 | return (data * self.std) + self.mean 103 | 104 | 105 | # ------------------------ Model tools ------------------------ 106 | 107 | # Get L_tilde or L_sym-I or -D^-(1/2)*A*D^-(1/2) 108 | 109 | def calculate_scaled_laplacian_torch(adj): 110 | """ 111 | torch version 112 | """ 113 | D = torch.diag(torch.sum(adj, dim=1)) 114 | L = D - adj 115 | 116 | eigenvalues, _ = torch.linalg.eig(L) 117 | lambda_max = eigenvalues.real.max() 118 | 119 | return (2 * L) / lambda_max - torch.eye(adj.shape[0]).cuda() 120 | 121 | 122 | def calculate_scaled_laplacian_np(adj): 123 | """ 124 | numpy version 125 | """ 126 | # if adj is tensor(cuda) 127 | adj = adj.cpu() 128 | adj = adj.detach().numpy() 129 | adj = np.array(adj) 130 | 131 | D = np.diag(np.sum(adj, axis=1)) 132 | L = D - adj 133 | 134 | lambda_max = eigs(L, k=1, which='LR')[0].real 135 | 136 | return (2 * L) / lambda_max - np.identity(adj.shape[0]) 137 | 138 | 139 | # Get chebyshev polynomials 140 | def get_Tk(L_tilde, K): 141 | L_tilde = L_tilde.cpu() 142 | L_tilde = L_tilde.detach().numpy() 143 | L_tilde = np.array(L_tilde) 144 | 145 | T_ks = [] 146 | N = L_tilde.shape[0] 147 | T_ks.append(np.identity(N)) 148 | T_ks.append(L_tilde) 149 | 150 | for i in range(2, K): 151 | T_ks.append(2 * L_tilde * T_ks[i - 1] - T_ks[i - 2]) 152 | 153 | return T_ks 154 | 155 | 156 | # Get current io matrix 157 | def IoMatrixSelect(io_adj, past_hours): 158 | number_segments = io_adj.shape[0] 159 | base_hour_interval = int(24 / number_segments) 160 | past_hours = past_hours 161 | 162 | return io_adj[int(past_hours / base_hour_interval)] 163 | 164 | 165 | # ------------------------ Metrics ------------------------ 166 | # MASK MAE 167 | def masked_mae(preds, labels, null_val=np.nan): 168 | if np.isnan(null_val): 169 | mask = ~torch.isnan(labels) 170 | else: 171 | mask = (labels != null_val) 172 | mask = mask.float() 173 | mask /= torch.mean((mask)) 174 | mask = torch.where(torch.isnan(mask), torch.zeros_like(mask), mask) 175 | loss = torch.abs(preds - labels) 176 | loss = loss * mask 177 | loss = torch.where(torch.isnan(loss), torch.zeros_like(loss), loss) 178 | return torch.mean(loss) 179 | 180 | 181 | def masked_mae_np(preds, labels, null_val=np.nan): 182 | if np.isnan(null_val): 183 | mask = ~np.isnan(labels) 184 | else: 185 | mask = (labels != null_val) 186 | mask = mask.astype(float) 187 | mask /= np.mean(mask) 188 | mask = np.where(np.isnan(mask), np.zeros_like(mask), mask) 189 | loss = np.abs(preds - labels) 190 | loss = loss * mask 191 | loss = np.where(np.isnan(loss), np.zeros_like(loss), loss) 192 | return np.mean(loss) 193 | 194 | 195 | # MASK MAPE 196 | def masked_mape(preds, labels, null_val=np.nan): 197 | if np.isnan(null_val): 198 | mask = ~torch.isnan(labels) 199 | else: 200 | mask = (labels != null_val) 201 | mask = mask.float() 202 | mask /= torch.mean((mask)) 203 | mask = torch.where(torch.isnan(mask), torch.zeros_like(mask), mask) 204 | loss = torch.abs(preds - labels) / labels 205 | loss = loss * mask 206 | loss = torch.where(torch.isnan(loss), torch.zeros_like(loss), loss) 207 | return torch.mean(loss) 208 | 209 | 210 | def masked_mape_np(preds, labels, null_val=np.nan): 211 | if np.isnan(null_val): 212 | mask = ~np.isnan(labels) 213 | else: 214 | mask = (labels != null_val) 215 | mask = mask.astype(float) 216 | mask /= np.mean(mask) 217 | mask = np.where(np.isnan(mask), np.zeros_like(mask), mask) 218 | loss = np.abs(preds - labels) / labels 219 | loss = loss * mask 220 | loss = np.where(np.isnan(loss), np.zeros_like(loss), loss) 221 | return np.mean(loss) 222 | 223 | 224 | # MASK MSE/RMSE 225 | def masked_mse(preds, labels, null_val=np.nan): 226 | if np.isnan(null_val): 227 | mask = ~torch.isnan(labels) 228 | else: 229 | mask = (labels != null_val) 230 | mask = mask.float() 231 | mask /= torch.mean((mask)) 232 | mask = torch.where(torch.isnan(mask), torch.zeros_like(mask), mask) 233 | loss = (preds - labels) ** 2 234 | loss = loss * mask 235 | loss = torch.where(torch.isnan(loss), torch.zeros_like(loss), loss) 236 | return torch.mean(loss) 237 | 238 | 239 | def masked_rmse(preds, labels, null_val=np.nan): 240 | return torch.sqrt(masked_mse(preds=preds, labels=labels, null_val=null_val)) 241 | 242 | 243 | def masked_mse_np(preds, labels, null_val=np.nan): 244 | if np.isnan(null_val): 245 | mask = ~np.isnan(labels) 246 | else: 247 | mask = (labels != null_val) 248 | mask = mask.astype(float) 249 | mask /= np.mean(mask) 250 | mask = np.where(np.isnan(mask), np.zeros_like(mask), mask) 251 | loss = (preds - labels) ** 2 252 | loss = loss * mask 253 | loss = np.where(np.isnan(loss), np.zeros_like(loss), loss) 254 | return np.mean(loss) 255 | 256 | 257 | def masked_rmse_np(preds, labels, null_val=np.nan): 258 | return np.sqrt(masked_mse_np(preds=preds, labels=labels, null_val=null_val)) 259 | -------------------------------------------------------------------------------- /models/ma2gcn.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from utils import calculate_scaled_laplacian_torch, get_Tk 6 | 7 | 8 | # ------------------------ Graph convolution ------------------------ 9 | 10 | class GraphConv(nn.Module): 11 | """ 12 | - GraphConv: 13 | - Params: in_channels*, out_channels*, device, K 14 | - Input: x*(b, t_in, n, c_in), Tks* 15 | - Output: x(b, t_in, n, c_out) 16 | """ 17 | 18 | def __init__(self, K, in_channels, out_channels, device): 19 | super(GraphConv, self).__init__() 20 | self.K = K 21 | self.device = device 22 | self.in_channels = in_channels 23 | self.out_channels = out_channels 24 | self.Theta = nn.ParameterList( 25 | [nn.Parameter(torch.FloatTensor(in_channels, out_channels)).to(device) for _ in range(K)] 26 | ) 27 | 28 | def forward(self, x, Tks): 29 | batch_size, seq_length, num_of_vertices, in_channels = x.shape 30 | output = torch.zeros(batch_size, seq_length, num_of_vertices, self.out_channels).to(self.device) 31 | 32 | for k in range(self.K): 33 | T_k = Tks[k] 34 | theta_k = self.Theta[k] 35 | 36 | temp = torch.einsum('vv,btvc->btvc', (T_k, x)) 37 | output = output + torch.einsum('btvc,co->btvo', (temp, theta_k)) 38 | 39 | return torch.Tensor(output) 40 | 41 | 42 | # ------------------------ Adjacency attention mechanism------------------------ 43 | 44 | class AdjAttention(nn.Module): 45 | """ 46 | - AdjAttention: 47 | - Params: adj_input_dim, adj_hidden_dim(For weights) 48 | - Input: adj_list*:[origin_adj(n, n), origin_adj_2(n, n), dynamic_adj(n, n), cur_io_adj(n, n)] 49 | - Output: adj_aggregated(n, n) 50 | """ 51 | 52 | def __init__(self, adj_input_dim, adj_hidden_dim): 53 | super(AdjAttention, self).__init__() 54 | 55 | self.W = nn.Linear(adj_input_dim, adj_hidden_dim) 56 | self.V = nn.Linear(adj_hidden_dim, 1) 57 | 58 | def forward(self, adj_list): 59 | # adj_list: a list of adjacency matrices with shape (num_nodes, num_nodes) 60 | num_nodes, _ = adj_list[0].shape 61 | 62 | # Compute weights for each adjacency matrix 63 | weights = [] 64 | for adj in adj_list: 65 | x = F.relu(self.W(adj)) 66 | x = self.V(x) 67 | x = x.view(num_nodes, 1) 68 | alpha = x.mean() 69 | weights.append(alpha) 70 | 71 | weights = F.softmax(torch.Tensor(weights), dim=0) 72 | # Compute weighted sum of adjacency matrices 73 | adj_aggregated = torch.zeros_like(adj_list[0]) 74 | for i in range(len(weights)): 75 | adj_aggregated += adj_list[i] * weights[i] 76 | 77 | return adj_aggregated 78 | 79 | 80 | # ------------------------ Gated TCN module ------------------------ 81 | 82 | class GatedTCN(nn.Module): 83 | """ 84 | - GatedTCN: 85 | - Params: in_tcn_dim, out_tcn_dim, cur_dilation, kernel_size 86 | - Input: x*(b, c, t_in, n) 87 | - Output: x(b, c, t_out, n) 88 | """ 89 | 90 | def __init__(self, in_tcn_dim, out_tcn_dim, cur_dilation, kernel_size): 91 | super(GatedTCN, self).__init__() 92 | self.filter_convs = nn.Conv2d(in_channels=in_tcn_dim, 93 | out_channels=out_tcn_dim, 94 | kernel_size=kernel_size, dilation=cur_dilation) 95 | 96 | self.gated_convs = nn.Conv2d(in_channels=in_tcn_dim, 97 | out_channels=out_tcn_dim, 98 | kernel_size=kernel_size, dilation=cur_dilation) 99 | 100 | def forward(self, x): 101 | x = torch.einsum('btvf->bfvt', x) 102 | filter = self.filter_convs(x) 103 | filter = torch.tanh(filter) 104 | gate = self.gated_convs(x) 105 | gate = torch.sigmoid(gate) 106 | x = filter * gate 107 | x = torch.einsum('bfvt->btvf', x) 108 | return x 109 | 110 | 111 | # ------------------------ Dynamic adjacency generator ------------------------ 112 | 113 | class DynamicAdjGen(nn.Module): 114 | """ 115 | - DynamicAdjGen: 116 | - Params: nodes_num, all_feature_dim*, device 117 | - Input: x*(b, t, n, c), origin_adj(n, n), cur_io_adj*(n, n) 118 | - Output: adj_list:[origin_adj(n, n), origin_adj_2(n, n), dynamic_adj(n, n), cur_io_adj(n, n)] 119 | """ 120 | 121 | def __init__(self, nodes_num, all_feature_dim, device): 122 | super(DynamicAdjGen, self).__init__() 123 | self.all_feature_dim = all_feature_dim 124 | self.nodes_num = nodes_num 125 | self.node_vec1 = nn.Parameter(torch.FloatTensor(self.all_feature_dim, self.nodes_num // 10)).to(device) 126 | self.node_vec2 = nn.Parameter(torch.FloatTensor(self.nodes_num, self.nodes_num)).to(device) 127 | nn.init.uniform_(self.node_vec1) 128 | nn.init.uniform_(self.node_vec2) 129 | 130 | def forward(self, x, origin_adj, cur_io_adj): 131 | b, t, n, c = x.shape 132 | x = x.reshape(b, n, t * c) 133 | x_emb = x @ self.node_vec1 134 | x_emb_T = torch.einsum('bnf->bfn', x_emb) 135 | x_selfdot = torch.mean(torch.einsum('bnf,bfm->bnm', (x_emb, x_emb_T)), dim=0) 136 | x_selfdot = x_selfdot.reshape(n, n) 137 | dynamic_adj = x_selfdot @ self.node_vec2 138 | dynamic_adj = F.softmax(dynamic_adj, dim=1) 139 | origin_adj = F.softmax(origin_adj, dim=1) 140 | origin_adj_2 = origin_adj @ origin_adj.T 141 | cur_io_adj = F.softmax(cur_io_adj, dim=1) 142 | 143 | return [origin_adj, origin_adj_2, dynamic_adj, cur_io_adj] 144 | # return [origin_adj, origin_adj_2, cur_io_adj] # without dynamic adjacency matrix 145 | 146 | 147 | # ------------------------ Multi adjacency attention convolution Block ------------------------ 148 | 149 | 150 | class MA2ConvBlock(nn.Module): 151 | """ 152 | - MA2ConvBlock: 153 | - Params: adj_input_dim/adj_hidden_dim/origin_adj(n, n)(for AdjAttention layer); 154 | device(for GraphConv layer and DynamicAdjGen layer); 155 | K/in_channel/out_channel(for GraphConv layer); 156 | in_tcn_dim/out_tcn_dim/cur_dilation/kernel_size(for GatedTCN layer); 157 | nodes_num/all_feature_dim(for DynamicAdjGen layer); 158 | - Input: x*(b, t_in, n, c_in), cur_io_adj*(for DynamicAdjGen layer); 159 | - Output: x(b, t_out, n, c_out) 160 | """ 161 | 162 | def __init__(self, adj_input_dim, adj_hidden_dim, K, in_channel, out_channel, 163 | in_tcn_dim, out_tcn_dim, origin_adj, nodes_num, all_feature_dim, 164 | cur_dilation, kernel_size, droprate, device): 165 | super(MA2ConvBlock, self).__init__() 166 | 167 | self.adj_gen = DynamicAdjGen(nodes_num, all_feature_dim, device) 168 | self.adj_att = AdjAttention(adj_input_dim, adj_hidden_dim) 169 | self.gated_tcn = GatedTCN(in_tcn_dim, out_tcn_dim, cur_dilation, kernel_size) 170 | self.cheb_gcn = GraphConv(K, in_channel, out_channel, device) 171 | self.residual = nn.Conv2d(in_channels=in_channel, out_channels=out_channel, 172 | kernel_size=(1, 1)) 173 | self.origin_adj = origin_adj 174 | self.K = K 175 | self.device = device 176 | self.dropout = nn.Dropout(droprate) 177 | 178 | def forward(self, x_with_io): 179 | x = x_with_io[0] 180 | residual = x 181 | cur_io_adj = x_with_io[1] 182 | x_t = self.gated_tcn(x) 183 | adj_list = self.adj_gen(x, self.origin_adj, cur_io_adj) 184 | adj_agg = self.adj_att(adj_list) 185 | 186 | # # [start] without attention module 187 | # adj_mean = torch.zeros_like(adj_list[0]) 188 | # ratio = 1/len(adj_list) 189 | # for adj in adj_list: 190 | # adj_mean = adj_mean + ratio * adj 191 | # L_tilde = calculate_scaled_laplacian_torch(adj_mean) 192 | # # [end] without attention module 193 | 194 | L_tilde = calculate_scaled_laplacian_torch(adj_agg) 195 | 196 | # # [start]without attention module and dynamic adjacency matrix 197 | # L_tilde = calculate_scaled_laplacian_torch(self.origin_adj) 198 | # # [end]without attention module and dynamic adjacency matrix 199 | 200 | Tks = torch.Tensor(np.array(get_Tk(L_tilde, self.K))).to(self.device) 201 | x_t_g = self.cheb_gcn(x_t, Tks) 202 | x_t_g = self.dropout(x_t_g) 203 | 204 | residual = residual.permute(0, 3, 1, 2) 205 | x_residual = self.residual(residual) 206 | x_residual = x_residual.permute(0, 2, 3, 1) 207 | x_t_g = x_t_g + x_residual[:, -x_t_g.size(1):, :, :] 208 | 209 | return (x_t_g, cur_io_adj) 210 | 211 | 212 | # ------------------------ Main architecture[Multi adjacency attention GCN] ------------------------ 213 | 214 | class MA2GCN(nn.Module): 215 | """ 216 | - MA2GCN: 217 | - Params: blocks, adj_input_dim/adj_hidden_dim/origin_adj(n, n)(for AdjAttention layer); 218 | device(for GraphConv layer and DynamicAdjGen layer); 219 | K/in_channel_list/out_channel_list(for GraphConv layer); 220 | in_tcn_dim_list/out_tcn_dim_list/cur_dilation/kernel_size(for GatedTCN layer); 221 | nodes_num/all_feature_dim(for DynamicAdjGen layer); 222 | - Input: x*(b, t_in, n, c), cur_io_adj*(for DynamicAdjGen layer); 223 | - Output: y(b, t_out, n, c) 224 | """ 225 | 226 | def __init__(self, blocks_num, K, adj_input_dim, adj_hidden_dim, origin_adj, 227 | in_channel_list, out_channel_list, in_tcn_dim_list, 228 | out_tcn_dim_list, cur_dilation_list, kernel_size, 229 | nodes_num, all_feature_dim, receptive_field, droprate, device): 230 | super(MA2GCN, self).__init__() 231 | modules = [] 232 | for i in range(blocks_num): 233 | modules.append(MA2ConvBlock(adj_input_dim, adj_hidden_dim, K, in_channel_list[i], out_channel_list[i], 234 | in_tcn_dim_list[i], out_tcn_dim_list[i], origin_adj, nodes_num, 235 | all_feature_dim[i], cur_dilation_list[i], kernel_size, droprate, device)) 236 | self.blocks = nn.Sequential(*modules) 237 | self.receptive_field = receptive_field 238 | 239 | def forward(self, x_with_io): 240 | x = x_with_io[0] 241 | input_seq = x.shape[1] 242 | # b,t,n,c -> b,c,n,t 243 | x = x.permute(0, 3, 2, 1) 244 | if input_seq < self.receptive_field: 245 | x_pad = F.pad(x, (self.receptive_field - input_seq, 0, 0, 0)) 246 | else: 247 | x_pad = input 248 | x_pad = x_pad.permute(0, 3, 2, 1) 249 | x_with_io = (x_pad, x_with_io[1]) 250 | return self.blocks(x_with_io) 251 | -------------------------------------------------------------------------------- /models/ASTGCN_r.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from utils import calculate_scaled_laplacian_torch, get_Tk 5 | 6 | 7 | class Spatial_Attention_layer(nn.Module): 8 | """ 9 | compute spatial attention scores 10 | """ 11 | 12 | def __init__(self, DEVICE, in_channels, num_of_vertices, num_of_timesteps): 13 | super(Spatial_Attention_layer, self).__init__() 14 | self.W1 = nn.Parameter(torch.FloatTensor(num_of_timesteps).to(DEVICE)) 15 | self.W2 = nn.Parameter(torch.FloatTensor(in_channels, num_of_timesteps).to(DEVICE)) 16 | self.W3 = nn.Parameter(torch.FloatTensor(in_channels).to(DEVICE)) 17 | self.bs = nn.Parameter(torch.FloatTensor(1, num_of_vertices, num_of_vertices).to(DEVICE)) 18 | self.Vs = nn.Parameter(torch.FloatTensor(num_of_vertices, num_of_vertices).to(DEVICE)) 19 | 20 | def forward(self, x): 21 | """ 22 | :param x: (batch_size, N, F_in, T) 23 | :return: (B,N,N) 24 | """ 25 | 26 | lhs = torch.matmul(torch.matmul(x, self.W1), self.W2) # (b,N,F,T)(T)->(b,N,F)(F,T)->(b,N,T) 27 | 28 | rhs = torch.matmul(self.W3, x).transpose(-1, -2) # (F)(b,N,F,T)->(b,N,T)->(b,T,N) 29 | 30 | product = torch.matmul(lhs, rhs) # (b,N,T)(b,T,N) -> (B, N, N) 31 | 32 | S = torch.matmul(self.Vs, torch.sigmoid(product + self.bs)) # (N,N)(B, N, N)->(B,N,N) 33 | 34 | S_normalized = F.softmax(S, dim=1) 35 | 36 | return S_normalized 37 | 38 | 39 | class cheb_conv_withSAt(nn.Module): 40 | """ 41 | K-order chebyshev graph convolution 42 | """ 43 | 44 | def __init__(self, K, cheb_polynomials, in_channels, out_channels): 45 | """ 46 | :param K: int 47 | :param in_channles: int, num of channels in the input sequence 48 | :param out_channels: int, num of channels in the output sequence 49 | """ 50 | super(cheb_conv_withSAt, self).__init__() 51 | self.K = K 52 | self.cheb_polynomials = cheb_polynomials 53 | self.in_channels = in_channels 54 | self.out_channels = out_channels 55 | self.DEVICE = cheb_polynomials[0].device 56 | self.Theta = nn.ParameterList( 57 | [nn.Parameter(torch.FloatTensor(in_channels, out_channels).to(self.DEVICE)) for _ in range(K)]) 58 | 59 | def forward(self, x, spatial_attention): 60 | """ 61 | Chebyshev graph convolution operation 62 | :param x: (batch_size, N, F_in, T) 63 | :return: (batch_size, N, F_out, T) 64 | """ 65 | 66 | batch_size, num_of_vertices, in_channels, num_of_timesteps = x.shape 67 | 68 | outputs = [] 69 | 70 | for time_step in range(num_of_timesteps): 71 | 72 | graph_signal = x[:, :, :, time_step] # (b, N, F_in) 73 | 74 | output = torch.zeros(batch_size, num_of_vertices, self.out_channels).to(self.DEVICE) # (b, N, F_out) 75 | 76 | for k in range(self.K): 77 | T_k = self.cheb_polynomials[k] # (N,N) 78 | 79 | T_k_with_at = T_k.mul(spatial_attention) # (N,N)*(N,N) = (N,N) 多行和为1, 按着列进行归一化 80 | 81 | theta_k = self.Theta[k] # (in_channel, out_channel) 82 | 83 | rhs = T_k_with_at.permute(0, 2, 1).matmul( 84 | graph_signal) # (N, N)(b, N, F_in) = (b, N, F_in) 因为是左乘，所以多行和为1变为多列和为1，即一行之和为1，进行左乘 85 | 86 | # temp = torch.einsum('nvv,nvc->nvc', (T_k_with_at, graph_signal)) 87 | # b = (rhs == temp) 88 | 89 | output = output + rhs.matmul(theta_k) # (b, N, F_in)(F_in, F_out) = (b, N, F_out) 90 | 91 | outputs.append(output.unsqueeze(-1)) # (b, N, F_out, 1) 92 | 93 | return F.relu(torch.cat(outputs, dim=-1)) # (b, N, F_out, T) 94 | 95 | 96 | class Temporal_Attention_layer(nn.Module): 97 | def __init__(self, DEVICE, in_channels, num_of_vertices, num_of_timesteps): 98 | super(Temporal_Attention_layer, self).__init__() 99 | self.U1 = nn.Parameter(torch.FloatTensor(num_of_vertices).to(DEVICE)) 100 | self.U2 = nn.Parameter(torch.FloatTensor(in_channels, num_of_vertices).to(DEVICE)) 101 | self.U3 = nn.Parameter(torch.FloatTensor(in_channels).to(DEVICE)) 102 | self.be = nn.Parameter(torch.FloatTensor(1, num_of_timesteps, num_of_timesteps).to(DEVICE)) 103 | self.Ve = nn.Parameter(torch.FloatTensor(num_of_timesteps, num_of_timesteps).to(DEVICE)) 104 | 105 | def forward(self, x): 106 | """ 107 | :param x: (batch_size, N, F_in, T) 108 | :return: (B, T, T) 109 | """ 110 | _, num_of_vertices, num_of_features, num_of_timesteps = x.shape 111 | 112 | lhs = torch.matmul(torch.matmul(x.permute(0, 3, 2, 1), self.U1), self.U2) 113 | # x:(B, N, F_in, T) -> (B, T, F_in, N) 114 | # (B, T, F_in, N)(N) -> (B,T,F_in) 115 | # (B,T,F_in)(F_in,N)->(B,T,N) 116 | 117 | rhs = torch.matmul(self.U3, x) # (F)(B,N,F,T)->(B, N, T) 118 | 119 | product = torch.matmul(lhs, rhs) # (B,T,N)(B,N,T)->(B,T,T) 120 | 121 | E = torch.matmul(self.Ve, torch.sigmoid(product + self.be)) # (B, T, T) 122 | 123 | E_normalized = F.softmax(E, dim=1) 124 | 125 | return E_normalized 126 | 127 | 128 | class cheb_conv(nn.Module): 129 | """ 130 | K-order chebyshev graph convolution 131 | """ 132 | 133 | def __init__(self, K, cheb_polynomials, in_channels, out_channels): 134 | ''' 135 | :param K: int 136 | :param in_channles: int, num of channels in the input sequence 137 | :param out_channels: int, num of channels in the output sequence 138 | ''' 139 | super(cheb_conv, self).__init__() 140 | self.K = K 141 | self.cheb_polynomials = cheb_polynomials 142 | self.in_channels = in_channels 143 | self.out_channels = out_channels 144 | self.DEVICE = cheb_polynomials[0].device 145 | self.Theta = nn.ParameterList( 146 | [nn.Parameter(torch.FloatTensor(in_channels, out_channels).to(self.DEVICE)) for _ in range(K)]) 147 | 148 | def forward(self, x): 149 | """ 150 | Chebyshev graph convolution operation 151 | :param x: (batch_size, N, F_in, T) 152 | :return: (batch_size, N, F_out, T) 153 | """ 154 | 155 | batch_size, num_of_vertices, in_channels, num_of_timesteps = x.shape 156 | 157 | outputs = [] 158 | 159 | for time_step in range(num_of_timesteps): 160 | 161 | graph_signal = x[:, :, :, time_step] # (b, N, F_in) 162 | 163 | output = torch.zeros(batch_size, num_of_vertices, self.out_channels).to(self.DEVICE) # (b, N, F_out) 164 | 165 | for k in range(self.K): 166 | T_k = self.cheb_polynomials[k] # (N,N) 167 | 168 | theta_k = self.Theta[k] # (in_channel, out_channel) 169 | 170 | rhs = graph_signal.permute(0, 2, 1).matmul(T_k).permute(0, 2, 1) 171 | 172 | output = output + rhs.matmul(theta_k) 173 | 174 | outputs.append(output.unsqueeze(-1)) 175 | 176 | return F.relu(torch.cat(outputs, dim=-1)) 177 | 178 | 179 | class ASTGCN_block(nn.Module): 180 | 181 | def __init__(self, DEVICE, in_channels, K, nb_chev_filter, nb_time_filter, time_strides, cheb_polynomials, 182 | num_of_vertices, num_of_timesteps): 183 | super(ASTGCN_block, self).__init__() 184 | self.TAt = Temporal_Attention_layer(DEVICE, in_channels, num_of_vertices, num_of_timesteps) 185 | self.SAt = Spatial_Attention_layer(DEVICE, in_channels, num_of_vertices, num_of_timesteps) 186 | self.cheb_conv_SAt = cheb_conv_withSAt(K, cheb_polynomials, in_channels, nb_chev_filter) 187 | self.time_conv = nn.Conv2d(nb_chev_filter, nb_time_filter, kernel_size=(1, 3), stride=(1, time_strides), 188 | padding=(0, 1)) 189 | self.residual_conv = nn.Conv2d(in_channels, nb_time_filter, kernel_size=(1, 1), stride=(1, time_strides)) 190 | self.ln = nn.LayerNorm(nb_time_filter) # 需要将channel放到最后一个维度上 191 | 192 | def forward(self, x): 193 | """ 194 | :param x: (batch_size, N, F_in, T) 195 | :return: (batch_size, N, nb_time_filter, T) 196 | """ 197 | batch_size, num_of_vertices, num_of_features, num_of_timesteps = x.shape 198 | 199 | # TAt 200 | temporal_At = self.TAt(x) # (b, T, T) 201 | 202 | x_TAt = torch.matmul(x.reshape(batch_size, -1, num_of_timesteps), temporal_At).reshape(batch_size, 203 | num_of_vertices, 204 | num_of_features, 205 | num_of_timesteps) 206 | # SAt 207 | spatial_At = self.SAt(x_TAt) 208 | 209 | # cheb gcn 210 | spatial_gcn = self.cheb_conv_SAt(x, spatial_At) # (b,N,F,T) 211 | # spatial_gcn = self.cheb_conv(x) 212 | 213 | # convolution along the time axis 214 | time_conv_output = self.time_conv( 215 | spatial_gcn.permute(0, 2, 1, 3)) # (b,N,F,T)->(b,F,N,T) 用(1,3)的卷积核去做->(b,F,N,T) 216 | 217 | # residual shortcut 218 | x_residual = self.residual_conv(x.permute(0, 2, 1, 3)) # (b,N,F,T)->(b,F,N,T) 用(1,1)的卷积核去做->(b,F,N,T) 219 | 220 | x_residual = self.ln(F.relu(x_residual + time_conv_output).permute(0, 3, 2, 1)).permute(0, 2, 3, 1) 221 | # (b,F,N,T)->(b,T,N,F) -ln-> (b,T,N,F)->(b,N,F,T) 222 | 223 | return x_residual 224 | 225 | 226 | class ASTGCN_submodule(nn.Module): 227 | 228 | def __init__(self, DEVICE, nb_block, in_channels, K, nb_chev_filter, nb_time_filter, time_strides, cheb_polynomials, 229 | num_for_predict, len_input, num_of_vertices): 230 | """ 231 | :param nb_block: 232 | :param in_channels: 233 | :param K: 234 | :param nb_chev_filter: 235 | :param nb_time_filter: 236 | :param time_strides: 237 | :param cheb_polynomials: 238 | :param nb_predict_step: 239 | """ 240 | 241 | super(ASTGCN_submodule, self).__init__() 242 | 243 | self.BlockList = nn.ModuleList([ASTGCN_block(DEVICE, in_channels, K, nb_chev_filter, nb_time_filter, 244 | time_strides, cheb_polynomials, num_of_vertices, len_input)]) 245 | 246 | self.BlockList.extend([ASTGCN_block(DEVICE, nb_time_filter, K, nb_chev_filter, nb_time_filter, 1, 247 | cheb_polynomials, num_of_vertices, len_input // time_strides) for _ in 248 | range(nb_block - 1)]) 249 | 250 | self.final_conv = nn.Conv2d(int(len_input / time_strides), num_for_predict, 251 | kernel_size=(1, nb_time_filter - in_channels + 1)) 252 | 253 | self.DEVICE = DEVICE 254 | 255 | self.to(DEVICE) 256 | 257 | def forward(self, x): 258 | """ 259 | :param x: (B, N_nodes, F_in, T_in) 260 | :return: (B, N_nodes, T_out) 261 | """ 262 | for block in self.BlockList: 263 | x = block(x) 264 | 265 | output = self.final_conv(x.permute(0, 3, 1, 2)) 266 | 267 | return output 268 | 269 | 270 | def make_model(DEVICE, nb_block, in_channels, K, nb_chev_filter, nb_time_filter, time_strides, adj_mx, num_for_predict, 271 | len_input, num_of_vertices): 272 | """ 273 | 274 | :param DEVICE: 275 | :param nb_block: 276 | :param in_channels: 277 | :param K: 278 | :param nb_chev_filter: 279 | :param nb_time_filter: 280 | :param time_strides: 281 | :param cheb_polynomials: 282 | :param nb_predict_step: 283 | :param len_input 284 | :return: 285 | """ 286 | L_tilde = calculate_scaled_laplacian_torch(adj_mx) 287 | cheb_polynomials = [torch.from_numpy(i).type(torch.FloatTensor).to(DEVICE) for i in get_Tk(L_tilde, K)] 288 | model = ASTGCN_submodule(DEVICE, nb_block, in_channels, K, nb_chev_filter, nb_time_filter, time_strides, 289 | cheb_polynomials, num_for_predict, len_input, num_of_vertices) 290 | 291 | for p in model.parameters(): 292 | if p.dim() > 1: 293 | nn.init.xavier_uniform_(p) 294 | else: 295 | nn.init.uniform_(p) 296 | 297 | return model 298 | -------------------------------------------------------------------------------- /models/stgcn/layers.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import torch.nn.init as init 6 | 7 | 8 | class Align(nn.Module): 9 | def __init__(self, c_in, c_out): 10 | super(Align, self).__init__() 11 | self.c_in = c_in 12 | self.c_out = c_out 13 | self.align_conv = nn.Conv2d(in_channels=c_in, out_channels=c_out, kernel_size=(1, 1)) 14 | 15 | def forward(self, x): 16 | if self.c_in > self.c_out: 17 | x = self.align_conv(x) 18 | elif self.c_in < self.c_out: 19 | batch_size, _, timestep, n_vertex = x.shape 20 | x = torch.cat([x, torch.zeros([batch_size, self.c_out - self.c_in, timestep, n_vertex]).to(x)], dim=1) 21 | else: 22 | x = x 23 | 24 | return x 25 | 26 | 27 | class CausalConv1d(nn.Conv1d): 28 | def __init__(self, in_channels, out_channels, kernel_size, stride=1, enable_padding=False, dilation=1, groups=1, 29 | bias=True): 30 | if enable_padding == True: 31 | self.__padding = (kernel_size - 1) * dilation 32 | else: 33 | self.__padding = 0 34 | super(CausalConv1d, self).__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride, 35 | padding=self.__padding, dilation=dilation, groups=groups, bias=bias) 36 | 37 | def forward(self, input): 38 | result = super(CausalConv1d, self).forward(input) 39 | if self.__padding != 0: 40 | return result[:, :, : -self.__padding] 41 | 42 | return result 43 | 44 | 45 | class CausalConv2d(nn.Conv2d): 46 | def __init__(self, in_channels, out_channels, kernel_size, stride=1, enable_padding=False, dilation=1, groups=1, 47 | bias=True): 48 | kernel_size = nn.modules.utils._pair(kernel_size) 49 | stride = nn.modules.utils._pair(stride) 50 | dilation = nn.modules.utils._pair(dilation) 51 | if enable_padding == True: 52 | self.__padding = [int((kernel_size[i] - 1) * dilation[i]) for i in range(len(kernel_size))] 53 | else: 54 | self.__padding = 0 55 | self.left_padding = nn.modules.utils._pair(self.__padding) 56 | super(CausalConv2d, self).__init__(in_channels, out_channels, kernel_size, stride=stride, padding=0, 57 | dilation=dilation, groups=groups, bias=bias) 58 | 59 | def forward(self, input): 60 | if self.__padding != 0: 61 | input = F.pad(input, (self.left_padding[1], 0, self.left_padding[0], 0)) 62 | result = super(CausalConv2d, self).forward(input) 63 | 64 | return result 65 | 66 | 67 | class TemporalConvLayer(nn.Module): 68 | 69 | # Temporal Convolution Layer (GLU) 70 | # 71 | # |--------------------------------| * Residual Connection * 72 | # | | 73 | # | |--->--- CasualConv2d ----- + -------| 74 | # -------|----| ⊙ ------> 75 | # |--->--- CasualConv2d --- Sigmoid ---| 76 | # 77 | 78 | # param x: tensor, [bs, c_in, ts, n_vertex] 79 | 80 | def __init__(self, Kt, c_in, c_out, n_vertex, act_func): 81 | super(TemporalConvLayer, self).__init__() 82 | self.Kt = Kt 83 | self.c_in = c_in 84 | self.c_out = c_out 85 | self.n_vertex = n_vertex 86 | self.align = Align(c_in, c_out) 87 | if act_func == 'glu' or act_func == 'gtu': 88 | self.causal_conv = CausalConv2d(in_channels=c_in, out_channels=2 * c_out, kernel_size=(Kt, 1), 89 | enable_padding=False, dilation=1) 90 | else: 91 | self.causal_conv = CausalConv2d(in_channels=c_in, out_channels=c_out, kernel_size=(Kt, 1), 92 | enable_padding=False, dilation=1) 93 | self.act_func = act_func 94 | self.sigmoid = nn.Sigmoid() 95 | self.tanh = nn.Tanh() 96 | self.relu = nn.ReLU() 97 | self.leaky_relu = nn.LeakyReLU() 98 | self.silu = nn.SiLU() 99 | 100 | def forward(self, x): 101 | x_in = self.align(x)[:, :, self.Kt - 1:, :] 102 | x_causal_conv = self.causal_conv(x) 103 | 104 | if self.act_func == 'glu' or self.act_func == 'gtu': 105 | x_p = x_causal_conv[:, : self.c_out, :, :] 106 | x_q = x_causal_conv[:, -self.c_out:, :, :] 107 | 108 | if self.act_func == 'glu': 109 | # GLU was first purposed in 110 | # *Language Modeling with Gated Convolutional Networks*. 111 | # URL: https://arxiv.org/abs/1612.08083 112 | # Input tensor X is split by a certain dimension into tensor X_a and X_b. 113 | # In PyTorch, GLU is defined as X_a ⊙ Sigmoid(X_b). 114 | # URL: https://pytorch.org/docs/master/nn.functional.html#torch.nn.functional.glu 115 | # (x_p + x_in) ⊙ Sigmoid(x_q) 116 | x = torch.mul((x_p + x_in), self.sigmoid(x_q)) 117 | 118 | else: 119 | # Tanh(x_p + x_in) ⊙ Sigmoid(x_q) 120 | x = torch.mul(self.tanh(x_p + x_in), self.sigmoid(x_q)) 121 | 122 | elif self.act_func == 'relu': 123 | x = self.relu(x_causal_conv + x_in) 124 | 125 | elif self.act_func == 'leaky_relu': 126 | x = self.leaky_relu(x_causal_conv + x_in) 127 | 128 | elif self.act_func == 'silu': 129 | x = self.silu(x_causal_conv + x_in) 130 | 131 | else: 132 | raise NotImplementedError(f'ERROR: The activation function {self.act_func} is not implemented.') 133 | 134 | return x 135 | 136 | 137 | class ChebGraphConv(nn.Module): 138 | def __init__(self, c_in, c_out, Ks, gso, bias): 139 | super(ChebGraphConv, self).__init__() 140 | self.c_in = c_in 141 | self.c_out = c_out 142 | self.Ks = Ks 143 | self.gso = gso 144 | self.weight = nn.Parameter(torch.FloatTensor(Ks, c_in, c_out)) 145 | if bias: 146 | self.bias = nn.Parameter(torch.FloatTensor(c_out)) 147 | else: 148 | self.register_parameter('bias', None) 149 | self.reset_parameters() 150 | 151 | def reset_parameters(self): 152 | init.kaiming_uniform_(self.weight, a=math.sqrt(5)) 153 | if self.bias is not None: 154 | fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) 155 | bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 156 | init.uniform_(self.bias, -bound, bound) 157 | 158 | def forward(self, x): 159 | # bs, c_in, ts, n_vertex = x.shape 160 | x = torch.permute(x, (0, 2, 3, 1)) 161 | 162 | if self.Ks - 1 < 0: 163 | raise ValueError( 164 | f'ERROR: the graph convolution kernel size Ks has to be a positive integer, but received {self.Ks}.') 165 | elif self.Ks - 1 == 0: 166 | x_0 = x 167 | x_list = [x_0] 168 | elif self.Ks - 1 == 1: 169 | x_0 = x 170 | x_1 = torch.einsum('hi,btij->bthj', self.gso, x) 171 | x_list = [x_0, x_1] 172 | elif self.Ks - 1 >= 2: 173 | x_0 = x 174 | x_1 = torch.einsum('hi,btij->bthj', self.gso, x) 175 | x_list = [x_0, x_1] 176 | for k in range(2, self.Ks): 177 | x_list.append(torch.einsum('hi,btij->bthj', 2 * self.gso, x_list[k - 1]) - x_list[k - 2]) 178 | 179 | x = torch.stack(x_list, dim=2) 180 | 181 | cheb_graph_conv = torch.einsum('btkhi,kij->bthj', x, self.weight) 182 | 183 | if self.bias is not None: 184 | cheb_graph_conv = torch.add(cheb_graph_conv, self.bias) 185 | else: 186 | cheb_graph_conv = cheb_graph_conv 187 | 188 | return cheb_graph_conv 189 | 190 | 191 | class GraphConv(nn.Module): 192 | def __init__(self, c_in, c_out, gso, bias): 193 | super(GraphConv, self).__init__() 194 | self.c_in = c_in 195 | self.c_out = c_out 196 | self.gso = gso 197 | self.weight = nn.Parameter(torch.FloatTensor(c_in, c_out)) 198 | if bias: 199 | self.bias = nn.Parameter(torch.FloatTensor(c_out)) 200 | else: 201 | self.register_parameter('bias', None) 202 | self.reset_parameters() 203 | 204 | def reset_parameters(self): 205 | init.kaiming_uniform_(self.weight, a=math.sqrt(5)) 206 | if self.bias is not None: 207 | fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) 208 | bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 209 | init.uniform_(self.bias, -bound, bound) 210 | 211 | def forward(self, x): 212 | # bs, c_in, ts, n_vertex = x.shape 213 | x = torch.permute(x, (0, 2, 3, 1)) 214 | 215 | first_mul = torch.einsum('hi,btij->bthj', self.gso, x) 216 | second_mul = torch.einsum('bthi,ij->bthj', first_mul, self.weight) 217 | 218 | if self.bias is not None: 219 | graph_conv = torch.add(second_mul, self.bias) 220 | else: 221 | graph_conv = second_mul 222 | 223 | return graph_conv 224 | 225 | 226 | class GraphConvLayer(nn.Module): 227 | def __init__(self, graph_conv_type, c_in, c_out, Ks, gso, bias): 228 | super(GraphConvLayer, self).__init__() 229 | self.graph_conv_type = graph_conv_type 230 | self.c_in = c_in 231 | self.c_out = c_out 232 | self.align = Align(c_in, c_out) 233 | self.Ks = Ks 234 | self.gso = gso 235 | if self.graph_conv_type == 'cheb_graph_conv': 236 | self.cheb_graph_conv = ChebGraphConv(c_out, c_out, Ks, gso, bias) 237 | elif self.graph_conv_type == 'graph_conv': 238 | self.graph_conv = GraphConv(c_out, c_out, gso, bias) 239 | 240 | def forward(self, x): 241 | x_gc_in = self.align(x) 242 | if self.graph_conv_type == 'cheb_graph_conv': 243 | x_gc = self.cheb_graph_conv(x_gc_in) 244 | elif self.graph_conv_type == 'graph_conv': 245 | x_gc = self.graph_conv(x_gc_in) 246 | x_gc = x_gc.permute(0, 3, 1, 2) 247 | x_gc_out = torch.add(x_gc, x_gc_in) 248 | 249 | return x_gc_out 250 | 251 | 252 | class STConvBlock(nn.Module): 253 | # STConv Block contains 'TGTND' structure 254 | # T: Gated Temporal Convolution Layer (GLU or GTU) 255 | # G: Graph Convolution Layer (ChebGraphConv or GraphConv) 256 | # T: Gated Temporal Convolution Layer (GLU or GTU) 257 | # N: Layer Normolization 258 | # D: Dropout 259 | 260 | def __init__(self, Kt, Ks, n_vertex, last_block_channel, channels, act_func, graph_conv_type, gso, bias, droprate): 261 | super(STConvBlock, self).__init__() 262 | self.tmp_conv1 = TemporalConvLayer(Kt, last_block_channel, channels[0], n_vertex, act_func) 263 | self.graph_conv = GraphConvLayer(graph_conv_type, channels[0], channels[1], Ks, gso, bias) 264 | self.tmp_conv2 = TemporalConvLayer(Kt, channels[1], channels[2], n_vertex, act_func) 265 | self.tc2_ln = nn.LayerNorm([n_vertex, channels[2]]) 266 | self.relu = nn.ReLU() 267 | self.dropout = nn.Dropout(p=droprate) 268 | 269 | def forward(self, x): 270 | x = self.tmp_conv1(x) 271 | x = self.graph_conv(x) 272 | x = self.relu(x) 273 | x = self.tmp_conv2(x) 274 | x = self.tc2_ln(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) 275 | x = self.dropout(x) 276 | 277 | return x 278 | 279 | 280 | class OutputBlock(nn.Module): 281 | # Output block contains 'TNFF' structure 282 | # T: Gated Temporal Convolution Layer (GLU or GTU) 283 | # N: Layer Normolization 284 | # F: Fully-Connected Layer 285 | # F: Fully-Connected Layer 286 | 287 | def __init__(self, Ko, last_block_channel, channels, end_channel, n_vertex, act_func, bias, droprate): 288 | super(OutputBlock, self).__init__() 289 | self.tmp_conv1 = TemporalConvLayer(Ko, last_block_channel, channels[0], n_vertex, act_func) 290 | self.fc1 = nn.Linear(in_features=channels[0], out_features=channels[1], bias=bias) 291 | self.fc2 = nn.Linear(in_features=channels[1], out_features=end_channel, bias=bias) 292 | self.tc1_ln = nn.LayerNorm([n_vertex, channels[0]]) 293 | self.relu = nn.ReLU() 294 | self.leaky_relu = nn.LeakyReLU() 295 | self.silu = nn.SiLU() 296 | self.dropout = nn.Dropout(p=droprate) 297 | 298 | def forward(self, x): 299 | x = self.tmp_conv1(x) 300 | x = self.tc1_ln(x.permute(0, 2, 3, 1)) 301 | x = self.fc1(x) 302 | x = self.relu(x) 303 | x = self.fc2(x).permute(0, 3, 1, 2) 304 | 305 | return x 306 | --------------------------------------------------------------------------------