/
loss.py
51 lines (31 loc) · 1.35 KB
/
loss.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
import numpy as np
import torch
from torch.autograd import Variable
import torch.nn.functional as F
from metrics import dice_coef
BCE = torch.nn.BCELoss()
def weighted_loss(pred, mask):
weit = 1 + 5*torch.abs(F.avg_pool2d(mask, kernel_size=31, stride=1, padding=15) - mask)
wbce = F.binary_cross_entropy_with_logits(pred, mask, reduce='none')
wbce = (weit*wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3))
# pred = torch.sigmoid(pred)
inter = ((pred * mask)*weit).sum(dim=(2, 3))
union = ((pred + mask)*weit).sum(dim=(2, 3))
wiou = 1 - (inter + 1)/(union - inter+1)
return (wbce + wiou).mean()
def calc_loss(pred, target, bce_weight=0.5):
bce = weighted_loss(pred, target)
# dl = 1 - dice_coef(pred, target)
# loss = bce * bce_weight + dl * bce_weight
return bce
def loss_sup(logit_S1, logit_S2, labels_S1, labels_S2):
loss1 = calc_loss(logit_S1, labels_S1)
loss2 = calc_loss(logit_S2, labels_S2)
return loss1 + loss2
def loss_diff(u_prediction_1, u_prediction_2, batch_size):
a = weighted_loss(u_prediction_1, Variable(u_prediction_2, requires_grad=False))
a = a.item()
b = weighted_loss(u_prediction_2, Variable(u_prediction_1, requires_grad=False))
b = b.item()
loss_diff_avg = (a + b)
return loss_diff_avg / batch_size