-
Notifications
You must be signed in to change notification settings - Fork 4
/
pimlturb1d_estimator.py
300 lines (231 loc) · 12 KB
/
pimlturb1d_estimator.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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
"""
Physics-Informed Convolutional model to predict
ratio of turbulent pressure to gas pressure in 1D.
The loss function here is a custom combination of
statistical (Kolmogorov-Smirnov) and spatial (Smooth L1)
losses.
"""
import numpy as np
import sys
import os
import torch
from sapsan.core.models import EstimatorConfig
from sapsan.lib.estimator.torch_backend import TorchBackend
from sapsan.lib.data import get_loader_shape
from sapsan.lib.estimator.torch_modules import Interp1d
from sapsan.lib.estimator.torch_modules import Gaussian
class PIMLTurb1DModel(torch.nn.ModuleDict):
def __init__(self, D_in = 1, D_out = 1, activ = "ReLU", sigma=1):
super(PIMLTurb1DModel, self).__init__()
self.conv1d = torch.nn.Conv1d(D_in, D_in*2, kernel_size=3, stride=2)
self.pool = torch.nn.MaxPool1d(kernel_size=2)
self.activation = getattr(torch.nn, activ)()
self.linear = torch.nn.Linear(D_in*98, D_in*196)
self.linear2 = torch.nn.Linear(D_in*196, D_out)
self.gaussian = Gaussian(sigma=sigma, axis=1)
def forward(self, x):
x = x.float()
x_shape = x.size()
x = self.conv1d(x)
x = self.pool(x)
x = self.activation(x)
x = x.view(x.size(0), -1)
x = self.linear(x)
x = self.activation(x)
x = self.linear2(x)
x_shape = list(x_shape)
x_shape[1] = 1
x = torch.reshape(x,x_shape)
x = torch.mul(x,x)
x = self.gaussian(x)
return x
class SmoothL1_KSLoss(torch.nn.Module):
'''
The loss functions combines a statistical (Kolmogorov-Smirnov)
and a spatial loss (Smooth L1).
Corresponding 'train_l1ks_log.txt' and 'valid_l1ks_log.txt'
are written out to include the individual loss evolutions.
'''
def __init__(self, ks_stop, ks_frac, ks_scale, l1_scale, beta, train_size, valid_size):
super(SmoothL1_KSLoss, self).__init__()
self.first_write = True
self.first_iter = True
self.ks_stop = ks_stop
self.ks_frac = ks_frac
self.ks_scale = ks_scale
self.l1_scale = l1_scale
self.beta = beta
self.train_size = train_size
self.valid_size = valid_size
self.stop = False
def write_log(self, losses):
if self.first_write:
if os.path.exists(self.log_fname): os.remove(self.log_fname)
with open(self.log_fname,'a') as f:
if self.first_write:
f.write(f"mean(L1_loss) \t mean(KS_loss) "\
f"\t mean(L1_loss)*{self.l1_scale:.3e} \t mean(KS_loss)*{self.ks_scale:.3e} \n")
np.savetxt(f, [losses.detach().cpu().numpy()], fmt='%.3e \t %.3e \t %.3e \t %.3e')
def forward(self, predictions, targets, write, write_idx):
if predictions.is_cuda:
self.device = torch.device('cuda:%d'%predictions.get_device())
else: self.device = torch.device('cpu')
#-----SmoothL1----
l1_loss = 0
self.beta = 0.1*targets.max()
diff = predictions-targets
mask = (diff.abs() < self.beta)
l1_loss += mask * (0.5*diff**2 / self.beta)
l1_loss += (~mask) * (diff.abs() - 0.5*self.beta)
#--------KS-------
distr = torch.distributions.normal.Normal(loc=targets.mean(),
scale=targets.std(),
validate_args=False)
cdf,idx = distr.cdf(targets).flatten().to(self.device).sort()
distr = torch.distributions.normal.Normal(loc=predictions.mean(),
scale=predictions.std(),
validate_args=False)
cdf_pred,idx_pred = distr.cdf(predictions).flatten().to(self.device).sort()
vals,idx = targets.flatten().to(self.device).sort()
vals_pred,idx_pred = predictions.flatten().to(self.device).sort()
cdf_interp = Interp1d()(vals_pred, cdf_pred, vals)[0]
ks_loss = torch.max(torch.abs(cdf-cdf_interp)).to(self.device)
#--Print Metrics--
if self.first_iter:
self.maxl1 = l1_loss.mean().detach()
self.maxks = ks_loss.mean().detach()
self.max_ratio = self.maxks/self.maxl1
print('---Initial (max) L1, KS losses and their ratio--')
print(self.maxl1, self.maxks, self.max_ratio)
print('------------------------------------------------')
self.first_iter = False
#----Calc Loss----
#fractions that KS and L1 contribute to the total loss
l1_frac = 1-self.ks_frac
loss = (l1_frac*l1_loss.mean()*self.l1_scale +
self.ks_frac*ks_loss.mean()*self.ks_scale)
#---Save Batch----
if write_idx == 0:
if write == 'train':
self.batch_size = self.train_size
self.log_fname = "train_l1ks_log.txt"
elif write == 'valid':
self.batch_size = self.valid_size
self.log_fname = "valid_l1ks_log.txt"
self.iter_losses_l1 = torch.zeros(self.batch_size,requires_grad=False).to(self.device)
self.iter_losses_ks = torch.zeros(self.batch_size,requires_grad=False).to(self.device)
self.iter_losses_l1[write_idx] = l1_loss.mean().detach()
self.iter_losses_ks[write_idx] = ks_loss.mean().detach()
#---Write Logs----
if write_idx == (self.batch_size-1):
mean_iter_l1 = self.iter_losses_l1.mean()
mean_iter_ks = self.iter_losses_ks.mean()
if mean_iter_ks < self.ks_stop or self.stop==True:
self.stop = True
self.write_log(torch.tensor([mean_iter_l1, mean_iter_ks,
mean_iter_l1*self.l1_scale,
mean_iter_ks*self.ks_scale]))
if write == 'valid': self.first_write = False
return loss, self.stop
class PIMLTurb1DConfig(EstimatorConfig):
def __init__(self,
n_epochs: int = 1,
patience: int = 10,
min_delta: float = 1e-5,
logdir: str = "./logs/",
lr: float = 1e-3,
min_lr = None,
*args, **kwargs):
self.n_epochs = n_epochs
self.logdir = logdir
self.patience = patience
self.min_delta = min_delta
self.lr = lr
if min_lr==None: self.min_lr = lr*1e-2
else: self.min_lr = min_lr
self.kwargs = kwargs
self.parameters = {f'model - {k}': v for k, v in self.__dict__.items() if k != 'kwargs'}
if bool(self.kwargs): self.parameters.update({f'model - {k}': v for k, v in self.kwargs.items()})
class PIMLTurb1D(TorchBackend):
def __init__(self, activ,
loss,
loaders,
ks_stop = 0.1,
ks_frac = 0.5,
ks_scale = 1,
l1_scale = 1,
l1_beta = 1,
sigma = 1,
config = PIMLTurb1DConfig(),
model = PIMLTurb1DModel()):
super().__init__(config, model)
self.config = config
self.loaders = loaders
self.device = self.config.kwargs['device']
self.ks_stop = ks_stop
self.ks_frac = ks_frac
self.ks_scale = ks_scale
self.l1_scale = l1_scale
self.beta = l1_beta
x_shape, y_shape = get_loader_shape(self.loaders)
self.model = PIMLTurb1DModel(x_shape[1], y_shape[2], activ, sigma)
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config.lr)
if loss == "SmoothL1_KSLoss":
self.loss_func = SmoothL1_KSLoss(ks_stop = self.ks_stop, ks_frac = self.ks_frac,
ks_scale = self.ks_scale, l1_scale = self.l1_scale,
beta = self.beta,
train_size = len(self.loaders['train']),
valid_size = len(self.loaders['valid']))
else:
print("This network only supports the custom SmoothL1_KSLoss. Please set the 'loss' parameter in config")
sys.exit()
self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer,
patience=self.config.patience,
min_lr=self.config.min_lr)
self.config.parameters["model - loss"] = str(self.loss_func).partition("(")[0]
self.config.parameters["model - activ"] = activ
def write_loss(self, losses, epoch, fname_train = "train_log.txt", fname_valid='valid_log.txt'):
for idx, fname in enumerate([fname_train, fname_valid]):
if epoch==1:
if os.path.exists(fname): os.remove(fname)
with open(fname,'a') as f:
if epoch == 1: f.write("epoch \t loss \n")
np.savetxt(f, [[losses[0],losses[idx+1]]], fmt='%d \t %.3e')
def train(self):
self.model = self.model#.double()
self.model.to(self.device)
if 'cuda' in str(self.device):
self.optimizer_to(self.optimizer, self.device)
for epoch in np.linspace(1,int(self.config.n_epochs),int(self.config.n_epochs), dtype='int'):
iter_loss = np.zeros(len(self.loaders['train']))
for idx, (x, y) in enumerate(self.loaders['train']):
x = x.to(self.device)
y = y.to(self.device)
y_pred = self.model(x)
loss, stop = self.loss_func(y_pred, y, 'train', idx)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
iter_loss[idx] = loss.item()
epoch_loss = iter_loss.mean()
print(f'train ({epoch}/{int(self.config.n_epochs)}) loss: {epoch_loss:.4e}')
with torch.set_grad_enabled(False):
iter_loss_valid = np.zeros(len(self.loaders['valid']))
for idx, (x, y) in enumerate(self.loaders['valid']):
x = x.to(self.device)
y = y.to(self.device)
y_pred = self.model(x)
loss_valid, stop = self.loss_func(y_pred, y, 'valid', idx)
iter_loss_valid[idx] = loss_valid.item()
epoch_loss_valid = iter_loss_valid.mean()
print(f'valid ({epoch}/{int(self.config.n_epochs)}) loss: {epoch_loss_valid:.4e}')
self.write_loss([epoch, epoch_loss, epoch_loss_valid], epoch)
if stop:
print('Reached sufficient KS Loss, stopping...')
break
print('-----')
with open('model_details.txt', 'w') as file:
file.write('%s\n\n%s\n\n%s'%(str(self.model),
str(self.optimizer),
str(self.scheduler)))
return self.model