Graph-WaveNet-WaveBound/engine.py at master · carrtesy/Graph-WaveNet-WaveBound · GitHub

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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import torch.optim as optim
from model import *
import util
from utils.ema import EMAUpdater
import numpy as np
import torch
import torch.nn as nn

class trainer():
    def __init__(self, scaler, in_dim, seq_length, num_nodes,
                 nhid , dropout, lrate, wdecay, device, supports,
                 gcn_bool, addaptadj, aptinit, args):

        self.args = args
        self.iter_count = 0
        self.source_model = gwnet(device, num_nodes, dropout,
                                  supports=supports, gcn_bool=gcn_bool, addaptadj=addaptadj,
                                  aptinit=aptinit, in_dim=in_dim, out_dim=seq_length, residual_channels=nhid,
                                  dilation_channels=nhid, skip_channels=nhid * 8, end_channels=nhid * 16)
        self.source_model.to(device)

        # ema
        self.use_ema = self.args.use_ema
        if self.use_ema:
            self.target_model = gwnet(device, num_nodes, dropout,
                                      supports=supports, gcn_bool=gcn_bool, addaptadj=addaptadj,
                                      aptinit=aptinit, in_dim=in_dim, out_dim=seq_length, residual_channels=nhid,
                                      dilation_channels=nhid, skip_channels=nhid * 8, end_channels=nhid * 16)
            self.target_model.to(device)
            self.ema_updater = EMAUpdater(self.target_model, self.source_model,
                                          self.args.moving_average_decay,
                                          self.args.start_iter)
            self.eval_model = self.target_model
            self.eval_model.to(device)
        else:
            self.eval_model = self.source_model
            self.eval_model.to(device)

        self.optimizer = optim.Adam(self.source_model.parameters(), lr=lrate, weight_decay=wdecay)
        self.loss = util.masked_mae
        self.scaler = scaler
        self.clip = 5


    def train(self, input, real_val):
        if self.args.use_ema:
            self.target_model.train()
        self.source_model.train()

        input = nn.functional.pad(input,(1,0,0,0))

        # get source, target output
        real = torch.unsqueeze(real_val, dim=1)
        source_output = self.source_model(input)
        source_output = source_output.transpose(1, 3)
        source_pred = self.scaler.inverse_transform(source_output)
        if self.use_ema and self.iter_count >= self.args.start_iter:
            target_output = self.target_model(input)
            target_output = target_output.transpose(1, 3)
            target_pred = self.scaler.inverse_transform(target_output)
            loss = self.compute_loss(source_pred, target_pred, real, null_val=0.0)
        else:
            loss = self.loss(source_pred, real, 0.0)

        self.optimizer.zero_grad()
        # compute loss
        loss.backward()

        # optim
        if self.clip is not None:
            torch.nn.utils.clip_grad_norm_(self.source_model.parameters(), self.clip)
            #if self.use_ema:
            #    torch.nn.utils.clip_grad_norm_(self.target_model.parameters(), self.clip)
        self.optimizer.step()

        # ema update
        if self.use_ema:
            self.ema_updater.update(iter_count=self.iter_count)

        # training log
        mape = util.masked_mape(source_pred, real, 0.0).item()
        rmse = util.masked_rmse(source_pred, real, 0.0).item()
        return loss.item(), mape, rmse

    def compute_loss(self, source_output, target_output, y, null_val=np.nan):
        '''
        if np.isnan(null_val):
            mask = ~torch.isnan(y)
        else:
            mask = (y != null_val)
        mask = mask.float()
        mask /= torch.mean((mask))
        mask = torch.where(torch.isnan(mask), torch.zeros_like(mask), mask)
        '''

        source_loss = self.loss(source_output, y, null_val=null_val, reduce=False).mean(0)
        target_loss = self.loss(target_output, y, null_val=null_val, reduce=False).mean(0)
        loss = torch.abs(source_loss - target_loss + self.args.epsilon).mean()
        return loss

    @torch.no_grad()
    def eval(self, input, real_val):
        self.eval_model.eval()
        input = nn.functional.pad(input,(1,0,0,0))
        output = self.eval_model(input)
        output = output.transpose(1,3)
        #output = [batch_size,12,num_nodes,1]
        real = torch.unsqueeze(real_val,dim=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