-
Notifications
You must be signed in to change notification settings - Fork 2
/
test.py
89 lines (75 loc) · 2.49 KB
/
test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import torch
from torch.utils.data import DataLoader
from models.stgsp import STGSP
from data.dataset import DatasetFactory
import numpy as np
seed = 777
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
class DataConfiguration:
# Data
name = 'TaxiBJ'
portion = 1. # portion of data
len_close = 3
len_period = 1
len_trend = 1
pad_forward_period = 0
pad_back_period = 0
pad_forward_trend = 0
pad_back_trend = 0
len_all_close = len_close * 1
len_all_period = len_period * (1 + pad_back_period + pad_forward_period)
len_all_trend = len_trend * (1 + pad_back_trend + pad_forward_trend)
len_seq = len_all_close + len_all_period + len_all_trend
cpt = [len_all_close, len_all_period, len_all_trend]
interval_period = 1
interval_trend = 7
ext_flag = True
timeenc_flag = 'w' # 'm', 'w', 'd'
rm_incomplete_flag = True
fourty_eight = True
previous_meteorol = True
ext_dim = 77 # 77
dim_flow = 2
dim_h = 32
dim_w = 32
def set_seed(seed):
torch.manual_seed(seed) # cpu
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
def train():
set_seed(seed)
dconf = DataConfiguration()
ds_factory = DatasetFactory(dconf)
select_pre = 0
test_ds = ds_factory.get_test_dataset(select_pre)
test_loader = DataLoader(
dataset=test_ds,
batch_size=32,
shuffle=False
)
model = STGSP(dconf)
model.load("checkpoints/TaxiBJ/model_finetune.pth")
model = model.to(device)
model.eval()
trues = []
preds = []
with torch.no_grad():
for _, (X, X_ext, Y, Y_ext) in enumerate(test_loader):
X = X.to(device)
X_ext = X_ext.to(device)
Y = Y.to(device)
Y_ext = Y_ext.to(device)
outputs = model(X, X_ext, Y_ext)
true = ds_factory.ds.mmn.inverse_transform(Y.detach().cpu().numpy())
pred = ds_factory.ds.mmn.inverse_transform(outputs.detach().cpu().numpy())
trues.append(true)
preds.append(pred)
trues = np.concatenate(trues, 0)
preds = np.concatenate(preds, 0)
mae = np.mean(np.abs(preds-trues))
rmse = np.sqrt(np.mean((preds-trues)**2))
print("test_rmse: %.4f, test_mae: %.4f" % (rmse, mae))
if __name__ == '__main__':
print("l_c:", DataConfiguration.len_close, "l_p:", DataConfiguration.len_period, "l_t:", DataConfiguration.len_trend)
train()