Repository navigation
Expand file tree
/
Copy pathcriterion.py
More file actions
105 lines (91 loc) · 4.07 KB
/
Copy pathcriterion.py
File metadata and controls
105 lines (91 loc) · 4.07 KB
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
import torch
import torch.nn as nn
import torch.nn.functional as F
from pytorch_msssim import SSIM, MS_SSIM
class GeoGradient(nn.Module):
def __init__(self, device='cuda'):
super().__init__()
self.conv_x = nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False)
self.conv_y = nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False)
scharr_x = torch.tensor([[3., 0., -3.],
[10., 0., -10.],
[3., 0., -3.]])
scharr_y = torch.tensor([[3., 10., 3.],
[0., 0., 0.],
[-3., -10., -3.]])
self.conv_x.weight.data = scharr_x.view(1, 1, 3, 3).to(device)
self.conv_y.weight.data = scharr_y.view(1, 1, 3, 3).to(device)
self.conv_x.weight.requires_grad_(False)
self.conv_y.weight.requires_grad_(False)
def forward(self, x):
x = x.unsqueeze(1)
x_pad = F.pad(x, (1, 1, 1, 1), mode='reflect')
gx = self.conv_x(x_pad)
gy = self.conv_y(x_pad)
gx = gx[:, :, 1:-1, 1:-1] # (B, 1, H, W)
gy = gy[:, :, 1:-1, 1:-1] # (B, 1, H, W)
return torch.cat([gx, gy], dim=1) # (B, 2, H, W)
class SalinityCriterion(nn.Module):
def __init__(self,
grad_weight=0.3,
ssim_weight=0.2,
use_msssim=False,
window_size=11,
device='cuda'):
super().__init__()
self.grad_weight = grad_weight
self.ssim_weight = ssim_weight
self.base_loss = nn.SmoothL1Loss(reduction='none')
self.gradient = GeoGradient(device=device)
if use_msssim:
self.ssim = MS_SSIM(
data_range=5.0,
win_size=window_size,
channel=1
)
else:
self.ssim = SSIM(
data_range=5.0,
win_size=window_size,
channel=1
)
def gradient_loss(self, pred_data, label_data, mask):
pred_grad = self.gradient(pred_data)
label_grad = self.gradient(label_data)
cos_sim = F.cosine_similarity(pred_grad, label_grad, dim=1)
valid_mask = mask & ~torch.isnan(label_data)
loss = 1 - cos_sim
masked_loss = torch.where(valid_mask, loss, torch.zeros_like(loss))
return masked_loss.sum() / (valid_mask.sum() + 1e-8)
def masked_ssim(self, pred_data, label_data, mask):
pred_norm = (pred_data - pred_data.min()) / (pred_data.max() - pred_data.min() + 1e-8) * 5.0
label_norm = (label_data - label_data.min()) / (label_data.max() - label_data.min() + 1e-8) * 5.0
ssim_val = self.ssim(pred_norm, label_norm)
valid_ratio = mask.float().mean()
return (1 - ssim_val) * valid_ratio
def forward(self, pred_data, label_data, mask, epoch=None, total_epochs=None):
pred_3d = pred_data.squeeze(1)
label_3d = label_data.squeeze(1)
mask_3d = mask.squeeze(1)
base_loss = self.base_loss(pred_3d, label_3d)
base_loss = torch.where(mask, base_loss, torch.zeros_like(base_loss))
base_loss = base_loss.sum() / (mask.sum() + 1e-8)
grad_loss = self.gradient_loss(pred_3d, label_3d, mask_3d)
ssim_loss = self.masked_ssim(pred_data, label_data, mask)
if epoch is not None and total_epochs is not None:
epoch_ratio = epoch / float(total_epochs)
if epoch_ratio < 0.3:
w_grad = 0.0
w_ssim = 0.0
elif epoch_ratio < 0.7:
progress = (epoch_ratio - 0.3) / 0.4
w_grad = self.grad_weight * progress
w_ssim = self.ssim_weight * progress
else:
w_grad = self.grad_weight
w_ssim = self.ssim_weight
else:
w_grad = self.grad_weight
w_ssim = self.ssim_weight
total_loss = base_loss + w_grad * grad_loss + w_ssim * ssim_loss
return total_loss