"""Custom losses.""" import torch import torch.nn as nn import torch.nn.functional as F from torch.autograd import Variable __all__ = ['MixSoftmaxCrossEntropyLoss', 'MixSoftmaxCrossEntropyOHEMLoss', 'EncNetLoss', 'ICNetLoss', 'get_segmentation_loss'] # TODO: optim function class MixSoftmaxCrossEntropyLoss(nn.CrossEntropyLoss): def __init__(self, aux=True, aux_weight=0.2, ignore_index=-1, **kwargs): super(MixSoftmaxCrossEntropyLoss, self).__init__(ignore_index=ignore_index) self.aux = aux self.aux_weight = aux_weight def _aux_forward(self, *inputs, **kwargs): *preds, target = tuple(inputs) loss = super(MixSoftmaxCrossEntropyLoss, self).forward(preds[0], target) for i in range(1, len(preds)): aux_loss = super(MixSoftmaxCrossEntropyLoss, self).forward(preds[i], target) loss += self.aux_weight * aux_loss return loss def forward(self, *inputs, **kwargs): preds, target = tuple(inputs) inputs = tuple(list(preds) + [target]) if self.aux: return dict(loss=self._aux_forward(*inputs)) else: return dict(loss=super(MixSoftmaxCrossEntropyLoss, self).forward(*inputs)) # reference: https://github.com/zhanghang1989/PyTorch-Encoding/blob/master/encoding/nn/loss.py class EncNetLoss(nn.CrossEntropyLoss): """2D Cross Entropy Loss with SE Loss""" def __init__(self, se_loss=True, se_weight=0.2, nclass=19, aux=False, aux_weight=0.4, weight=None, ignore_index=-1, **kwargs): super(EncNetLoss, self).__init__(weight, None, ignore_index) self.se_loss = se_loss self.aux = aux self.nclass = nclass self.se_weight = se_weight self.aux_weight = aux_weight self.bceloss = nn.BCELoss(weight) def forward(self, *inputs): preds, target = tuple(inputs) inputs = tuple(list(preds) + [target]) if not self.se_loss and not self.aux: return super(EncNetLoss, self).forward(*inputs) elif not self.se_loss: pred1, pred2, target = tuple(inputs) loss1 = super(EncNetLoss, self).forward(pred1, target) loss2 = super(EncNetLoss, self).forward(pred2, target) return dict(loss=loss1 + self.aux_weight * loss2) elif not self.aux: print (inputs) pred, se_pred, target = tuple(inputs) se_target = self._get_batch_label_vector(target, nclass=self.nclass).type_as(pred) loss1 = super(EncNetLoss, self).forward(pred, target) loss2 = self.bceloss(torch.sigmoid(se_pred), se_target) return dict(loss=loss1 + self.se_weight * loss2) else: pred1, se_pred, pred2, target = tuple(inputs) se_target = self._get_batch_label_vector(target, nclass=self.nclass).type_as(pred1) loss1 = super(EncNetLoss, self).forward(pred1, target) loss2 = super(EncNetLoss, self).forward(pred2, target) loss3 = self.bceloss(torch.sigmoid(se_pred), se_target) return dict(loss=loss1 + self.aux_weight * loss2 + self.se_weight * loss3) @staticmethod def _get_batch_label_vector(target, nclass): # target is a 3D Variable BxHxW, output is 2D BxnClass batch = target.size(0) tvect = Variable(torch.zeros(batch, nclass)) for i in range(batch): hist = torch.histc(target[i].cpu().data.float(), bins=nclass, min=0, max=nclass - 1) vect = hist > 0 tvect[i] = vect return tvect # TODO: optim function class ICNetLoss(nn.CrossEntropyLoss): """Cross Entropy Loss for ICNet""" def __init__(self, nclass, aux_weight=0.4, ignore_index=-1, **kwargs): super(ICNetLoss, self).__init__(ignore_index=ignore_index) self.nclass = nclass self.aux_weight = aux_weight def forward(self, *inputs): preds, target = tuple(inputs) inputs = tuple(list(preds) + [target]) pred, pred_sub4, pred_sub8, pred_sub16, target = tuple(inputs) # [batch, W, H] -> [batch, 1, W, H] target = target.unsqueeze(1).float() target_sub4 = F.interpolate(target, pred_sub4.size()[2:], mode='bilinear', align_corners=True).squeeze(1).long() target_sub8 = F.interpolate(target, pred_sub8.size()[2:], mode='bilinear', align_corners=True).squeeze(1).long() target_sub16 = F.interpolate(target, pred_sub16.size()[2:], mode='bilinear', align_corners=True).squeeze( 1).long() loss1 = super(ICNetLoss, self).forward(pred_sub4, target_sub4) loss2 = super(ICNetLoss, self).forward(pred_sub8, target_sub8) loss3 = super(ICNetLoss, self).forward(pred_sub16, target_sub16) return dict(loss=loss1 + loss2 * self.aux_weight + loss3 * self.aux_weight) class OhemCrossEntropy2d(nn.Module): def __init__(self, ignore_index=-1, thresh=0.7, min_kept=100000, use_weight=True, **kwargs): super(OhemCrossEntropy2d, self).__init__() self.ignore_index = ignore_index self.thresh = float(thresh) self.min_kept = int(min_kept) if use_weight: weight = torch.FloatTensor([0.8373, 0.918, 0.866, 1.0345, 1.0166, 0.9969, 0.9754, 1.0489, 0.8786, 1.0023, 0.9539, 0.9843, 1.1116, 0.9037, 1.0865, 1.0955, 1.0865, 1.1529, 1.0507]) self.criterion = torch.nn.CrossEntropyLoss(weight=weight, ignore_index=ignore_index) else: self.criterion = torch.nn.CrossEntropyLoss(ignore_index=ignore_index) def forward(self, pred, target): n, c, h, w = pred.size() target = target.view(-1) valid_mask = target.ne(self.ignore_index) target = target * valid_mask.long() num_valid = valid_mask.sum() prob = F.softmax(pred, dim=1) prob = prob.transpose(0, 1).reshape(c, -1) if self.min_kept > num_valid: print("Lables: {}".format(num_valid)) elif num_valid > 0: prob = prob.masked_fill_(1 - valid_mask, 1) mask_prob = prob[target, torch.arange(len(target), dtype=torch.long)] threshold = self.thresh if self.min_kept > 0: index = mask_prob.argsort() threshold_index = index[min(len(index), self.min_kept) - 1] if mask_prob[threshold_index] > self.thresh: threshold = mask_prob[threshold_index] kept_mask = mask_prob.le(threshold) valid_mask = valid_mask * kept_mask target = target * kept_mask.long() target = target.masked_fill_(1 - valid_mask, self.ignore_index) target = target.view(n, h, w) return self.criterion(pred, target) class MixSoftmaxCrossEntropyOHEMLoss(OhemCrossEntropy2d): def __init__(self, aux=False, aux_weight=0.4, weight=None, ignore_index=-1, **kwargs): super(MixSoftmaxCrossEntropyOHEMLoss, self).__init__(ignore_index=ignore_index) self.aux = aux self.aux_weight = aux_weight self.bceloss = nn.BCELoss(weight) def _aux_forward(self, *inputs, **kwargs): *preds, target = tuple(inputs) loss = super(MixSoftmaxCrossEntropyOHEMLoss, self).forward(preds[0], target) for i in range(1, len(preds)): aux_loss = super(MixSoftmaxCrossEntropyOHEMLoss, self).forward(preds[i], target) loss += self.aux_weight * aux_loss return loss def forward(self, *inputs): preds, target = tuple(inputs) inputs = tuple(list(preds) + [target]) if self.aux: return dict(loss=self._aux_forward(*inputs)) else: return dict(loss=super(MixSoftmaxCrossEntropyOHEMLoss, self).forward(*inputs)) def get_segmentation_loss(model, use_ohem=False, **kwargs): if use_ohem: return MixSoftmaxCrossEntropyOHEMLoss(**kwargs) model = model.lower() if model == 'encnet': return EncNetLoss(**kwargs) elif model == 'icnet': return ICNetLoss(nclass=4, **kwargs) else: return MixSoftmaxCrossEntropyLoss(**kwargs)