-
Notifications
You must be signed in to change notification settings - Fork 5
/
exp_elas_interp.py
120 lines (96 loc) · 3.67 KB
/
exp_elas_interp.py
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
import torch.nn.functional as F
import matplotlib.pyplot as plt
from timeit import default_timer
from utils.utilities3 import *
from utils.adam import Adam
from utils.params import get_args
from model_dict import get_model
import math
import os
torch.manual_seed(0)
np.random.seed(0)
torch.cuda.manual_seed(0)
torch.backends.cudnn.deterministic = True
################################################################
# configs
################################################################
args = get_args()
INPUT_PATH = os.path.join(args.data_path, './Interp/Random_UnitCell_mask_10_interp.npy')
OUTPUT_PATH = os.path.join(args.data_path, './Interp/Random_UnitCell_sigma_10_interp.npy')
ntrain = args.ntrain
ntest = args.ntest
N = args.ntotal
in_channels = args.in_dim
out_channels = args.out_dim
r1 = args.h_down
r2 = args.w_down
s1 = int(((args.h - 1) / r1) + 1)
s2 = int(((args.w - 1) / r2) + 1)
batch_size = args.batch_size
learning_rate = args.learning_rate
epochs = args.epochs
step_size = args.step_size
gamma = args.gamma
model_save_path = args.model_save_path
model_save_name = args.model_save_name
################################################################
# models
################################################################
model = get_model(args)
print(count_params(model))
################################################################
# load data and data normalization
################################################################
input = np.load(INPUT_PATH)
input = torch.tensor(input, dtype=torch.float).permute(2, 0, 1)
output = np.load(OUTPUT_PATH)
output = torch.tensor(output, dtype=torch.float).permute(2, 0, 1)
x_train = input[:N][:ntrain, ::r1, ::r2][:, :s1, :s2]
y_train = output[:N][:ntrain, ::r1, ::r2][:, :s1, :s2]
x_test = input[:N][-ntest:, ::r1, ::r2][:, :s1, :s2]
y_test = output[:N][-ntest:, ::r1, ::r2][:, :s1, :s2]
x_train = x_train.reshape(ntrain, s1, s2, 1)
x_test = x_test.reshape(ntest, s1, s2, 1)
train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_train, y_train), batch_size=batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_test, y_test), batch_size=batch_size,
shuffle=False)
################################################################
# training and evaluation
################################################################
optimizer = Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)
myloss = LpLoss(size_average=False)
for ep in range(epochs):
model.train()
t1 = default_timer()
train_l2 = 0
for x, y in train_loader:
x, y = x.cuda(), y.cuda()
mask = x.clone()
optimizer.zero_grad()
out = model(x)
out = out * mask
loss = myloss(out.view(batch_size, -1), y.view(batch_size, -1))
loss.backward()
optimizer.step()
train_l2 += loss.item()
scheduler.step()
model.eval()
test_l2 = 0.0
with torch.no_grad():
for x, y in test_loader:
x, y = x.cuda(), y.cuda()
mask = x.clone()
out = model(x)
out2 = out * mask
test_l2 += myloss(out2.view(batch_size, -1), y.view(batch_size, -1)).item()
train_l2 /= ntrain
test_l2 /= ntest
t2 = default_timer()
print(ep, t2 - t1, train_l2, test_l2)
if ep % step_size == 0:
if not os.path.exists(model_save_path):
os.makedirs(model_save_path)
print('save model')
torch.save(model.state_dict(), os.path.join(model_save_path, model_save_name))