-
Notifications
You must be signed in to change notification settings - Fork 278
/
Copy pathcriterion.py
75 lines (61 loc) · 2.97 KB
/
criterion.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
import torch.nn as nn
import math
import torch
import numpy as np
from torch.nn import functional as F
from torch.autograd import Variable
from .loss import OhemCrossEntropy2d
from .lovasz_losses import lovasz_softmax
import scipy.ndimage as nd
class CriterionDSN(nn.Module):
'''
DSN : We need to consider two supervision for the model.
'''
def __init__(self, ignore_index=255, use_weight=True, reduction='mean'):
super(CriterionDSN, self).__init__()
self.ignore_index = ignore_index
self.criterion = torch.nn.CrossEntropyLoss(ignore_index=ignore_index, reduction=reduction)
if not reduction:
print("disabled the reduction.")
def forward(self, preds, target):
h, w = target.size(1), target.size(2)
if len(preds) >= 2:
scale_pred = F.interpolate(input=preds[0], size=(h, w), mode='bilinear', align_corners=True)
loss1 = self.criterion(scale_pred, target)
scale_pred = F.interpolate(input=preds[1], size=(h, w), mode='bilinear', align_corners=True)
loss2 = self.criterion(scale_pred, target)
return loss1 + loss2*0.4
else:
scale_pred = F.interpolate(input=preds[0], size=(h, w), mode='bilinear', align_corners=True)
loss = self.criterion(scale_pred, target)
return loss
class CriterionOhemDSN(nn.Module):
'''
DSN : We need to consider two supervision for the model.
'''
def __init__(self, ignore_index=255, thresh=0.7, min_kept=100000, use_weight=True, reduction='mean'):
super(CriterionOhemDSN, self).__init__()
self.ignore_index = ignore_index
self.criterion1 = OhemCrossEntropy2d(ignore_index, thresh, min_kept)
self.criterion2 = torch.nn.CrossEntropyLoss(ignore_index=ignore_index, reduction=reduction)
def forward(self, preds, target):
h, w = target.size(1), target.size(2)
scale_pred = F.interpolate(input=preds[0], size=(h, w), mode='bilinear', align_corners=True)
loss1 = self.criterion1(scale_pred, target)
scale_pred = F.interpolate(input=preds[1], size=(h, w), mode='bilinear', align_corners=True)
loss2 = self.criterion2(scale_pred, target)
return loss1 + loss2*0.4
class CriterionOhemDSN2(nn.Module):
'''
DSN : We need to consider two supervision for the model.
'''
def __init__(self, ignore_index=255, thresh=0.7, min_kept=100000, use_weight=True, reduction='mean'):
super(CriterionOhemDSN2, self).__init__()
self.ignore_index = ignore_index
self.criterion = torch.nn.CrossEntropyLoss(ignore_index=ignore_index, reduction=reduction)
def forward(self, preds, target):
h, w = target.size(1), target.size(2)
scale_pred = F.interpolate(input=preds[0], size=(h, w), mode='bilinear', align_corners=True)
loss1 = self.criterion(scale_pred, target)
loss2 = lovasz_softmax(F.softmax(scale_pred, dim=1), target, ignore=self.ignore_index)
return loss1 + loss2