ModSync/eval.py at main · heejokong/ModSync · 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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
import os
import numpy as np
import torch

from data_sampler import Burgers_sampler, Helmholtz_sampler, Klein_Gordon_sampler
from models import PINN_Burgers, PINN_Helmholtz, PINN_Klein_Gordon
from utils.parser import get_parser

import plotly.graph_objects as go
import plotly.express as px


def main():
    args = get_parser()

    os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"]= args.gpu_id

    if torch.cuda.is_available():
        args.device = torch.device('cuda')
    else:
        args.device = torch.device('cpu')

    #
    args.depth = 5
    args.width = 128
    args.lr = 1e-3
    seed_list = [0, 1, 2]
    optim_list_1 = ["adam", "lra", "ntk", "pcgrad", "cagrad", "multiadam", "dcgd", "config",]
    optim_list_2 = ["dcgd_sparse", "config_sparse"]

    loss_sum_dict = {}
    loss_res_dict = {}
    loss_ibc_dict = {}
    abs_error_dict = {}
    rel_error_dict = {}

    for optim in optim_list_1:
        loss_sum_dict[optim] = []
        loss_res_dict[optim] = []
        loss_ibc_dict[optim] = []
        abs_error_dict[optim] = []
        rel_error_dict[optim] = []

    for optim in optim_list_2:
        loss_sum_dict[optim] = []
        loss_res_dict[optim] = []
        loss_ibc_dict[optim] = []
        abs_error_dict[optim] = []
        rel_error_dict[optim] = []

    #
    args.backbone = "base"
    for optim in optim_list_1:
        for seed in seed_list:
            args.seed = seed
            np.random.seed(args.seed)
            torch.manual_seed(args.seed)

            if args.equation == 'burgers':
                samplers = Burgers_sampler()
                test_data = samplers.testset()
                model = PINN_Burgers(args, samplers, test_data)

            elif args.equation == 'helmholtz':
                H_samplers = Helmholtz_sampler()
                samplers = H_samplers.samplers()
                test_data = H_samplers.testset()
                model = PINN_Helmholtz(args, samplers, test_data)

            elif args.equation == 'klein_gordon':
                K_samplers = Klein_Gordon_sampler()
                samplers = K_samplers.samplers()
                test_data = K_samplers.testset()
                model = PINN_Klein_Gordon(args, samplers, test_data)

            resume_dir = f"../results/{args.equation}/{optim}/depth_{args.depth}_width_{args.width}/lr_1e-3_seed_{seed}"
            # resume_idx = os.path.join(resume_dir, "model_best.pth")
            resume_idx = os.path.join(resume_dir, "model_checkpoint_50000.pth")
            results_dict = model.test(resume_idx)

            loss_sum_dict[optim].append(results_dict['loss_total'])
            loss_res_dict[optim].append(results_dict['loss_r'])
            loss_ibc_dict[optim].append(results_dict['loss_b'])
            abs_error_dict[optim].append(results_dict['abs_error'])
            rel_error_dict[optim].append(results_dict['rel_error'])

    #
    args.backbone = "dcst"
    for optim in optim_list_2:
        for seed in seed_list:
            args.seed = seed
            np.random.seed(args.seed)
            torch.manual_seed(args.seed)

            if args.equation == 'burgers':
                samplers = Burgers_sampler()
                test_data = samplers.testset()
                model = PINN_Burgers(args, samplers, test_data)

            elif args.equation == 'helmholtz':
                H_samplers = Helmholtz_sampler()
                samplers = H_samplers.samplers()
                test_data = H_samplers.testset()
                model = PINN_Helmholtz(args, samplers, test_data)

            elif args.equation == 'klein_gordon':
                K_samplers = Klein_Gordon_sampler()
                samplers = K_samplers.samplers()
                test_data = K_samplers.testset()
                model = PINN_Klein_Gordon(args, samplers, test_data)

            resume_dir = f"../results/{args.equation}/{optim}/depth_{args.depth}_width_{args.width}/lr_1e-3_seed_{seed}"
            # resume_idx = os.path.join(resume_dir, "model_best.pth")
            resume_idx = os.path.join(resume_dir, "model_checkpoint_50000.pth")
            results_dict = model.test(resume_idx)

            loss_sum_dict[optim].append(results_dict['loss_total'])
            loss_res_dict[optim].append(results_dict['loss_r'])
            loss_ibc_dict[optim].append(results_dict['loss_b'])
            abs_error_dict[optim].append(results_dict['abs_error'])
            rel_error_dict[optim].append(results_dict['rel_error'])

    for kidx in abs_error_dict.keys():
        abs_idx = sum(abs_error_dict[kidx]) / len(seed_list)
        rel_idx = sum(rel_error_dict[kidx]) / len(seed_list)
        print(f"[{kidx.upper():15}] Absolute: {abs_idx:.10}, Relative: {rel_idx:.10}")

if __name__ == "__main__":
    main()