DDSC/train.py at main · NUIST-GenAI-Lab/DDSC · 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
import math
import os

import numpy as np
from torch.utils.data import DataLoader
from torch.nn import DataParallel
from tqdm import tqdm
import torch


def train_one_epoch(net, train_loader, device, criterion, optimizer):
    training_loss = 0.
    for data in train_loader:
        x = data[0]
        y = data[1]

        x = x.data.to(device)
        y = y.data.to(device)

        y_pred = net(x)[0]
        loss = criterion(y_pred, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        training_loss += loss.item()

    return training_loss / len(train_loader)


def validate_one_epoch(net, test_data, device, criterion):
    validating_loss = 0.
    with torch.no_grad():
        x_test = test_data.nwp_data
        y_test = test_data.era5_data

        x_test, y_test = x_test.data.to(device), y_test.data.to(device)

        y_pred = net(x_test)[0]
        test_c_loss = criterion(y_pred, y_test)
        validating_loss += test_c_loss.item()

    return validating_loss


def train_model(task, device_ids, viz, test_num):
    criterion = task.loss_function
    optimizer = torch.optim.Adam(task.model.parameters(), lr=task.adjustable_parameters["lr"])
    device = torch.device(f"cuda:{device_ids[0]}" if torch.cuda.is_available() else "cpu")

    net = DataParallel(task.model, device_ids=device_ids).to(device)

    train_loader = DataLoader(dataset=task.data[0], batch_size=task.adjustable_parameters["batch_size"], shuffle=True,
                              pin_memory=True)
    model_dir = f'../models/model_{task.info[0]}_{task.info[1]}'
    os.makedirs(model_dir, exist_ok=True)
    best_valid_loss = 1e8

    for epoch in tqdm(range(task.adjustable_parameters["epoch"])):

        lr_init = task.adjustable_parameters["lr"]
        lr_min = task.adjustable_parameters["lr"] / 100
        warmup_epoch = 15
        if epoch == 0:
            lr = lr_init
        elif epoch <= warmup_epoch:
            lr = lr_init * epoch / warmup_epoch
        else:
            lr = lr_min + (lr_init - lr_min) * (1 + math.cos(math.pi * epoch / task.adjustable_parameters["epoch"])) / 2
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        net.train()
        training_one_loss = train_one_epoch(net, train_loader, device, criterion, optimizer)
        net.eval()

        validating_one_loss = validate_one_epoch(net, task.data[1], device, criterion)

        viz.line([[training_one_loss, validating_one_loss]], [epoch],
                 win=f'tt_loss_{test_num}', update='append',
                 opts=dict(title=f'train_test_loss_{test_num}', lenend=['train_loss', 'test_loss']))

        if validating_one_loss < best_valid_loss:
            best_valid_loss = validating_one_loss
            model_save_path = os.path.join(model_dir, f'{test_num}-model.pth')
            torch.save(net, model_save_path)

def test_model(net, test_data, device, diff_std, diff_mean):
    with torch.no_grad():
        x_test = test_data.nwp_data
        x_test = x_test.to(device)
        y_pred = net(x_test)[0]
        y_pred = y_pred.view(5, 6, 25, y_pred.size(1), y_pred.size(2), y_pred.size(3))
        corr_std = diff_std[:, :, :1, np.newaxis, ...] * test_data.nwp_std[:, :, :1, np.newaxis, ...]
        corr_mean = diff_mean[:, :, :1, np.newaxis, ...] + test_data.nwp_mean[:, :, :1, np.newaxis, ...]
        y = np.array(y_pred.to('cpu')) * corr_std + corr_mean

    return y


if __name__ == '__main__':
    pass