├── LICENSE
├── README.md
├── data
    └── sensor_graph
    │   ├── adj_mx.pkl
    │   └── adj_mx_bay.pkl
├── generate_training_data.py
├── layer.py
├── net.py
├── requirements.txt
├── train_multi_step.py
├── train_single_step.py
├── trainer.py
└── util.py


/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2020 Zonghan Wu
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # MTGNN
 2 | This is a PyTorch implementation of the paper: [Connecting the Dots: Multivariate Time Series Forecasting with Graph Neural Networks](https://arxiv.org/abs/2005.11650), published in KDD-2020.
 3 | 
 4 | ## Requirements
 5 | The model is implemented using Python3 with dependencies specified in requirements.txt
 6 | ## Data Preparation
 7 | ### Multivariate time series datasets
 8 | 
 9 | Download Solar-Energy, Traffic, Electricity, Exchange-rate datasets from [https://github.com/laiguokun/multivariate-time-series-data](https://github.com/laiguokun/multivariate-time-series-data). Uncompress them and move them to the data folder.
10 | 
11 | ### Traffic datasets
12 | Download the METR-LA and PEMS-BAY dataset from [Google Drive](https://drive.google.com/open?id=10FOTa6HXPqX8Pf5WRoRwcFnW9BrNZEIX) or [Baidu Yun](https://pan.baidu.com/s/14Yy9isAIZYdU__OYEQGa_g) provided by [Li et al.](https://github.com/liyaguang/DCRNN.git) . Move them into the data folder. 
13 | 
14 | ```
15 | 
16 | # Create data directories
17 | mkdir -p data/{METR-LA,PEMS-BAY}
18 | 
19 | # METR-LA
20 | python generate_training_data.py --output_dir=data/METR-LA --traffic_df_filename=data/metr-la.h5
21 | 
22 | # PEMS-BAY
23 | python generate_training_data.py --output_dir=data/PEMS-BAY --traffic_df_filename=data/pems-bay.h5
24 | 
25 | ```
26 | 
27 | ## Model Training
28 | 
29 | ### Single-step
30 | 
31 | * Solar-Energy
32 | 
33 | ```
34 | python train_single_step.py --save ./model-solar-3.pt --data ./data/solar_AL.txt --num_nodes 137 --batch_size 4 --epochs 30 --horizon 3
35 | #sampling
36 | python train_single_step.py --num_split 3 --save ./model-solar-sampling-3.pt --data ./data/solar_AL.txt --num_nodes 137 --batch_size 16 --epochs 30 --horizon 3
37 | ```
38 | * Traffic 
39 | 
40 | ```
41 | python train_single_step.py --save ./model-traffic3.pt --data ./data/traffic.txt --num_nodes 862 --batch_size 16 --epochs 30 --horizon 3
42 | #sampling
43 | python train_single_step.py --num_split 3 --save ./model-traffic-sampling-3.pt --data ./data/traffic --num_nodes 321 --batch_size 16 --epochs 30 --horizon 3
44 | ```
45 | 
46 | * Electricity
47 | 
48 | ```
49 | python train_single_step.py --save ./model-electricity-3.pt --data ./data/electricity.txt --num_nodes 321 --batch_size 4 --epochs 30 --horizon 3
50 | #sampling 
51 | python train_single_step.py --num_split 3 --save ./model-electricity-sampling-3.pt --data ./data/electricity.txt --num_nodes 321 --batch_size 16 --epochs 30 --horizon 3
52 | ```
53 | 
54 | * Exchange-Rate
55 | 
56 | ```
57 | python train_single_step.py --save ./model/model-exchange-3.pt --data ./data/exchange_rate.txt --num_nodes 8 --subgraph_size 8  --batch_size 4 --epochs 30 --horizon 3
58 | #sampling
59 | python train_single_step.py --num_split 3 --save ./model-exchange-3.pt --data ./data/exchange_rate.txt --num_nodes 8 --subgraph_size 2  --batch_size 16 --epochs 30 --horizon 3
60 | ```
61 | ### Multi-step
62 | * METR-LA
63 | 
64 | ```
65 | python train_multi_step.py --adj_data ./data/sensor_graph/adj_mx.pkl --data ./data/METR-LA --num_nodes 207
66 | ```
67 | * PEMS-BAY
68 | 
69 | ```
70 | python train_multi_step.py --adj_data ./data/sensor_graph/adj_mx_bay.pkl --data ./data/PEMS-BAY/ --num_nodes 325
71 | ```
72 | 
73 | ## Citation
74 | 
75 | ```
76 | @inproceedings{wu2020connecting,
77 |   title={Connecting the Dots: Multivariate Time Series Forecasting with Graph Neural Networks},
78 |   author={Wu, Zonghan and Pan, Shirui and Long, Guodong and Jiang, Jing and Chang, Xiaojun and Zhang, Chengqi},
79 |   booktitle={Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery \& Data Mining},
80 |   year={2020}
81 | }
82 | ```
83 | 


--------------------------------------------------------------------------------
/generate_training_data.py:
--------------------------------------------------------------------------------
  1 | from __future__ import absolute_import
  2 | from __future__ import division
  3 | from __future__ import print_function
  4 | from __future__ import unicode_literals
  5 | 
  6 | import argparse
  7 | import numpy as np
  8 | import os
  9 | import pandas as pd
 10 | 
 11 | 
 12 | def generate_graph_seq2seq_io_data(
 13 |         df, x_offsets, y_offsets, add_time_in_day=True, add_day_in_week=False, scaler=None
 14 | ):
 15 |     """
 16 |     Generate samples from
 17 |     :param df:
 18 |     :param x_offsets:
 19 |     :param y_offsets:
 20 |     :param add_time_in_day:
 21 |     :param add_day_in_week:
 22 |     :param scaler:
 23 |     :return:
 24 |     # x: (epoch_size, input_length, num_nodes, input_dim)
 25 |     # y: (epoch_size, output_length, num_nodes, output_dim)
 26 |     """
 27 | 
 28 |     num_samples, num_nodes = df.shape
 29 |     data = np.expand_dims(df.values, axis=-1)
 30 |     data_list = [data]
 31 |     if add_time_in_day:
 32 |         time_ind = (df.index.values - df.index.values.astype("datetime64[D]")) / np.timedelta64(1, "D")
 33 |         time_in_day = np.tile(time_ind, [1, num_nodes, 1]).transpose((2, 1, 0))
 34 |         data_list.append(time_in_day)
 35 |     if add_day_in_week:
 36 |         day_in_week = np.zeros(shape=(num_samples, num_nodes, 7))
 37 |         day_in_week[np.arange(num_samples), :, df.index.dayofweek] = 1
 38 |         data_list.append(day_in_week)
 39 | 
 40 |     data = np.concatenate(data_list, axis=-1)
 41 |     # epoch_len = num_samples + min(x_offsets) - max(y_offsets)
 42 |     x, y = [], []
 43 |     # t is the index of the last observation.
 44 |     min_t = abs(min(x_offsets))
 45 |     max_t = abs(num_samples - abs(max(y_offsets)))  # Exclusive
 46 |     for t in range(min_t, max_t):
 47 |         x_t = data[t + x_offsets, ...]
 48 |         y_t = data[t + y_offsets, ...]
 49 |         x.append(x_t)
 50 |         y.append(y_t)
 51 |     x = np.stack(x, axis=0)
 52 |     y = np.stack(y, axis=0)
 53 |     return x, y
 54 | 
 55 | 
 56 | def generate_train_val_test(args):
 57 |     df = pd.read_hdf(args.traffic_df_filename)
 58 |     # 0 is the latest observed sample.
 59 |     x_offsets = np.sort(
 60 |         # np.concatenate(([-week_size + 1, -day_size + 1], np.arange(-11, 1, 1)))
 61 |         np.concatenate((np.arange(-11, 1, 1),))
 62 |     )
 63 |     # Predict the next one hour
 64 |     y_offsets = np.sort(np.arange(1, 13, 1))
 65 |     # x: (num_samples, input_length, num_nodes, input_dim)
 66 |     # y: (num_samples, output_length, num_nodes, output_dim)
 67 |     x, y = generate_graph_seq2seq_io_data(
 68 |         df,
 69 |         x_offsets=x_offsets,
 70 |         y_offsets=y_offsets,
 71 |         add_time_in_day=True,
 72 |         add_day_in_week=False,
 73 |     )
 74 | 
 75 |     print("x shape: ", x.shape, ", y shape: ", y.shape)
 76 |     # Write the data into npz file.
 77 |     # num_test = 6831, using the last 6831 examples as testing.
 78 |     # for the rest: 7/8 is used for training, and 1/8 is used for validation.
 79 |     num_samples = x.shape[0]
 80 |     num_test = round(num_samples * 0.2)
 81 |     num_train = round(num_samples * 0.7)
 82 |     num_val = num_samples - num_test - num_train
 83 | 
 84 |     # train
 85 |     x_train, y_train = x[:num_train], y[:num_train]
 86 |     # val
 87 |     x_val, y_val = (
 88 |         x[num_train: num_train + num_val],
 89 |         y[num_train: num_train + num_val],
 90 |     )
 91 |     # test
 92 |     x_test, y_test = x[-num_test:], y[-num_test:]
 93 | 
 94 |     for cat in ["train", "val", "test"]:
 95 |         _x, _y = locals()["x_" + cat], locals()["y_" + cat]
 96 |         print(cat, "x: ", _x.shape, "y:", _y.shape)
 97 |         np.savez_compressed(
 98 |             os.path.join(args.output_dir, "%s.npz" % cat),
 99 |             x=_x,
100 |             y=_y,
101 |             x_offsets=x_offsets.reshape(list(x_offsets.shape) + [1]),
102 |             y_offsets=y_offsets.reshape(list(y_offsets.shape) + [1]),
103 |         )
104 | 
105 | 
106 | def main(args):
107 |     print("Generating training data")
108 |     generate_train_val_test(args)
109 | 
110 | 
111 | if __name__ == "__main__":
112 |     parser = argparse.ArgumentParser()
113 |     parser.add_argument(
114 |         "--output_dir", type=str, default="data/", help="Output directory."
115 |     )
116 |     parser.add_argument(
117 |         "--traffic_df_filename",
118 |         type=str,
119 |         default="data/metr-la.h5",
120 |         help="Raw traffic readings.",
121 |     )
122 |     args = parser.parse_args()
123 |     main(args)
124 | 


--------------------------------------------------------------------------------
/layer.py:
--------------------------------------------------------------------------------
  1 | from __future__ import division
  2 | import torch
  3 | import torch.nn as nn
  4 | from torch.nn import init
  5 | import numbers
  6 | import torch.nn.functional as F
  7 | 
  8 | 
  9 | class nconv(nn.Module):
 10 |     def __init__(self):
 11 |         super(nconv,self).__init__()
 12 | 
 13 |     def forward(self,x, A):
 14 |         x = torch.einsum('ncwl,vw->ncvl',(x,A))
 15 |         return x.contiguous()
 16 | 
 17 | class dy_nconv(nn.Module):
 18 |     def __init__(self):
 19 |         super(dy_nconv,self).__init__()
 20 | 
 21 |     def forward(self,x, A):
 22 |         x = torch.einsum('ncvl,nvwl->ncwl',(x,A))
 23 |         return x.contiguous()
 24 | 
 25 | class linear(nn.Module):
 26 |     def __init__(self,c_in,c_out,bias=True):
 27 |         super(linear,self).__init__()
 28 |         self.mlp = torch.nn.Conv2d(c_in, c_out, kernel_size=(1, 1), padding=(0,0), stride=(1,1), bias=bias)
 29 | 
 30 |     def forward(self,x):
 31 |         return self.mlp(x)
 32 | 
 33 | 
 34 | class prop(nn.Module):
 35 |     def __init__(self,c_in,c_out,gdep,dropout,alpha):
 36 |         super(prop, self).__init__()
 37 |         self.nconv = nconv()
 38 |         self.mlp = linear(c_in,c_out)
 39 |         self.gdep = gdep
 40 |         self.dropout = dropout
 41 |         self.alpha = alpha
 42 | 
 43 |     def forward(self,x,adj):
 44 |         adj = adj + torch.eye(adj.size(0)).to(x.device)
 45 |         d = adj.sum(1)
 46 |         h = x
 47 |         dv = d
 48 |         a = adj / dv.view(-1, 1)
 49 |         for i in range(self.gdep):
 50 |             h = self.alpha*x + (1-self.alpha)*self.nconv(h,a)
 51 |         ho = self.mlp(h)
 52 |         return ho
 53 | 
 54 | 
 55 | class mixprop(nn.Module):
 56 |     def __init__(self,c_in,c_out,gdep,dropout,alpha):
 57 |         super(mixprop, self).__init__()
 58 |         self.nconv = nconv()
 59 |         self.mlp = linear((gdep+1)*c_in,c_out)
 60 |         self.gdep = gdep
 61 |         self.dropout = dropout
 62 |         self.alpha = alpha
 63 | 
 64 | 
 65 |     def forward(self,x,adj):
 66 |         adj = adj + torch.eye(adj.size(0)).to(x.device)
 67 |         d = adj.sum(1)
 68 |         h = x
 69 |         out = [h]
 70 |         a = adj / d.view(-1, 1)
 71 |         for i in range(self.gdep):
 72 |             h = self.alpha*x + (1-self.alpha)*self.nconv(h,a)
 73 |             out.append(h)
 74 |         ho = torch.cat(out,dim=1)
 75 |         ho = self.mlp(ho)
 76 |         return ho
 77 | 
 78 | class dy_mixprop(nn.Module):
 79 |     def __init__(self,c_in,c_out,gdep,dropout,alpha):
 80 |         super(dy_mixprop, self).__init__()
 81 |         self.nconv = dy_nconv()
 82 |         self.mlp1 = linear((gdep+1)*c_in,c_out)
 83 |         self.mlp2 = linear((gdep+1)*c_in,c_out)
 84 | 
 85 |         self.gdep = gdep
 86 |         self.dropout = dropout
 87 |         self.alpha = alpha
 88 |         self.lin1 = linear(c_in,c_in)
 89 |         self.lin2 = linear(c_in,c_in)
 90 | 
 91 | 
 92 |     def forward(self,x):
 93 |         #adj = adj + torch.eye(adj.size(0)).to(x.device)
 94 |         #d = adj.sum(1)
 95 |         x1 = torch.tanh(self.lin1(x))
 96 |         x2 = torch.tanh(self.lin2(x))
 97 |         adj = self.nconv(x1.transpose(2,1),x2)
 98 |         adj0 = torch.softmax(adj, dim=2)
 99 |         adj1 = torch.softmax(adj.transpose(2,1), dim=2)
100 | 
101 |         h = x
102 |         out = [h]
103 |         for i in range(self.gdep):
104 |             h = self.alpha*x + (1-self.alpha)*self.nconv(h,adj0)
105 |             out.append(h)
106 |         ho = torch.cat(out,dim=1)
107 |         ho1 = self.mlp1(ho)
108 | 
109 | 
110 |         h = x
111 |         out = [h]
112 |         for i in range(self.gdep):
113 |             h = self.alpha * x + (1 - self.alpha) * self.nconv(h, adj1)
114 |             out.append(h)
115 |         ho = torch.cat(out, dim=1)
116 |         ho2 = self.mlp2(ho)
117 | 
118 |         return ho1+ho2
119 | 
120 | 
121 | 
122 | class dilated_1D(nn.Module):
123 |     def __init__(self, cin, cout, dilation_factor=2):
124 |         super(dilated_1D, self).__init__()
125 |         self.tconv = nn.ModuleList()
126 |         self.kernel_set = [2,3,6,7]
127 |         self.tconv = nn.Conv2d(cin,cout,(1,7),dilation=(1,dilation_factor))
128 | 
129 |     def forward(self,input):
130 |         x = self.tconv(input)
131 |         return x
132 | 
133 | class dilated_inception(nn.Module):
134 |     def __init__(self, cin, cout, dilation_factor=2):
135 |         super(dilated_inception, self).__init__()
136 |         self.tconv = nn.ModuleList()
137 |         self.kernel_set = [2,3,6,7]
138 |         cout = int(cout/len(self.kernel_set))
139 |         for kern in self.kernel_set:
140 |             self.tconv.append(nn.Conv2d(cin,cout,(1,kern),dilation=(1,dilation_factor)))
141 | 
142 |     def forward(self,input):
143 |         x = []
144 |         for i in range(len(self.kernel_set)):
145 |             x.append(self.tconv[i](input))
146 |         for i in range(len(self.kernel_set)):
147 |             x[i] = x[i][...,-x[-1].size(3):]
148 |         x = torch.cat(x,dim=1)
149 |         return x
150 | 
151 | 
152 | class graph_constructor(nn.Module):
153 |     def __init__(self, nnodes, k, dim, device, alpha=3, static_feat=None):
154 |         super(graph_constructor, self).__init__()
155 |         self.nnodes = nnodes
156 |         if static_feat is not None:
157 |             xd = static_feat.shape[1]
158 |             self.lin1 = nn.Linear(xd, dim)
159 |             self.lin2 = nn.Linear(xd, dim)
160 |         else:
161 |             self.emb1 = nn.Embedding(nnodes, dim)
162 |             self.emb2 = nn.Embedding(nnodes, dim)
163 |             self.lin1 = nn.Linear(dim,dim)
164 |             self.lin2 = nn.Linear(dim,dim)
165 | 
166 |         self.device = device
167 |         self.k = k
168 |         self.dim = dim
169 |         self.alpha = alpha
170 |         self.static_feat = static_feat
171 | 
172 |     def forward(self, idx):
173 |         if self.static_feat is None:
174 |             nodevec1 = self.emb1(idx)
175 |             nodevec2 = self.emb2(idx)
176 |         else:
177 |             nodevec1 = self.static_feat[idx,:]
178 |             nodevec2 = nodevec1
179 | 
180 |         nodevec1 = torch.tanh(self.alpha*self.lin1(nodevec1))
181 |         nodevec2 = torch.tanh(self.alpha*self.lin2(nodevec2))
182 | 
183 |         a = torch.mm(nodevec1, nodevec2.transpose(1,0))-torch.mm(nodevec2, nodevec1.transpose(1,0))
184 |         adj = F.relu(torch.tanh(self.alpha*a))
185 |         mask = torch.zeros(idx.size(0), idx.size(0)).to(self.device)
186 |         mask.fill_(float('0'))
187 |         s1,t1 = (adj + torch.rand_like(adj)*0.01).topk(self.k,1)
188 |         mask.scatter_(1,t1,s1.fill_(1))
189 |         adj = adj*mask
190 |         return adj
191 | 
192 |     def fullA(self, idx):
193 |         if self.static_feat is None:
194 |             nodevec1 = self.emb1(idx)
195 |             nodevec2 = self.emb2(idx)
196 |         else:
197 |             nodevec1 = self.static_feat[idx,:]
198 |             nodevec2 = nodevec1
199 | 
200 |         nodevec1 = torch.tanh(self.alpha*self.lin1(nodevec1))
201 |         nodevec2 = torch.tanh(self.alpha*self.lin2(nodevec2))
202 | 
203 |         a = torch.mm(nodevec1, nodevec2.transpose(1,0))-torch.mm(nodevec2, nodevec1.transpose(1,0))
204 |         adj = F.relu(torch.tanh(self.alpha*a))
205 |         return adj
206 | 
207 | class graph_global(nn.Module):
208 |     def __init__(self, nnodes, k, dim, device, alpha=3, static_feat=None):
209 |         super(graph_global, self).__init__()
210 |         self.nnodes = nnodes
211 |         self.A = nn.Parameter(torch.randn(nnodes, nnodes).to(device), requires_grad=True).to(device)
212 | 
213 |     def forward(self, idx):
214 |         return F.relu(self.A)
215 | 
216 | 
217 | class graph_undirected(nn.Module):
218 |     def __init__(self, nnodes, k, dim, device, alpha=3, static_feat=None):
219 |         super(graph_undirected, self).__init__()
220 |         self.nnodes = nnodes
221 |         if static_feat is not None:
222 |             xd = static_feat.shape[1]
223 |             self.lin1 = nn.Linear(xd, dim)
224 |         else:
225 |             self.emb1 = nn.Embedding(nnodes, dim)
226 |             self.lin1 = nn.Linear(dim,dim)
227 | 
228 |         self.device = device
229 |         self.k = k
230 |         self.dim = dim
231 |         self.alpha = alpha
232 |         self.static_feat = static_feat
233 | 
234 |     def forward(self, idx):
235 |         if self.static_feat is None:
236 |             nodevec1 = self.emb1(idx)
237 |             nodevec2 = self.emb1(idx)
238 |         else:
239 |             nodevec1 = self.static_feat[idx,:]
240 |             nodevec2 = nodevec1
241 | 
242 |         nodevec1 = torch.tanh(self.alpha*self.lin1(nodevec1))
243 |         nodevec2 = torch.tanh(self.alpha*self.lin1(nodevec2))
244 | 
245 |         a = torch.mm(nodevec1, nodevec2.transpose(1,0))
246 |         adj = F.relu(torch.tanh(self.alpha*a))
247 |         mask = torch.zeros(idx.size(0), idx.size(0)).to(self.device)
248 |         mask.fill_(float('0'))
249 |         s1,t1 = adj.topk(self.k,1)
250 |         mask.scatter_(1,t1,s1.fill_(1))
251 |         adj = adj*mask
252 |         return adj
253 | 
254 | 
255 | 
256 | class graph_directed(nn.Module):
257 |     def __init__(self, nnodes, k, dim, device, alpha=3, static_feat=None):
258 |         super(graph_directed, self).__init__()
259 |         self.nnodes = nnodes
260 |         if static_feat is not None:
261 |             xd = static_feat.shape[1]
262 |             self.lin1 = nn.Linear(xd, dim)
263 |             self.lin2 = nn.Linear(xd, dim)
264 |         else:
265 |             self.emb1 = nn.Embedding(nnodes, dim)
266 |             self.emb2 = nn.Embedding(nnodes, dim)
267 |             self.lin1 = nn.Linear(dim,dim)
268 |             self.lin2 = nn.Linear(dim,dim)
269 | 
270 |         self.device = device
271 |         self.k = k
272 |         self.dim = dim
273 |         self.alpha = alpha
274 |         self.static_feat = static_feat
275 | 
276 |     def forward(self, idx):
277 |         if self.static_feat is None:
278 |             nodevec1 = self.emb1(idx)
279 |             nodevec2 = self.emb2(idx)
280 |         else:
281 |             nodevec1 = self.static_feat[idx,:]
282 |             nodevec2 = nodevec1
283 | 
284 |         nodevec1 = torch.tanh(self.alpha*self.lin1(nodevec1))
285 |         nodevec2 = torch.tanh(self.alpha*self.lin2(nodevec2))
286 | 
287 |         a = torch.mm(nodevec1, nodevec2.transpose(1,0))
288 |         adj = F.relu(torch.tanh(self.alpha*a))
289 |         mask = torch.zeros(idx.size(0), idx.size(0)).to(self.device)
290 |         mask.fill_(float('0'))
291 |         s1,t1 = adj.topk(self.k,1)
292 |         mask.scatter_(1,t1,s1.fill_(1))
293 |         adj = adj*mask
294 |         return adj
295 | 
296 | 
297 | class LayerNorm(nn.Module):
298 |     __constants__ = ['normalized_shape', 'weight', 'bias', 'eps', 'elementwise_affine']
299 |     def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
300 |         super(LayerNorm, self).__init__()
301 |         if isinstance(normalized_shape, numbers.Integral):
302 |             normalized_shape = (normalized_shape,)
303 |         self.normalized_shape = tuple(normalized_shape)
304 |         self.eps = eps
305 |         self.elementwise_affine = elementwise_affine
306 |         if self.elementwise_affine:
307 |             self.weight = nn.Parameter(torch.Tensor(*normalized_shape))
308 |             self.bias = nn.Parameter(torch.Tensor(*normalized_shape))
309 |         else:
310 |             self.register_parameter('weight', None)
311 |             self.register_parameter('bias', None)
312 |         self.reset_parameters()
313 | 
314 | 
315 |     def reset_parameters(self):
316 |         if self.elementwise_affine:
317 |             init.ones_(self.weight)
318 |             init.zeros_(self.bias)
319 | 
320 |     def forward(self, input, idx):
321 |         if self.elementwise_affine:
322 |             return F.layer_norm(input, tuple(input.shape[1:]), self.weight[:,idx,:], self.bias[:,idx,:], self.eps)
323 |         else:
324 |             return F.layer_norm(input, tuple(input.shape[1:]), self.weight, self.bias, self.eps)
325 | 
326 |     def extra_repr(self):
327 |         return '{normalized_shape}, eps={eps}, ' \
328 |             'elementwise_affine={elementwise_affine}'.format(**self.__dict__)
329 | 


--------------------------------------------------------------------------------
/net.py:
--------------------------------------------------------------------------------
  1 | from layer import *
  2 | 
  3 | 
  4 | class gtnet(nn.Module):
  5 |     def __init__(self, gcn_true, buildA_true, gcn_depth, num_nodes, device, predefined_A=None, static_feat=None, dropout=0.3, subgraph_size=20, node_dim=40, dilation_exponential=1, conv_channels=32, residual_channels=32, skip_channels=64, end_channels=128, seq_length=12, in_dim=2, out_dim=12, layers=3, propalpha=0.05, tanhalpha=3, layer_norm_affline=True):
  6 |         super(gtnet, self).__init__()
  7 |         self.gcn_true = gcn_true
  8 |         self.buildA_true = buildA_true
  9 |         self.num_nodes = num_nodes
 10 |         self.dropout = dropout
 11 |         self.predefined_A = predefined_A
 12 |         self.filter_convs = nn.ModuleList()
 13 |         self.gate_convs = nn.ModuleList()
 14 |         self.residual_convs = nn.ModuleList()
 15 |         self.skip_convs = nn.ModuleList()
 16 |         self.gconv1 = nn.ModuleList()
 17 |         self.gconv2 = nn.ModuleList()
 18 |         self.norm = nn.ModuleList()
 19 |         self.start_conv = nn.Conv2d(in_channels=in_dim,
 20 |                                     out_channels=residual_channels,
 21 |                                     kernel_size=(1, 1))
 22 |         self.gc = graph_constructor(num_nodes, subgraph_size, node_dim, device, alpha=tanhalpha, static_feat=static_feat)
 23 | 
 24 |         self.seq_length = seq_length
 25 |         kernel_size = 7
 26 |         if dilation_exponential>1:
 27 |             self.receptive_field = int(1+(kernel_size-1)*(dilation_exponential**layers-1)/(dilation_exponential-1))
 28 |         else:
 29 |             self.receptive_field = layers*(kernel_size-1) + 1
 30 | 
 31 |         for i in range(1):
 32 |             if dilation_exponential>1:
 33 |                 rf_size_i = int(1 + i*(kernel_size-1)*(dilation_exponential**layers-1)/(dilation_exponential-1))
 34 |             else:
 35 |                 rf_size_i = i*layers*(kernel_size-1)+1
 36 |             new_dilation = 1
 37 |             for j in range(1,layers+1):
 38 |                 if dilation_exponential > 1:
 39 |                     rf_size_j = int(rf_size_i + (kernel_size-1)*(dilation_exponential**j-1)/(dilation_exponential-1))
 40 |                 else:
 41 |                     rf_size_j = rf_size_i+j*(kernel_size-1)
 42 | 
 43 |                 self.filter_convs.append(dilated_inception(residual_channels, conv_channels, dilation_factor=new_dilation))
 44 |                 self.gate_convs.append(dilated_inception(residual_channels, conv_channels, dilation_factor=new_dilation))
 45 |                 self.residual_convs.append(nn.Conv2d(in_channels=conv_channels,
 46 |                                                     out_channels=residual_channels,
 47 |                                                  kernel_size=(1, 1)))
 48 |                 if self.seq_length>self.receptive_field:
 49 |                     self.skip_convs.append(nn.Conv2d(in_channels=conv_channels,
 50 |                                                     out_channels=skip_channels,
 51 |                                                     kernel_size=(1, self.seq_length-rf_size_j+1)))
 52 |                 else:
 53 |                     self.skip_convs.append(nn.Conv2d(in_channels=conv_channels,
 54 |                                                     out_channels=skip_channels,
 55 |                                                     kernel_size=(1, self.receptive_field-rf_size_j+1)))
 56 | 
 57 |                 if self.gcn_true:
 58 |                     self.gconv1.append(mixprop(conv_channels, residual_channels, gcn_depth, dropout, propalpha))
 59 |                     self.gconv2.append(mixprop(conv_channels, residual_channels, gcn_depth, dropout, propalpha))
 60 | 
 61 |                 if self.seq_length>self.receptive_field:
 62 |                     self.norm.append(LayerNorm((residual_channels, num_nodes, self.seq_length - rf_size_j + 1),elementwise_affine=layer_norm_affline))
 63 |                 else:
 64 |                     self.norm.append(LayerNorm((residual_channels, num_nodes, self.receptive_field - rf_size_j + 1),elementwise_affine=layer_norm_affline))
 65 | 
 66 |                 new_dilation *= dilation_exponential
 67 | 
 68 |         self.layers = layers
 69 |         self.end_conv_1 = nn.Conv2d(in_channels=skip_channels,
 70 |                                              out_channels=end_channels,
 71 |                                              kernel_size=(1,1),
 72 |                                              bias=True)
 73 |         self.end_conv_2 = nn.Conv2d(in_channels=end_channels,
 74 |                                              out_channels=out_dim,
 75 |                                              kernel_size=(1,1),
 76 |                                              bias=True)
 77 |         if self.seq_length > self.receptive_field:
 78 |             self.skip0 = nn.Conv2d(in_channels=in_dim, out_channels=skip_channels, kernel_size=(1, self.seq_length), bias=True)
 79 |             self.skipE = nn.Conv2d(in_channels=residual_channels, out_channels=skip_channels, kernel_size=(1, self.seq_length-self.receptive_field+1), bias=True)
 80 | 
 81 |         else:
 82 |             self.skip0 = nn.Conv2d(in_channels=in_dim, out_channels=skip_channels, kernel_size=(1, self.receptive_field), bias=True)
 83 |             self.skipE = nn.Conv2d(in_channels=residual_channels, out_channels=skip_channels, kernel_size=(1, 1), bias=True)
 84 | 
 85 | 
 86 |         self.idx = torch.arange(self.num_nodes).to(device)
 87 | 
 88 | 
 89 |     def forward(self, input, idx=None):
 90 |         seq_len = input.size(3)
 91 |         assert seq_len==self.seq_length, 'input sequence length not equal to preset sequence length'
 92 | 
 93 |         if self.seq_length<self.receptive_field:
 94 |             input = nn.functional.pad(input,(self.receptive_field-self.seq_length,0,0,0))
 95 | 
 96 | 
 97 | 
 98 |         if self.gcn_true:
 99 |             if self.buildA_true:
100 |                 if idx is None:
101 |                     adp = self.gc(self.idx)
102 |                 else:
103 |                     adp = self.gc(idx)
104 |             else:
105 |                 adp = self.predefined_A
106 | 
107 |         x = self.start_conv(input)
108 |         skip = self.skip0(F.dropout(input, self.dropout, training=self.training))
109 |         for i in range(self.layers):
110 |             residual = x
111 |             filter = self.filter_convs[i](x)
112 |             filter = torch.tanh(filter)
113 |             gate = self.gate_convs[i](x)
114 |             gate = torch.sigmoid(gate)
115 |             x = filter * gate
116 |             x = F.dropout(x, self.dropout, training=self.training)
117 |             s = x
118 |             s = self.skip_convs[i](s)
119 |             skip = s + skip
120 |             if self.gcn_true:
121 |                 x = self.gconv1[i](x, adp)+self.gconv2[i](x, adp.transpose(1,0))
122 |             else:
123 |                 x = self.residual_convs[i](x)
124 | 
125 |             x = x + residual[:, :, :, -x.size(3):]
126 |             if idx is None:
127 |                 x = self.norm[i](x,self.idx)
128 |             else:
129 |                 x = self.norm[i](x,idx)
130 | 
131 |         skip = self.skipE(x) + skip
132 |         x = F.relu(skip)
133 |         x = F.relu(self.end_conv_1(x))
134 |         x = self.end_conv_2(x)
135 |         return x
136 | 


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | matplotlib==3.1.1
2 | numpy==1.17.4
3 | pandas==0.25.3
4 | scipy==1.4.1
5 | torch==1.2.0
6 | scikit_learn==0.23.1
7 | 


--------------------------------------------------------------------------------
/train_multi_step.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import numpy as np
  3 | import argparse
  4 | import time
  5 | from util import *
  6 | from trainer import Trainer
  7 | from net import gtnet
  8 | 
  9 | def str_to_bool(value):
 10 |     if isinstance(value, bool):
 11 |         return value
 12 |     if value.lower() in {'false', 'f', '0', 'no', 'n'}:
 13 |         return False
 14 |     elif value.lower() in {'true', 't', '1', 'yes', 'y'}:
 15 |         return True
 16 |     raise ValueError(f'{value} is not a valid boolean value')
 17 | 
 18 | 
 19 | parser = argparse.ArgumentParser()
 20 | 
 21 | parser.add_argument('--device',type=str,default='cuda:1',help='')
 22 | parser.add_argument('--data',type=str,default='data/METR-LA',help='data path')
 23 | 
 24 | parser.add_argument('--adj_data', type=str,default='data/sensor_graph/adj_mx.pkl',help='adj data path')
 25 | parser.add_argument('--gcn_true', type=str_to_bool, default=True, help='whether to add graph convolution layer')
 26 | parser.add_argument('--buildA_true', type=str_to_bool, default=True,help='whether to construct adaptive adjacency matrix')
 27 | parser.add_argument('--load_static_feature', type=str_to_bool, default=False,help='whether to load static feature')
 28 | parser.add_argument('--cl', type=str_to_bool, default=True,help='whether to do curriculum learning')
 29 | 
 30 | parser.add_argument('--gcn_depth',type=int,default=2,help='graph convolution depth')
 31 | parser.add_argument('--num_nodes',type=int,default=207,help='number of nodes/variables')
 32 | parser.add_argument('--dropout',type=float,default=0.3,help='dropout rate')
 33 | parser.add_argument('--subgraph_size',type=int,default=20,help='k')
 34 | parser.add_argument('--node_dim',type=int,default=40,help='dim of nodes')
 35 | parser.add_argument('--dilation_exponential',type=int,default=1,help='dilation exponential')
 36 | 
 37 | parser.add_argument('--conv_channels',type=int,default=32,help='convolution channels')
 38 | parser.add_argument('--residual_channels',type=int,default=32,help='residual channels')
 39 | parser.add_argument('--skip_channels',type=int,default=64,help='skip channels')
 40 | parser.add_argument('--end_channels',type=int,default=128,help='end channels')
 41 | 
 42 | 
 43 | parser.add_argument('--in_dim',type=int,default=2,help='inputs dimension')
 44 | parser.add_argument('--seq_in_len',type=int,default=12,help='input sequence length')
 45 | parser.add_argument('--seq_out_len',type=int,default=12,help='output sequence length')
 46 | 
 47 | parser.add_argument('--layers',type=int,default=3,help='number of layers')
 48 | parser.add_argument('--batch_size',type=int,default=64,help='batch size')
 49 | parser.add_argument('--learning_rate',type=float,default=0.001,help='learning rate')
 50 | parser.add_argument('--weight_decay',type=float,default=0.0001,help='weight decay rate')
 51 | parser.add_argument('--clip',type=int,default=5,help='clip')
 52 | parser.add_argument('--step_size1',type=int,default=2500,help='step_size')
 53 | parser.add_argument('--step_size2',type=int,default=100,help='step_size')
 54 | 
 55 | 
 56 | parser.add_argument('--epochs',type=int,default=100,help='')
 57 | parser.add_argument('--print_every',type=int,default=50,help='')
 58 | parser.add_argument('--seed',type=int,default=101,help='random seed')
 59 | parser.add_argument('--save',type=str,default='./save/',help='save path')
 60 | parser.add_argument('--expid',type=int,default=1,help='experiment id')
 61 | 
 62 | parser.add_argument('--propalpha',type=float,default=0.05,help='prop alpha')
 63 | parser.add_argument('--tanhalpha',type=float,default=3,help='adj alpha')
 64 | 
 65 | parser.add_argument('--num_split',type=int,default=1,help='number of splits for graphs')
 66 | 
 67 | parser.add_argument('--runs',type=int,default=10,help='number of runs')
 68 | 
 69 | 
 70 | 
 71 | args = parser.parse_args()
 72 | torch.set_num_threads(3)
 73 | 
 74 | 
 75 | def main(runid):
 76 |     # torch.manual_seed(args.seed)
 77 |     # torch.backends.cudnn.deterministic = True
 78 |     # torch.backends.cudnn.benchmark = False
 79 |     # np.random.seed(args.seed)
 80 |     #load data
 81 |     device = torch.device(args.device)
 82 |     dataloader = load_dataset(args.data, args.batch_size, args.batch_size, args.batch_size)
 83 |     scaler = dataloader['scaler']
 84 | 
 85 |     predefined_A = load_adj(args.adj_data)
 86 |     predefined_A = torch.tensor(predefined_A)-torch.eye(args.num_nodes)
 87 |     predefined_A = predefined_A.to(device)
 88 | 
 89 |     # if args.load_static_feature:
 90 |     #     static_feat = load_node_feature('data/sensor_graph/location.csv')
 91 |     # else:
 92 |     #     static_feat = None
 93 | 
 94 |     model = gtnet(args.gcn_true, args.buildA_true, args.gcn_depth, args.num_nodes,
 95 |                   device, predefined_A=predefined_A,
 96 |                   dropout=args.dropout, subgraph_size=args.subgraph_size,
 97 |                   node_dim=args.node_dim,
 98 |                   dilation_exponential=args.dilation_exponential,
 99 |                   conv_channels=args.conv_channels, residual_channels=args.residual_channels,
100 |                   skip_channels=args.skip_channels, end_channels= args.end_channels,
101 |                   seq_length=args.seq_in_len, in_dim=args.in_dim, out_dim=args.seq_out_len,
102 |                   layers=args.layers, propalpha=args.propalpha, tanhalpha=args.tanhalpha, layer_norm_affline=True)
103 | 
104 |     print(args)
105 |     print('The recpetive field size is', model.receptive_field)
106 |     nParams = sum([p.nelement() for p in model.parameters()])
107 |     print('Number of model parameters is', nParams)
108 | 
109 |     engine = Trainer(model, args.learning_rate, args.weight_decay, args.clip, args.step_size1, args.seq_out_len, scaler, device, args.cl)
110 | 
111 |     print("start training...",flush=True)
112 |     his_loss =[]
113 |     val_time = []
114 |     train_time = []
115 |     minl = 1e5
116 |     for i in range(1,args.epochs+1):
117 |         train_loss = []
118 |         train_mape = []
119 |         train_rmse = []
120 |         t1 = time.time()
121 |         dataloader['train_loader'].shuffle()
122 |         for iter, (x, y) in enumerate(dataloader['train_loader'].get_iterator()):
123 |             trainx = torch.Tensor(x).to(device)
124 |             trainx= trainx.transpose(1, 3)
125 |             trainy = torch.Tensor(y).to(device)
126 |             trainy = trainy.transpose(1, 3)
127 |             if iter%args.step_size2==0:
128 |                 perm = np.random.permutation(range(args.num_nodes))
129 |             num_sub = int(args.num_nodes/args.num_split)
130 |             for j in range(args.num_split):
131 |                 if j != args.num_split-1:
132 |                     id = perm[j * num_sub:(j + 1) * num_sub]
133 |                 else:
134 |                     id = perm[j * num_sub:]
135 |                 id = torch.tensor(id).to(device)
136 |                 tx = trainx[:, :, id, :]
137 |                 ty = trainy[:, :, id, :]
138 |                 metrics = engine.train(tx, ty[:,0,:,:],id)
139 |                 train_loss.append(metrics[0])
140 |                 train_mape.append(metrics[1])
141 |                 train_rmse.append(metrics[2])
142 |             if iter % args.print_every == 0 :
143 |                 log = 'Iter: {:03d}, Train Loss: {:.4f}, Train MAPE: {:.4f}, Train RMSE: {:.4f}'
144 |                 print(log.format(iter, train_loss[-1], train_mape[-1], train_rmse[-1]),flush=True)
145 |         t2 = time.time()
146 |         train_time.append(t2-t1)
147 |         #validation
148 |         valid_loss = []
149 |         valid_mape = []
150 |         valid_rmse = []
151 | 
152 |         s1 = time.time()
153 |         for iter, (x, y) in enumerate(dataloader['val_loader'].get_iterator()):
154 |             testx = torch.Tensor(x).to(device)
155 |             testx = testx.transpose(1, 3)
156 |             testy = torch.Tensor(y).to(device)
157 |             testy = testy.transpose(1, 3)
158 |             metrics = engine.eval(testx, testy[:,0,:,:])
159 |             valid_loss.append(metrics[0])
160 |             valid_mape.append(metrics[1])
161 |             valid_rmse.append(metrics[2])
162 |         s2 = time.time()
163 |         log = 'Epoch: {:03d}, Inference Time: {:.4f} secs'
164 |         print(log.format(i,(s2-s1)))
165 |         val_time.append(s2-s1)
166 |         mtrain_loss = np.mean(train_loss)
167 |         mtrain_mape = np.mean(train_mape)
168 |         mtrain_rmse = np.mean(train_rmse)
169 | 
170 |         mvalid_loss = np.mean(valid_loss)
171 |         mvalid_mape = np.mean(valid_mape)
172 |         mvalid_rmse = np.mean(valid_rmse)
173 |         his_loss.append(mvalid_loss)
174 | 
175 |         log = 'Epoch: {:03d}, Train Loss: {:.4f}, Train MAPE: {:.4f}, Train RMSE: {:.4f}, Valid Loss: {:.4f}, Valid MAPE: {:.4f}, Valid RMSE: {:.4f}, Training Time: {:.4f}/epoch'
176 |         print(log.format(i, mtrain_loss, mtrain_mape, mtrain_rmse, mvalid_loss, mvalid_mape, mvalid_rmse, (t2 - t1)),flush=True)
177 | 
178 |         if mvalid_loss<minl:
179 |             torch.save(engine.model.state_dict(), args.save + "exp" + str(args.expid) + "_" + str(runid) +".pth")
180 |             minl = mvalid_loss
181 | 
182 |     print("Average Training Time: {:.4f} secs/epoch".format(np.mean(train_time)))
183 |     print("Average Inference Time: {:.4f} secs".format(np.mean(val_time)))
184 | 
185 | 
186 |     bestid = np.argmin(his_loss)
187 |     engine.model.load_state_dict(torch.load(args.save + "exp" + str(args.expid) + "_" + str(runid) +".pth"))
188 | 
189 |     print("Training finished")
190 |     print("The valid loss on best model is", str(round(his_loss[bestid],4)))
191 | 
192 |     #valid data
193 |     outputs = []
194 |     realy = torch.Tensor(dataloader['y_val']).to(device)
195 |     realy = realy.transpose(1,3)[:,0,:,:]
196 | 
197 |     for iter, (x, y) in enumerate(dataloader['val_loader'].get_iterator()):
198 |         testx = torch.Tensor(x).to(device)
199 |         testx = testx.transpose(1,3)
200 |         with torch.no_grad():
201 |             preds = engine.model(testx)
202 |             preds = preds.transpose(1,3)
203 |         outputs.append(preds.squeeze())
204 | 
205 |     yhat = torch.cat(outputs,dim=0)
206 |     yhat = yhat[:realy.size(0),...]
207 | 
208 | 
209 |     pred = scaler.inverse_transform(yhat)
210 |     vmae, vmape, vrmse = metric(pred,realy)
211 | 
212 |     #test data
213 |     outputs = []
214 |     realy = torch.Tensor(dataloader['y_test']).to(device)
215 |     realy = realy.transpose(1, 3)[:, 0, :, :]
216 | 
217 |     for iter, (x, y) in enumerate(dataloader['test_loader'].get_iterator()):
218 |         testx = torch.Tensor(x).to(device)
219 |         testx = testx.transpose(1, 3)
220 |         with torch.no_grad():
221 |             preds = engine.model(testx)
222 |             preds = preds.transpose(1, 3)
223 |         outputs.append(preds.squeeze())
224 | 
225 |     yhat = torch.cat(outputs, dim=0)
226 |     yhat = yhat[:realy.size(0), ...]
227 | 
228 |     mae = []
229 |     mape = []
230 |     rmse = []
231 |     for i in range(args.seq_out_len):
232 |         pred = scaler.inverse_transform(yhat[:, :, i])
233 |         real = realy[:, :, i]
234 |         metrics = metric(pred, real)
235 |         log = 'Evaluate best model on test data for horizon {:d}, Test MAE: {:.4f}, Test MAPE: {:.4f}, Test RMSE: {:.4f}'
236 |         print(log.format(i + 1, metrics[0], metrics[1], metrics[2]))
237 |         mae.append(metrics[0])
238 |         mape.append(metrics[1])
239 |         rmse.append(metrics[2])
240 |     return vmae, vmape, vrmse, mae, mape, rmse
241 | 
242 | if __name__ == "__main__":
243 | 
244 |     vmae = []
245 |     vmape = []
246 |     vrmse = []
247 |     mae = []
248 |     mape = []
249 |     rmse = []
250 |     for i in range(args.runs):
251 |         vm1, vm2, vm3, m1, m2, m3 = main(i)
252 |         vmae.append(vm1)
253 |         vmape.append(vm2)
254 |         vrmse.append(vm3)
255 |         mae.append(m1)
256 |         mape.append(m2)
257 |         rmse.append(m3)
258 | 
259 |     mae = np.array(mae)
260 |     mape = np.array(mape)
261 |     rmse = np.array(rmse)
262 | 
263 |     amae = np.mean(mae,0)
264 |     amape = np.mean(mape,0)
265 |     armse = np.mean(rmse,0)
266 | 
267 |     smae = np.std(mae,0)
268 |     smape = np.std(mape,0)
269 |     srmse = np.std(rmse,0)
270 | 
271 |     print('\n\nResults for 10 runs\n\n')
272 |     #valid data
273 |     print('valid\tMAE\tRMSE\tMAPE')
274 |     log = 'mean:\t{:.4f}\t{:.4f}\t{:.4f}'
275 |     print(log.format(np.mean(vmae),np.mean(vrmse),np.mean(vmape)))
276 |     log = 'std:\t{:.4f}\t{:.4f}\t{:.4f}'
277 |     print(log.format(np.std(vmae),np.std(vrmse),np.std(vmape)))
278 |     print('\n\n')
279 |     #test data
280 |     print('test|horizon\tMAE-mean\tRMSE-mean\tMAPE-mean\tMAE-std\tRMSE-std\tMAPE-std')
281 |     for i in [2,5,11]:
282 |         log = '{:d}\t{:.4f}\t{:.4f}\t{:.4f}\t{:.4f}\t{:.4f}\t{:.4f}'
283 |         print(log.format(i+1, amae[i], armse[i], amape[i], smae[i], srmse[i], smape[i]))
284 | 
285 | 
286 | 
287 | 
288 | 
289 | 


--------------------------------------------------------------------------------
/train_single_step.py:
--------------------------------------------------------------------------------
  1 | import argparse
  2 | import math
  3 | import time
  4 | 
  5 | import torch
  6 | import torch.nn as nn
  7 | from net import gtnet
  8 | import numpy as np
  9 | import importlib
 10 | 
 11 | from util import *
 12 | from trainer import Optim
 13 | 
 14 | 
 15 | def evaluate(data, X, Y, model, evaluateL2, evaluateL1, batch_size):
 16 |     model.eval()
 17 |     total_loss = 0
 18 |     total_loss_l1 = 0
 19 |     n_samples = 0
 20 |     predict = None
 21 |     test = None
 22 | 
 23 |     for X, Y in data.get_batches(X, Y, batch_size, False):
 24 |         X = torch.unsqueeze(X,dim=1)
 25 |         X = X.transpose(2,3)
 26 |         with torch.no_grad():
 27 |             output = model(X)
 28 |         output = torch.squeeze(output)
 29 |         if len(output.shape)==1:
 30 |             output = output.unsqueeze(dim=0)
 31 |         if predict is None:
 32 |             predict = output
 33 |             test = Y
 34 |         else:
 35 |             predict = torch.cat((predict, output))
 36 |             test = torch.cat((test, Y))
 37 | 
 38 |         scale = data.scale.expand(output.size(0), data.m)
 39 |         total_loss += evaluateL2(output * scale, Y * scale).item()
 40 |         total_loss_l1 += evaluateL1(output * scale, Y * scale).item()
 41 |         n_samples += (output.size(0) * data.m)
 42 | 
 43 |     rse = math.sqrt(total_loss / n_samples) / data.rse
 44 |     rae = (total_loss_l1 / n_samples) / data.rae
 45 | 
 46 |     predict = predict.data.cpu().numpy()
 47 |     Ytest = test.data.cpu().numpy()
 48 |     sigma_p = (predict).std(axis=0)
 49 |     sigma_g = (Ytest).std(axis=0)
 50 |     mean_p = predict.mean(axis=0)
 51 |     mean_g = Ytest.mean(axis=0)
 52 |     index = (sigma_g != 0)
 53 |     correlation = ((predict - mean_p) * (Ytest - mean_g)).mean(axis=0) / (sigma_p * sigma_g)
 54 |     correlation = (correlation[index]).mean()
 55 |     return rse, rae, correlation
 56 | 
 57 | 
 58 | def train(data, X, Y, model, criterion, optim, batch_size):
 59 |     model.train()
 60 |     total_loss = 0
 61 |     n_samples = 0
 62 |     iter = 0
 63 |     for X, Y in data.get_batches(X, Y, batch_size, True):
 64 |         model.zero_grad()
 65 |         X = torch.unsqueeze(X,dim=1)
 66 |         X = X.transpose(2,3)
 67 |         if iter % args.step_size == 0:
 68 |             perm = np.random.permutation(range(args.num_nodes))
 69 |         num_sub = int(args.num_nodes / args.num_split)
 70 | 
 71 |         for j in range(args.num_split):
 72 |             if j != args.num_split - 1:
 73 |                 id = perm[j * num_sub:(j + 1) * num_sub]
 74 |             else:
 75 |                 id = perm[j * num_sub:]
 76 |             id = torch.tensor(id).to(device)
 77 |             tx = X[:, :, id, :]
 78 |             ty = Y[:, id]
 79 |             output = model(tx,id)
 80 |             output = torch.squeeze(output)
 81 |             scale = data.scale.expand(output.size(0), data.m)
 82 |             scale = scale[:,id]
 83 |             loss = criterion(output * scale, ty * scale)
 84 |             loss.backward()
 85 |             total_loss += loss.item()
 86 |             n_samples += (output.size(0) * data.m)
 87 |             grad_norm = optim.step()
 88 | 
 89 |         if iter%100==0:
 90 |             print('iter:{:3d} | loss: {:.3f}'.format(iter,loss.item()/(output.size(0) * data.m)))
 91 |         iter += 1
 92 |     return total_loss / n_samples
 93 | 
 94 | 
 95 | parser = argparse.ArgumentParser(description='PyTorch Time series forecasting')
 96 | parser.add_argument('--data', type=str, default='./data/solar_AL.txt',
 97 |                     help='location of the data file')
 98 | parser.add_argument('--log_interval', type=int, default=2000, metavar='N',
 99 |                     help='report interval')
100 | parser.add_argument('--save', type=str, default='model/model.pt',
101 |                     help='path to save the final model')
102 | parser.add_argument('--optim', type=str, default='adam')
103 | parser.add_argument('--L1Loss', type=bool, default=True)
104 | parser.add_argument('--normalize', type=int, default=2)
105 | parser.add_argument('--device',type=str,default='cuda:1',help='')
106 | parser.add_argument('--gcn_true', type=bool, default=True, help='whether to add graph convolution layer')
107 | parser.add_argument('--buildA_true', type=bool, default=True, help='whether to construct adaptive adjacency matrix')
108 | parser.add_argument('--gcn_depth',type=int,default=2,help='graph convolution depth')
109 | parser.add_argument('--num_nodes',type=int,default=137,help='number of nodes/variables')
110 | parser.add_argument('--dropout',type=float,default=0.3,help='dropout rate')
111 | parser.add_argument('--subgraph_size',type=int,default=20,help='k')
112 | parser.add_argument('--node_dim',type=int,default=40,help='dim of nodes')
113 | parser.add_argument('--dilation_exponential',type=int,default=2,help='dilation exponential')
114 | parser.add_argument('--conv_channels',type=int,default=16,help='convolution channels')
115 | parser.add_argument('--residual_channels',type=int,default=16,help='residual channels')
116 | parser.add_argument('--skip_channels',type=int,default=32,help='skip channels')
117 | parser.add_argument('--end_channels',type=int,default=64,help='end channels')
118 | parser.add_argument('--in_dim',type=int,default=1,help='inputs dimension')
119 | parser.add_argument('--seq_in_len',type=int,default=24*7,help='input sequence length')
120 | parser.add_argument('--seq_out_len',type=int,default=1,help='output sequence length')
121 | parser.add_argument('--horizon', type=int, default=3)
122 | parser.add_argument('--layers',type=int,default=5,help='number of layers')
123 | 
124 | parser.add_argument('--batch_size',type=int,default=32,help='batch size')
125 | parser.add_argument('--lr',type=float,default=0.0001,help='learning rate')
126 | parser.add_argument('--weight_decay',type=float,default=0.00001,help='weight decay rate')
127 | 
128 | parser.add_argument('--clip',type=int,default=5,help='clip')
129 | 
130 | parser.add_argument('--propalpha',type=float,default=0.05,help='prop alpha')
131 | parser.add_argument('--tanhalpha',type=float,default=3,help='tanh alpha')
132 | 
133 | parser.add_argument('--epochs',type=int,default=1,help='')
134 | parser.add_argument('--num_split',type=int,default=1,help='number of splits for graphs')
135 | parser.add_argument('--step_size',type=int,default=100,help='step_size')
136 | 
137 | 
138 | args = parser.parse_args()
139 | device = torch.device(args.device)
140 | torch.set_num_threads(3)
141 | 
142 | def main():
143 | 
144 |     Data = DataLoaderS(args.data, 0.6, 0.2, device, args.horizon, args.seq_in_len, args.normalize)
145 | 
146 |     model = gtnet(args.gcn_true, args.buildA_true, args.gcn_depth, args.num_nodes,
147 |                   device, dropout=args.dropout, subgraph_size=args.subgraph_size,
148 |                   node_dim=args.node_dim, dilation_exponential=args.dilation_exponential,
149 |                   conv_channels=args.conv_channels, residual_channels=args.residual_channels,
150 |                   skip_channels=args.skip_channels, end_channels= args.end_channels,
151 |                   seq_length=args.seq_in_len, in_dim=args.in_dim, out_dim=args.seq_out_len,
152 |                   layers=args.layers, propalpha=args.propalpha, tanhalpha=args.tanhalpha, layer_norm_affline=False)
153 |     model = model.to(device)
154 | 
155 |     print(args)
156 |     print('The recpetive field size is', model.receptive_field)
157 |     nParams = sum([p.nelement() for p in model.parameters()])
158 |     print('Number of model parameters is', nParams, flush=True)
159 | 
160 |     if args.L1Loss:
161 |         criterion = nn.L1Loss(size_average=False).to(device)
162 |     else:
163 |         criterion = nn.MSELoss(size_average=False).to(device)
164 |     evaluateL2 = nn.MSELoss(size_average=False).to(device)
165 |     evaluateL1 = nn.L1Loss(size_average=False).to(device)
166 | 
167 | 
168 |     best_val = 10000000
169 |     optim = Optim(
170 |         model.parameters(), args.optim, args.lr, args.clip, lr_decay=args.weight_decay
171 |     )
172 | 
173 |     # At any point you can hit Ctrl + C to break out of training early.
174 |     try:
175 |         print('begin training')
176 |         for epoch in range(1, args.epochs + 1):
177 |             epoch_start_time = time.time()
178 |             train_loss = train(Data, Data.train[0], Data.train[1], model, criterion, optim, args.batch_size)
179 |             val_loss, val_rae, val_corr = evaluate(Data, Data.valid[0], Data.valid[1], model, evaluateL2, evaluateL1,
180 |                                                args.batch_size)
181 |             print(
182 |                 '| end of epoch {:3d} | time: {:5.2f}s | train_loss {:5.4f} | valid rse {:5.4f} | valid rae {:5.4f} | valid corr  {:5.4f}'.format(
183 |                     epoch, (time.time() - epoch_start_time), train_loss, val_loss, val_rae, val_corr), flush=True)
184 |             # Save the model if the validation loss is the best we've seen so far.
185 | 
186 |             if val_loss < best_val:
187 |                 with open(args.save, 'wb') as f:
188 |                     torch.save(model, f)
189 |                 best_val = val_loss
190 |             if epoch % 5 == 0:
191 |                 test_acc, test_rae, test_corr = evaluate(Data, Data.test[0], Data.test[1], model, evaluateL2, evaluateL1,
192 |                                                      args.batch_size)
193 |                 print("test rse {:5.4f} | test rae {:5.4f} | test corr {:5.4f}".format(test_acc, test_rae, test_corr), flush=True)
194 | 
195 |     except KeyboardInterrupt:
196 |         print('-' * 89)
197 |         print('Exiting from training early')
198 | 
199 |     # Load the best saved model.
200 |     with open(args.save, 'rb') as f:
201 |         model = torch.load(f)
202 | 
203 |     vtest_acc, vtest_rae, vtest_corr = evaluate(Data, Data.valid[0], Data.valid[1], model, evaluateL2, evaluateL1,
204 |                                          args.batch_size)
205 |     test_acc, test_rae, test_corr = evaluate(Data, Data.test[0], Data.test[1], model, evaluateL2, evaluateL1,
206 |                                          args.batch_size)
207 |     print("final test rse {:5.4f} | test rae {:5.4f} | test corr {:5.4f}".format(test_acc, test_rae, test_corr))
208 |     return vtest_acc, vtest_rae, vtest_corr, test_acc, test_rae, test_corr
209 | 
210 | if __name__ == "__main__":
211 |     vacc = []
212 |     vrae = []
213 |     vcorr = []
214 |     acc = []
215 |     rae = []
216 |     corr = []
217 |     for i in range(10):
218 |         val_acc, val_rae, val_corr, test_acc, test_rae, test_corr = main()
219 |         vacc.append(val_acc)
220 |         vrae.append(val_rae)
221 |         vcorr.append(val_corr)
222 |         acc.append(test_acc)
223 |         rae.append(test_rae)
224 |         corr.append(test_corr)
225 |     print('\n\n')
226 |     print('10 runs average')
227 |     print('\n\n')
228 |     print("valid\trse\trae\tcorr")
229 |     print("mean\t{:5.4f}\t{:5.4f}\t{:5.4f}".format(np.mean(vacc), np.mean(vrae), np.mean(vcorr)))
230 |     print("std\t{:5.4f}\t{:5.4f}\t{:5.4f}".format(np.std(vacc), np.std(vrae), np.std(vcorr)))
231 |     print('\n\n')
232 |     print("test\trse\trae\tcorr")
233 |     print("mean\t{:5.4f}\t{:5.4f}\t{:5.4f}".format(np.mean(acc), np.mean(rae), np.mean(corr)))
234 |     print("std\t{:5.4f}\t{:5.4f}\t{:5.4f}".format(np.std(acc), np.std(rae), np.std(corr)))
235 | 
236 | 


--------------------------------------------------------------------------------
/trainer.py:
--------------------------------------------------------------------------------
  1 | import torch.optim as optim
  2 | import math
  3 | from net import *
  4 | import util
  5 | class Trainer():
  6 |     def __init__(self, model, lrate, wdecay, clip, step_size, seq_out_len, scaler, device, cl=True):
  7 |         self.scaler = scaler
  8 |         self.model = model
  9 |         self.model.to(device)
 10 |         self.optimizer = optim.Adam(self.model.parameters(), lr=lrate, weight_decay=wdecay)
 11 |         self.loss = util.masked_mae
 12 |         self.clip = clip
 13 |         self.step = step_size
 14 |         self.iter = 1
 15 |         self.task_level = 1
 16 |         self.seq_out_len = seq_out_len
 17 |         self.cl = cl
 18 | 
 19 |     def train(self, input, real_val, idx=None):
 20 |         self.model.train()
 21 |         self.optimizer.zero_grad()
 22 |         output = self.model(input, idx=idx)
 23 |         output = output.transpose(1,3)
 24 |         real = torch.unsqueeze(real_val,dim=1)
 25 |         predict = self.scaler.inverse_transform(output)
 26 |         if self.iter%self.step==0 and self.task_level<=self.seq_out_len:
 27 |             self.task_level +=1
 28 |         if self.cl:
 29 |             loss = self.loss(predict[:, :, :, :self.task_level], real[:, :, :, :self.task_level], 0.0)
 30 |         else:
 31 |             loss = self.loss(predict, real, 0.0)
 32 | 
 33 |         loss.backward()
 34 | 
 35 |         if self.clip is not None:
 36 |             torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip)
 37 | 
 38 |         self.optimizer.step()
 39 |         # mae = util.masked_mae(predict,real,0.0).item()
 40 |         mape = util.masked_mape(predict,real,0.0).item()
 41 |         rmse = util.masked_rmse(predict,real,0.0).item()
 42 |         self.iter += 1
 43 |         return loss.item(),mape,rmse
 44 | 
 45 |     def eval(self, input, real_val):
 46 |         self.model.eval()
 47 |         output = self.model(input)
 48 |         output = output.transpose(1,3)
 49 |         real = torch.unsqueeze(real_val,dim=1)
 50 |         predict = self.scaler.inverse_transform(output)
 51 |         loss = self.loss(predict, real, 0.0)
 52 |         mape = util.masked_mape(predict,real,0.0).item()
 53 |         rmse = util.masked_rmse(predict,real,0.0).item()
 54 |         return loss.item(),mape,rmse
 55 | 
 56 | 
 57 | 
 58 | class Optim(object):
 59 | 
 60 |     def _makeOptimizer(self):
 61 |         if self.method == 'sgd':
 62 |             self.optimizer = optim.SGD(self.params, lr=self.lr, weight_decay=self.lr_decay)
 63 |         elif self.method == 'adagrad':
 64 |             self.optimizer = optim.Adagrad(self.params, lr=self.lr, weight_decay=self.lr_decay)
 65 |         elif self.method == 'adadelta':
 66 |             self.optimizer = optim.Adadelta(self.params, lr=self.lr, weight_decay=self.lr_decay)
 67 |         elif self.method == 'adam':
 68 |             self.optimizer = optim.Adam(self.params, lr=self.lr, weight_decay=self.lr_decay)
 69 |         else:
 70 |             raise RuntimeError("Invalid optim method: " + self.method)
 71 | 
 72 |     def __init__(self, params, method, lr, clip, lr_decay=1, start_decay_at=None):
 73 |         self.params = params  # careful: params may be a generator
 74 |         self.last_ppl = None
 75 |         self.lr = lr
 76 |         self.clip = clip
 77 |         self.method = method
 78 |         self.lr_decay = lr_decay
 79 |         self.start_decay_at = start_decay_at
 80 |         self.start_decay = False
 81 | 
 82 |         self._makeOptimizer()
 83 | 
 84 |     def step(self):
 85 |         # Compute gradients norm.
 86 |         grad_norm = 0
 87 |         if self.clip is not None:
 88 |             torch.nn.utils.clip_grad_norm_(self.params, self.clip)
 89 | 
 90 |         # for param in self.params:
 91 |         #     grad_norm += math.pow(param.grad.data.norm(), 2)
 92 |         #
 93 |         # grad_norm = math.sqrt(grad_norm)
 94 |         # if grad_norm > 0:
 95 |         #     shrinkage = self.max_grad_norm / grad_norm
 96 |         # else:
 97 |         #     shrinkage = 1.
 98 |         #
 99 |         # for param in self.params:
100 |         #     if shrinkage < 1:
101 |         #         param.grad.data.mul_(shrinkage)
102 |         self.optimizer.step()
103 |         return  grad_norm
104 | 
105 |     # decay learning rate if val perf does not improve or we hit the start_decay_at limit
106 |     def updateLearningRate(self, ppl, epoch):
107 |         if self.start_decay_at is not None and epoch >= self.start_decay_at:
108 |             self.start_decay = True
109 |         if self.last_ppl is not None and ppl > self.last_ppl:
110 |             self.start_decay = True
111 | 
112 |         if self.start_decay:
113 |             self.lr = self.lr * self.lr_decay
114 |             print("Decaying learning rate to %g" % self.lr)
115 |         #only decay for one epoch
116 |         self.start_decay = False
117 | 
118 |         self.last_ppl = ppl
119 | 
120 |         self._makeOptimizer()
121 | 


--------------------------------------------------------------------------------
/util.py:
--------------------------------------------------------------------------------
  1 | import pickle
  2 | import numpy as np
  3 | import os
  4 | import scipy.sparse as sp
  5 | import torch
  6 | from scipy.sparse import linalg
  7 | from torch.autograd import Variable
  8 | 
  9 | def normal_std(x):
 10 |     return x.std() * np.sqrt((len(x) - 1.)/(len(x)))
 11 | 
 12 | class DataLoaderS(object):
 13 |     # train and valid is the ratio of training set and validation set. test = 1 - train - valid
 14 |     def __init__(self, file_name, train, valid, device, horizon, window, normalize=2):
 15 |         self.P = window
 16 |         self.h = horizon
 17 |         fin = open(file_name)
 18 |         self.rawdat = np.loadtxt(fin, delimiter=',')
 19 |         self.dat = np.zeros(self.rawdat.shape)
 20 |         self.n, self.m = self.dat.shape
 21 |         self.normalize = 2
 22 |         self.scale = np.ones(self.m)
 23 |         self._normalized(normalize)
 24 |         self._split(int(train * self.n), int((train + valid) * self.n), self.n)
 25 | 
 26 |         self.scale = torch.from_numpy(self.scale).float()
 27 |         tmp = self.test[1] * self.scale.expand(self.test[1].size(0), self.m)
 28 | 
 29 |         self.scale = self.scale.to(device)
 30 |         self.scale = Variable(self.scale)
 31 | 
 32 |         self.rse = normal_std(tmp)
 33 |         self.rae = torch.mean(torch.abs(tmp - torch.mean(tmp)))
 34 | 
 35 |         self.device = device
 36 | 
 37 |     def _normalized(self, normalize):
 38 |         # normalized by the maximum value of entire matrix.
 39 | 
 40 |         if (normalize == 0):
 41 |             self.dat = self.rawdat
 42 | 
 43 |         if (normalize == 1):
 44 |             self.dat = self.rawdat / np.max(self.rawdat)
 45 | 
 46 |         # normlized by the maximum value of each row(sensor).
 47 |         if (normalize == 2):
 48 |             for i in range(self.m):
 49 |                 self.scale[i] = np.max(np.abs(self.rawdat[:, i]))
 50 |                 self.dat[:, i] = self.rawdat[:, i] / np.max(np.abs(self.rawdat[:, i]))
 51 | 
 52 |     def _split(self, train, valid, test):
 53 | 
 54 |         train_set = range(self.P + self.h - 1, train)
 55 |         valid_set = range(train, valid)
 56 |         test_set = range(valid, self.n)
 57 |         self.train = self._batchify(train_set, self.h)
 58 |         self.valid = self._batchify(valid_set, self.h)
 59 |         self.test = self._batchify(test_set, self.h)
 60 | 
 61 |     def _batchify(self, idx_set, horizon):
 62 |         n = len(idx_set)
 63 |         X = torch.zeros((n, self.P, self.m))
 64 |         Y = torch.zeros((n, self.m))
 65 |         for i in range(n):
 66 |             end = idx_set[i] - self.h + 1
 67 |             start = end - self.P
 68 |             X[i, :, :] = torch.from_numpy(self.dat[start:end, :])
 69 |             Y[i, :] = torch.from_numpy(self.dat[idx_set[i], :])
 70 |         return [X, Y]
 71 | 
 72 |     def get_batches(self, inputs, targets, batch_size, shuffle=True):
 73 |         length = len(inputs)
 74 |         if shuffle:
 75 |             index = torch.randperm(length)
 76 |         else:
 77 |             index = torch.LongTensor(range(length))
 78 |         start_idx = 0
 79 |         while (start_idx < length):
 80 |             end_idx = min(length, start_idx + batch_size)
 81 |             excerpt = index[start_idx:end_idx]
 82 |             X = inputs[excerpt]
 83 |             Y = targets[excerpt]
 84 |             X = X.to(self.device)
 85 |             Y = Y.to(self.device)
 86 |             yield Variable(X), Variable(Y)
 87 |             start_idx += batch_size
 88 | 
 89 | class DataLoaderM(object):
 90 |     def __init__(self, xs, ys, batch_size, pad_with_last_sample=True):
 91 |         """
 92 |         :param xs:
 93 |         :param ys:
 94 |         :param batch_size:
 95 |         :param pad_with_last_sample: pad with the last sample to make number of samples divisible to batch_size.
 96 |         """
 97 |         self.batch_size = batch_size
 98 |         self.current_ind = 0
 99 |         if pad_with_last_sample:
100 |             num_padding = (batch_size - (len(xs) % batch_size)) % batch_size
101 |             x_padding = np.repeat(xs[-1:], num_padding, axis=0)
102 |             y_padding = np.repeat(ys[-1:], num_padding, axis=0)
103 |             xs = np.concatenate([xs, x_padding], axis=0)
104 |             ys = np.concatenate([ys, y_padding], axis=0)
105 |         self.size = len(xs)
106 |         self.num_batch = int(self.size // self.batch_size)
107 |         self.xs = xs
108 |         self.ys = ys
109 | 
110 |     def shuffle(self):
111 |         permutation = np.random.permutation(self.size)
112 |         xs, ys = self.xs[permutation], self.ys[permutation]
113 |         self.xs = xs
114 |         self.ys = ys
115 | 
116 |     def get_iterator(self):
117 |         self.current_ind = 0
118 |         def _wrapper():
119 |             while self.current_ind < self.num_batch:
120 |                 start_ind = self.batch_size * self.current_ind
121 |                 end_ind = min(self.size, self.batch_size * (self.current_ind + 1))
122 |                 x_i = self.xs[start_ind: end_ind, ...]
123 |                 y_i = self.ys[start_ind: end_ind, ...]
124 |                 yield (x_i, y_i)
125 |                 self.current_ind += 1
126 | 
127 |         return _wrapper()
128 | 
129 | class StandardScaler():
130 |     """
131 |     Standard the input
132 |     """
133 |     def __init__(self, mean, std):
134 |         self.mean = mean
135 |         self.std = std
136 |     def transform(self, data):
137 |         return (data - self.mean) / self.std
138 |     def inverse_transform(self, data):
139 |         return (data * self.std) + self.mean
140 | 
141 | 
142 | def sym_adj(adj):
143 |     """Symmetrically normalize adjacency matrix."""
144 |     adj = sp.coo_matrix(adj)
145 |     rowsum = np.array(adj.sum(1))
146 |     d_inv_sqrt = np.power(rowsum, -0.5).flatten()
147 |     d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.
148 |     d_mat_inv_sqrt = sp.diags(d_inv_sqrt)
149 |     return adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt).astype(np.float32).todense()
150 | 
151 | def asym_adj(adj):
152 |     """Asymmetrically normalize adjacency matrix."""
153 |     adj = sp.coo_matrix(adj)
154 |     rowsum = np.array(adj.sum(1)).flatten()
155 |     d_inv = np.power(rowsum, -1).flatten()
156 |     d_inv[np.isinf(d_inv)] = 0.
157 |     d_mat= sp.diags(d_inv)
158 |     return d_mat.dot(adj).astype(np.float32).todense()
159 | 
160 | def calculate_normalized_laplacian(adj):
161 |     """
162 |     # L = D^-1/2 (D-A) D^-1/2 = I - D^-1/2 A D^-1/2
163 |     # D = diag(A 1)
164 |     :param adj:
165 |     :return:
166 |     """
167 |     adj = sp.coo_matrix(adj)
168 |     d = np.array(adj.sum(1))
169 |     d_inv_sqrt = np.power(d, -0.5).flatten()
170 |     d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.
171 |     d_mat_inv_sqrt = sp.diags(d_inv_sqrt)
172 |     normalized_laplacian = sp.eye(adj.shape[0]) - adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt).tocoo()
173 |     return normalized_laplacian
174 | 
175 | def calculate_scaled_laplacian(adj_mx, lambda_max=2, undirected=True):
176 |     if undirected:
177 |         adj_mx = np.maximum.reduce([adj_mx, adj_mx.T])
178 |     L = calculate_normalized_laplacian(adj_mx)
179 |     if lambda_max is None:
180 |         lambda_max, _ = linalg.eigsh(L, 1, which='LM')
181 |         lambda_max = lambda_max[0]
182 |     L = sp.csr_matrix(L)
183 |     M, _ = L.shape
184 |     I = sp.identity(M, format='csr', dtype=L.dtype)
185 |     L = (2 / lambda_max * L) - I
186 |     return L.astype(np.float32).todense()
187 | 
188 | 
189 | def load_pickle(pickle_file):
190 |     try:
191 |         with open(pickle_file, 'rb') as f:
192 |             pickle_data = pickle.load(f)
193 |     except UnicodeDecodeError as e:
194 |         with open(pickle_file, 'rb') as f:
195 |             pickle_data = pickle.load(f, encoding='latin1')
196 |     except Exception as e:
197 |         print('Unable to load data ', pickle_file, ':', e)
198 |         raise
199 |     return pickle_data
200 | 
201 | def load_adj(pkl_filename):
202 |     sensor_ids, sensor_id_to_ind, adj = load_pickle(pkl_filename)
203 |     return adj
204 | 
205 | 
206 | def load_dataset(dataset_dir, batch_size, valid_batch_size= None, test_batch_size=None):
207 |     data = {}
208 |     for category in ['train', 'val', 'test']:
209 |         cat_data = np.load(os.path.join(dataset_dir, category + '.npz'))
210 |         data['x_' + category] = cat_data['x']
211 |         data['y_' + category] = cat_data['y']
212 |     scaler = StandardScaler(mean=data['x_train'][..., 0].mean(), std=data['x_train'][..., 0].std())
213 |     # Data format
214 |     for category in ['train', 'val', 'test']:
215 |         data['x_' + category][..., 0] = scaler.transform(data['x_' + category][..., 0])
216 | 
217 |     data['train_loader'] = DataLoaderM(data['x_train'], data['y_train'], batch_size)
218 |     data['val_loader'] = DataLoaderM(data['x_val'], data['y_val'], valid_batch_size)
219 |     data['test_loader'] = DataLoaderM(data['x_test'], data['y_test'], test_batch_size)
220 |     data['scaler'] = scaler
221 |     return data
222 | 
223 | 
224 | 
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 | def masked_rmse(preds, labels, null_val=np.nan):
239 |     return torch.sqrt(masked_mse(preds=preds, labels=labels, null_val=null_val))
240 | 
241 | 
242 | def masked_mae(preds, labels, null_val=np.nan):
243 |     if np.isnan(null_val):
244 |         mask = ~torch.isnan(labels)
245 |     else:
246 |         mask = (labels!=null_val)
247 |     mask = mask.float()
248 |     mask /=  torch.mean((mask))
249 |     mask = torch.where(torch.isnan(mask), torch.zeros_like(mask), mask)
250 |     loss = torch.abs(preds-labels)
251 |     loss = loss * mask
252 |     loss = torch.where(torch.isnan(loss), torch.zeros_like(loss), loss)
253 |     return torch.mean(loss)
254 | 
255 | def masked_mape(preds, labels, null_val=np.nan):
256 |     if np.isnan(null_val):
257 |         mask = ~torch.isnan(labels)
258 |     else:
259 |         mask = (labels!=null_val)
260 |     mask = mask.float()
261 |     mask /=  torch.mean((mask))
262 |     mask = torch.where(torch.isnan(mask), torch.zeros_like(mask), mask)
263 |     loss = torch.abs(preds-labels)/labels
264 |     loss = loss * mask
265 |     loss = torch.where(torch.isnan(loss), torch.zeros_like(loss), loss)
266 |     return torch.mean(loss)
267 | 
268 | 
269 | def metric(pred, real):
270 |     mae = masked_mae(pred,real,0.0).item()
271 |     mape = masked_mape(pred,real,0.0).item()
272 |     rmse = masked_rmse(pred,real,0.0).item()
273 |     return mae,mape,rmse
274 | 
275 | 
276 | def load_node_feature(path):
277 |     fi = open(path)
278 |     x = []
279 |     for li in fi:
280 |         li = li.strip()
281 |         li = li.split(",")
282 |         e = [float(t) for t in li[1:]]
283 |         x.append(e)
284 |     x = np.array(x)
285 |     mean = np.mean(x,axis=0)
286 |     std = np.std(x,axis=0)
287 |     z = torch.tensor((x-mean)/std,dtype=torch.float)
288 |     return z
289 | 
290 | 
291 | def normal_std(x):
292 |     return x.std() * np.sqrt((len(x) - 1.) / (len(x)))
293 | 
294 | 
295 | 
296 |             


--------------------------------------------------------------------------------