深度估计berHu损失函数和语义分割带权值交叉熵损失函数

深度估计berHu损失函数和语义分割带权值交叉熵损失函数

最近在做联合深度估计和语义分割的深度学习算法，深度估计默认使用的是L1_loss，语义分割使用的是普通的交叉熵损失函数，继续改进模型对于指标的提升微乎甚微，也不想再改了，就想着看看损失函数的小变化，能否对结果产生一点积极的影响。

berHu Loss

berHu损失函数包含了一范数损失和二范数损失两个部分，当预测值与真实值差值“不是很大的时候”采用一范数，过大时，采用二范数，因此对稳定网络训练有一定积极作用。
$$L^D\left(y_i,y_i^{\prime }\right)=\{\begin{array}{l}\left|y_i-y_i^{\prime }\right|,\left|y_i-y_i^{\prime }\right|\le \delta \\ \frac{{\left(y_i-y_i^{\prime }\right)}^2+{\delta }^2}{2\delta },\text{ otherwise }\end{array},$$
其中$y_i,y_i^{\prime }$分别表示预测值和真实值，$\delta =0.2max|y_i-y_i^{\prime }|$。具体的代码实现如下所示:

class BerHu_loss(nn.Module):
    def __init__(self, c=0.2, ignore_index=255):
        super(BerHu_loss, self).__init__()
        self.c = c
        self.ignore_index = ignore_index

    def forward(self, out, label, reduction='mean'):
        """out and label shape both are [B, 1, H, W], float type"""
        mask = (label != self.ignore_index).all(dim=1, keepdim=True)  
        n_valid = torch.sum(mask).item()
        # masked_select输出一个一维向量  
        masked_out = torch.masked_select(out, mask)         # 预测值
        masked_label = torch.masked_select(label, mask)     # 真实值

        diff = torch.abs(masked_out - masked_label)
        delta = self.c * torch.max(diff).item()             # delta is scaler
        berhu_loss = torch.where(diff <= delta, diff, (diff**2 + delta**2) / (2 * delta))

        if reduction == 'mean':
            # return torch.mean(berhu_loss)
            return berhu_loss.sum() / max(n_valid, 1)
        elif reduction == 'sum':
            return torch.sum(berhu_loss)
        elif reduction == 'none':
            return berhu_loss

采用一个简单的网络测试一下：

    class MyModel(nn.Module):
        def __init__(self):
            super(MyModel, self).__init__()
            self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
            self.conv2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
            self.conv3 = nn.Conv2d(64, 1, kernel_size=3, padding=1)
            self.relu = nn.ReLU()
            self.berhu_loss = BerHu_loss()

        def forward(self, x, target):
            x = self.relu(self.conv1(x))
            x = self.relu(self.conv2(x))
            x = self.conv3(x)
            loss = self.berhu_loss(x, target)
            return x, loss


    model = MyModel()

    x_input = torch.rand(2, 3, 64, 64)
    label = torch.rand(2, 1, 64, 64)

    out, loss = model(x_input, label)
    print(out.shape, loss)

带权重的交叉熵loss

带权重的交叉熵loss，公式为：
$$L=-\sum _{c=1}^M{w}_c{y}_c\log \left(p_c\right)$$
可以看到只是在交叉熵Loss的基础上为每一个语义类别添加了一个权重参数，其中$w_c$的计算公式为： $w_c=\frac{N-N_c}{N}$ 其中$N$表示总的像素个数，而$N_c$表示GT类别为$c$的像素个数。这样相比于原始的交叉熵Loss，在样本数量极不均衡的情况下可以获得更好的效果。

本次记录的是如何计算数据集类别权重$w_c$，NYUD-V2是一个很常用的多任务数据集，包含深度估计、语义分割、法线估计等任务。其中语义分割有40个类别（标签号码为1-40），加上背景类（0）一共有41个类别，但在实际训练过程中，nyud-v2数据集的背景类一般不参与计算，也不参与miou等指标计算。类别权重计算的代码如下：

# 类别名称
NYU_CATEGORY_NAMES = ['wall', 'floor', 'cabinet', 'bed', 'chair',
                      'sofa', 'table', 'door', 'window', 'bookshelf',
                      'picture', 'counter', 'blinds', 'desk', 'shelves',
                      'curtain', 'dresser', 'pillow', 'mirror', 'floor mat',
                      'clothes', 'ceiling', 'books', 'refridgerator', 'television',
                      'paper', 'towel', 'shower curtain', 'box', 'whiteboard',
                      'person', 'night stand', 'toilet', 'sink', 'lamp',
                      'bathtub', 'bag', 'otherstructure', 'otherfurniture', 'otherprop']


def calculate_class_weights(dataset):
    num_classes = len(NYU_CATEGORY_NAMES) + 1      # 41
    total_count = 0

    class_weight = torch.zeros(num_classes)
    class_counts = torch.zeros(num_classes)
    for i in range(len(dataset)):
        # image, semseg均为 tensor类型，且为进行transform
        image, label = dataset[i]['image'], dataset[i]['semseg']
        label_count = torch.bincount(label.view(-1), minlength=num_classes)
        class_counts += label_count
		# 记录除背景类以外的语义类像素总数
        total_count += torch.nonzero(label.view(-1)).size(0)

        if i % 50 == 0:
            print('current image index is:', i)

    total_count = torch.tensor(total_count, dtype=torch.float)
	# 在nyud数据集中背景类的权重是我们不需要的，去除即可
    class_weight = torch.div((total_count - class_counts.float()), total_count)[1:]

    return class_weight


# 计算结果保存为一个数组，分别对应40个语义类别
CLASS_WEIGHT = [0.7613, 0.8796, 0.9183, 0.9541, 0.9511, 0.9689, 0.9724, 0.9776, 0.9764,
                0.9774, 0.9785, 0.9821, 0.9824, 0.9868, 0.9849, 0.9894, 0.9882, 0.9889,
                0.9887, 0.9917, 0.9892, 0.9905, 0.9932, 0.9937, 0.9930, 0.9950, 0.9958,
                0.9955, 0.9946, 0.9967, 0.9963, 0.9966, 0.9961, 0.9967, 0.9965, 0.9971,
                0.9972, 0.9736, 0.9753, 0.9387]

使用时将数组转化为一维tensor格式，然后传入torch.nn.CrossEntorpyloss()即可。

总结

berHu损失对单目深度估计指标有一定提升，损失值下降曲线与L1_loss比起来也相对更稳定一点。nyud数据集语义像素数量分布相对较为均衡，带权交叉熵损失函数并没有取到预期的效果。

如有错误，欢迎指正