|
- import torch
- import numpy as np
- import argparse
- import time
- import util
- import matplotlib.pyplot as plt
- from engine import trainer
- import random
-
-
-
- def setup_seed(seed):
- torch.manual_seed(seed)
- torch.cuda.manual_seed_all(seed)
- np.random.seed(seed)
- random.seed(seed)
- torch.backends.cudnn.deterministic = True
-
-
- # 设置随机数种子
- setup_seed(20)
-
-
- def main(data,adjdata,seq_length,nhid,in_dim,num_nodes,batch_size,learning_rate,dropout,weight_decay,epochs,print_every,save,SE_file,test_output_file):
- # load data
- sensor_ids, sensor_id_to_ind, adj_mx = util.load_adj(adjdata) # adj_max=[[forward_adj_max(nroute,nroute)],[back_adj_matrix(nroute,nroute)]]
- dataloader = util.load_dataset(data, batch_size, batch_size, batch_size)
- scaler = dataloader['scaler'] # scaler.mean scaler.std
-
-
- graphSE = util.load_SE(SE_file)
-
-
- gpu_count = torch.cuda.device_count()
- supports = adj_mx
- adj_mx = [np.tile(i, (gpu_count, 1)) for i in adj_mx]
-
-
-
- adjinit = supports[0]
-
-
-
- engine = trainer(scaler, in_dim, seq_length, num_nodes, nhid, dropout,
- learning_rate, weight_decay, supports,
- adjinit, graphSE=graphSE)
-
- print("start training...", flush=True)
- ealystop_minval = float("inf")
- patient_count = 0
-
- his_loss = []
- val_time = []
- train_time = []
- lr = learning_rate
- for i in range(1, epochs + 1):
- if patient_count > 10:
- break
- train_loss = []
- train_mape = []
- train_rmse = []
- t1 = time.time()
- dataloader['train_loader'].shuffle()
- if i % 10 == 0:
- lr = lr / 2
- print("learning rate:", lr)
- for iter, (x, y) in enumerate(dataloader['train_loader'].get_iterator()):
- trainx = torch.Tensor(x).cuda() # [batchsize,nhis,nroute,dim]
- trainx = trainx.transpose(1, 3) # [batchsize,dim,nroute,nhis] dim-0 speed dim-1 dayofweek
- trainy = torch.Tensor(y).cuda()
- trainy = trainy.transpose(1, 3)
-
- supports = [torch.tensor(i).cuda()for i in adj_mx]
- metrics = engine.train(trainx, trainy, True, lr=lr, current_epoch=i, all_epoch=epochs,
- supports=supports)
- train_loss.append(metrics[1])
- train_mape.append(metrics[4])
- train_rmse.append(metrics[5])
- if iter % print_every == 0:
- log = 'epoch:{:03d},Iter: {:03d}, Train Loss: {:.4f},{:.4f},{:.4f},{:.4f}, Train MAPE: {:.4f}, Train RMSE: {:.4f},delta: {:.4f}'
- print(log.format(i, iter, metrics[0], metrics[1], metrics[2], metrics[3], metrics[4], metrics[5],
- metrics[6]), flush=True)
- t2 = time.time()
- train_time.append(t2 - t1)
- # validation
- valid_loss = []
- valid_mape = []
- valid_rmse = []
-
- s1 = time.time()
- for iter, (x, y) in enumerate(dataloader['val_loader'].get_iterator()):
- testx = torch.Tensor(x).cuda()
- testx = testx.transpose(1, 3)
- testy = torch.Tensor(y).cuda()
- testy = testy.transpose(1, 3)
- supports = [torch.tensor(i).cuda()for i in adj_mx]
-
- metrics = engine.eval(testx, testy, supports)
- valid_loss.append(metrics[0])
- valid_mape.append(metrics[1])
- valid_rmse.append(metrics[2])
- s2 = time.time()
- log = 'Epoch: {:03d}, Inference Time: {:.4f} secs'
- print(log.format(i, (s2 - s1)))
- val_time.append(s2 - s1)
- mtrain_loss = np.mean(train_loss)
- mtrain_mape = np.mean(train_mape)
- mtrain_rmse = np.mean(train_rmse)
-
- mvalid_loss = np.mean(valid_loss)
- mvalid_mape = np.mean(valid_mape)
- mvalid_rmse = np.mean(valid_rmse)
- his_loss.append(mvalid_loss)
-
- 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'
- print(log.format(i, mtrain_loss, mtrain_mape, mtrain_rmse, mvalid_loss, mvalid_mape, mvalid_rmse, (t2 - t1)),
- flush=True)
- torch.save(engine.model.state_dict(),
- save + "epoch_" + str(i) + "_" + str(round(mvalid_loss, 4)) + ".pth")
- if mvalid_loss < ealystop_minval:
- ealystop_minval = mvalid_loss
- patient_count = 0
- else:
- patient_count += 1
- print("patience count:", patient_count)
-
- print("Average Training Time: {:.4f} secs/epoch".format(np.mean(train_time)))
- print("Average Inference Time: {:.4f} secs".format(np.mean(val_time)))
-
-
- # testing
- bestid = np.argmin(his_loss)
- print("Training finished")
- print("The valid loss on best model is", str(round(his_loss[bestid], 4)))
-
- engine.model.load_state_dict(
- torch.load(save + "epoch_" + str(bestid + 1) + "_" + str(round(his_loss[bestid], 4)) + ".pth"))
- print("load model..", save + "epoch_" + str(bestid + 1) + "_" + str(round(his_loss[bestid], 4)) + ".pth")
-
- outputs = []
- realy = torch.Tensor(dataloader['y_test']).cuda()
- realy = realy.transpose(1, 3)[:, 0, :, :]
-
- for iter, (x, y) in enumerate(dataloader['test_loader'].get_iterator()):
- testx = torch.Tensor(x).cuda()
- testx = testx.transpose(1, 3)
- supports = [torch.tensor(i).cuda()for i in adj_mx]
- with torch.no_grad():
- preds = engine.model(testx, realy, False, supports=supports).transpose(1, 3)
- outputs.append(preds.squeeze())
-
- yhat = torch.cat(outputs, dim=0)
- f=open(test_output_file,'wb')
- np.savez(f,yhat.cpu().detach().numpy())
- yhat = yhat[:realy.size(0), ...]
-
-
- amae = []
- amape = []
- armse = []
- for i in range(12):
- pred = scaler.inverse_transform(yhat[:, :, i])
- real = realy[:, :, i]
- metrics = util.metric(pred, real)
- log = 'Evaluate best model on test data for horizon {:d}, Test MAE: {:.4f}, Test MAPE: {:.4f}, Test RMSE: {:.4f}'
- print(log.format(i + 1, metrics[0], metrics[1], metrics[2]))
- amae.append(metrics[0])
- amape.append(metrics[1])
- armse.append(metrics[2])
-
- log = 'On average over 12 horizons, Test MAE: {:.4f}, Test MAPE: {:.5f}, Test RMSE: {:.4f}'
- print(log.format(np.mean(amae), np.mean(amape), np.mean(armse)))
- torch.save(engine.model.state_dict(),
- save + "_exp_best_" + str(round(his_loss[bestid], 2)) + ".pth")
-
-
- if __name__ == "__main__":
- t1 = time.time()
- main()
- t2 = time.time()
- print("Total time spent: {:.4f}".format(t2 - t1))
|