|
- import torch.optim as optim
- from model_dagstn import model_dagstn
- import torch
- import torch.nn as nn
- import torch.nn.functional as F
- from torch.autograd import Variable
- from syn_batchnormal.func import convert_model
- import sys
- import util
- class trainer():
- def __init__(self, scaler, in_dim, seq_length, num_nodes, nhid , dropout, lrate, wdecay, supports, aptinit,training=True,graphSE=None):
- self.model = model_dagstn(scaler, in_dim, seq_length, num_nodes, nhid , dropout, lrate, wdecay, supports, aptinit,training,graphSE=graphSE)
- if torch.cuda.device_count() > 0:
- print("Let's use", torch.cuda.device_count(), "GPUs!")
- self.model = torch.nn.DataParallel(self.model) # device_ids will include all GPU devices by default
- self.model = self.model.cuda()
- self.model = convert_model(self.model)
- self.lr=lrate
- self.optimizer = optim.Adam(self.model.parameters(), lr=lrate, weight_decay=wdecay)
- self.loss = util.masked_mae
- self.scaler = scaler
- self.clip = 5
- self.istraining=training
-
- def train(self, input, real_val,istraining=True,lr=0.001,current_epoch=None,all_epoch=None,supports=None):# input:[batchsize,2,nroute,12] real_val:[batchsize,2,nroute,nhis]
-
- for param_group in self.optimizer.param_groups:
- param_group['lr']=lr
-
- if istraining:
- self.model.train()
- self.optimizer.zero_grad()
- forward_output=self.model(input,real_val,True,current_epoch,all_epoch,supports)
-
- forward_output = forward_output.transpose(1,3) #output:[batchsize,1,nroute,12]
- real = real_val[:,0:1,:,:] #real:[batchsize,1,nroute,nhist]
- predict = self.scaler.inverse_transform(forward_output)#predict:[batchsize,1,nroute,nhist]
- loss = self.loss(predict, real, 0.0)
-
-
- loss1=loss2=loss3=loss
- loss.backward()
- if self.clip is not None:
- torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip)
-
- self.optimizer.step()
- mape = util.masked_mape(predict,real,0.0).item()
- rmse = util.masked_rmse(predict,real,0.0).item()
- return loss.item(),loss1.item(),loss2.item(),loss3.item(),mape,rmse,None is None
- else:
- pass
-
- def eval(self, input, real_val,supports):
- self.model.eval()
- output = self.model(input,real_val,istraining=False,supports=supports)
- output = output.transpose(1,3)
- #output = [batch_size,12,num_nodes,1]
- real = real_val[:, 0:1, :, :]
- predict = self.scaler.inverse_transform(output)
- loss = self.loss(predict, real, 0.0)
- mape = util.masked_mape(predict,real,0.0).item()
- rmse = util.masked_rmse(predict,real,0.0).item()
- return loss.item(),mape,rmse
|