mmsegmentation 添加L1Loss

mmseg/models/losses/模块中添加L1Loss定义：

import torch
import torch.nn as nn
import torch.nn.functional as F
from ..builder import LOSSES


@LOSSES.register_module()
class L1Loss(nn.Module):
    # TODO: weight
    def __init__(self, loss_name='loss_l1', **kwargs):
        super(L1Loss, self).__init__()
        self._loss_name = loss_name

    def forward(self, pred, target, weight=None, ignore_index=None):
   		# pred: (n,c,h,w)   target: (n,h,w)
        classes = pred.shape[1]
        size = list(target.shape)
        size.append(classes)  # (n,h,w,c)
        target_one_hot = target.view(-1)  # (n*h*w)
        ones = torch.sparse.torch.eye(classes).to(target_one_hot.device)
        ones = ones.index_select(0, target_one_hot)  # (n*h*w, classes)
        ones = ones.view(*size)  # (n,h,w,c)
        target_one_hot = ones.permute(0, 3, 1, 2)  # (n,c,h,w)
        loss = nn.L1Loss()(pred, target_one_hot)
        return loss

	@property
    def loss_name(self):
        """Loss Name.

        This function must be implemented and will return the name of this
        loss function. This name will be used to combine different loss items
        by simple sum operation. In addition, if you want this loss item to be
        included into the backward graph, `loss_` must be the prefix of the
        name.

        Returns:
            str: The name of this loss item.
        """
        return self._loss_name
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

注意必须要有loss_name方法，并且返回的loss_name需要以loss_作为前缀。

传入的pred和target的shape不一致，需要转为一致才可以直接调用nn.L1Loss()方法。
pred.shape: (n,c,h,w)
target.shape: (n,h,w)
所以需要将target转one-hot。转one-hot方法：index_select。
(n,h,w) => (n,h,w,c) => (n,c,h,w)

拓展阅读：Pytorch中，将label变成one hot编码的两种方式