DSSTGNN/engine.py at main · Sephuroth/DSSTGNN · 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
# -*- coding:utf-8 -*-
"""
Filename: engine.py
Author: lisheng
Time: 2025-04-06
"""

import util
from models.DSSTGNN.DSSTGNN import DSSTGNN
from optimizer import Ranger
from util import *

class trainer:
    def __init__(
        self,
        scaler,
        batch_size,
        input_dim,
        channels,
        num_nodes,
        input_len,
        output_len,
        dropout,
        lrate,
        wdecay,
        device,
    ):
        self.model = DSSTGNN(device, batch_size, input_dim, channels, num_nodes, input_len, output_len, dropout)
        self.model.to(device)
        self.optimizer = Ranger(self.model.parameters(), lr=lrate, weight_decay=wdecay)
        self.loss = util.MAE_torch   # MAE_Freq_torch
        self.scaler = scaler
        self.clip = 5
        print(self.model)

    def train(self, input, real_val):

        self.model.train()
        self.optimizer.zero_grad()
        output, A_graph, D_graph = self.model(input)
        real = real_val
        predict = self.scaler.inverse_transform(output)
        loss = self.loss(predict, real, 0.0)
        loss.backward()
        if self.clip is not None:
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip)
        self.optimizer.step()
        mape = util.MAPE_torch(predict, real, 0.0).item()
        rmse = util.RMSE_torch(predict, real, 0.0).item()
        wmape = util.WMAPE_torch(predict, real, 0.0).item()
        return loss.item(), mape, rmse, wmape, A_graph.detach(), D_graph.detach()

    def eval(self, input, real_val):
        self.model.eval()
        output, A_graph, D_graph = self.model(input)
        real = real_val
        predict = self.scaler.inverse_transform(output)
        loss = self.loss(predict, real, 0.0)
        mape = util.MAPE_torch(predict, real, 0.0).item()
        rmse = util.RMSE_torch(predict, real, 0.0).item()
        wmape = util.WMAPE_torch(predict, real, 0.0).item()
        return loss.item(), mape, rmse, wmape, A_graph.detach(), D_graph.detach()