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
注意必须要有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)