File size: 5,676 Bytes
47162d0 |
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 |
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@Author : Peike Li
@Contact : [email protected]
@File : criterion.py
@Time : 8/30/19 8:59 PM
@Desc :
@License : This source code is licensed under the license found in the
LICENSE file in the root directory of this source tree.
"""
import torch.nn as nn
import torch
import numpy as np
from torch.nn import functional as F
from .lovasz_softmax import LovaszSoftmax
from .kl_loss import KLDivergenceLoss
from .consistency_loss import ConsistencyLoss
NUM_CLASSES = 20
class CriterionAll(nn.Module):
def __init__(self, use_class_weight=False, ignore_index=255, lambda_1=1, lambda_2=1, lambda_3=1,
num_classes=20):
super(CriterionAll, self).__init__()
self.ignore_index = ignore_index
self.use_class_weight = use_class_weight
self.criterion = torch.nn.CrossEntropyLoss(ignore_index=ignore_index)
self.lovasz = LovaszSoftmax(ignore_index=ignore_index)
self.kldiv = KLDivergenceLoss(ignore_index=ignore_index)
self.reg = ConsistencyLoss(ignore_index=ignore_index)
self.lamda_1 = lambda_1
self.lamda_2 = lambda_2
self.lamda_3 = lambda_3
self.num_classes = num_classes
def parsing_loss(self, preds, target, cycle_n=None):
"""
Loss function definition.
Args:
preds: [[parsing result1, parsing result2],[edge result]]
target: [parsing label, egde label]
soft_preds: [[parsing result1, parsing result2],[edge result]]
Returns:
Calculated Loss.
"""
h, w = target[0].size(1), target[0].size(2)
pos_num = torch.sum(target[1] == 1, dtype=torch.float)
neg_num = torch.sum(target[1] == 0, dtype=torch.float)
weight_pos = neg_num / (pos_num + neg_num)
weight_neg = pos_num / (pos_num + neg_num)
weights = torch.tensor([weight_neg, weight_pos]) # edge loss weight
loss = 0
# loss for segmentation
preds_parsing = preds[0]
for pred_parsing in preds_parsing:
scale_pred = F.interpolate(input=pred_parsing, size=(h, w),
mode='bilinear', align_corners=True)
loss += 0.5 * self.lamda_1 * self.lovasz(scale_pred, target[0])
if target[2] is None:
loss += 0.5 * self.lamda_1 * self.criterion(scale_pred, target[0])
else:
soft_scale_pred = F.interpolate(input=target[2], size=(h, w),
mode='bilinear', align_corners=True)
soft_scale_pred = moving_average(soft_scale_pred, to_one_hot(target[0], num_cls=self.num_classes),
1.0 / (cycle_n + 1.0))
loss += 0.5 * self.lamda_1 * self.kldiv(scale_pred, soft_scale_pred, target[0])
# loss for edge
preds_edge = preds[1]
for pred_edge in preds_edge:
scale_pred = F.interpolate(input=pred_edge, size=(h, w),
mode='bilinear', align_corners=True)
if target[3] is None:
loss += self.lamda_2 * F.cross_entropy(scale_pred, target[1],
weights.cuda(), ignore_index=self.ignore_index)
else:
soft_scale_edge = F.interpolate(input=target[3], size=(h, w),
mode='bilinear', align_corners=True)
soft_scale_edge = moving_average(soft_scale_edge, to_one_hot(target[1], num_cls=2),
1.0 / (cycle_n + 1.0))
loss += self.lamda_2 * self.kldiv(scale_pred, soft_scale_edge, target[0])
# consistency regularization
preds_parsing = preds[0]
preds_edge = preds[1]
for pred_parsing in preds_parsing:
scale_pred = F.interpolate(input=pred_parsing, size=(h, w),
mode='bilinear', align_corners=True)
scale_edge = F.interpolate(input=preds_edge[0], size=(h, w),
mode='bilinear', align_corners=True)
loss += self.lamda_3 * self.reg(scale_pred, scale_edge, target[0])
return loss
def forward(self, preds, target, cycle_n=None):
loss = self.parsing_loss(preds, target, cycle_n)
return loss
def _generate_weights(self, masks, num_classes):
"""
masks: torch.Tensor with shape [B, H, W]
"""
masks_label = masks.data.cpu().numpy().astype(np.int64)
pixel_nums = []
tot_pixels = 0
for i in range(num_classes):
pixel_num_of_cls_i = np.sum(masks_label == i).astype(np.float)
pixel_nums.append(pixel_num_of_cls_i)
tot_pixels += pixel_num_of_cls_i
weights = []
for i in range(num_classes):
weights.append(
(tot_pixels - pixel_nums[i]) / tot_pixels / (num_classes - 1)
)
weights = np.array(weights, dtype=np.float)
# weights = torch.from_numpy(weights).float().to(masks.device)
return weights
def moving_average(target1, target2, alpha=1.0):
target = 0
target += (1.0 - alpha) * target1
target += target2 * alpha
return target
def to_one_hot(tensor, num_cls, dim=1, ignore_index=255):
b, h, w = tensor.shape
tensor[tensor == ignore_index] = 0
onehot_tensor = torch.zeros(b, num_cls, h, w).cuda()
onehot_tensor.scatter_(dim, tensor.unsqueeze(dim), 1)
return onehot_tensor
|