【mmsegmentation】Loss模块（进阶）自定义自己的LOSS

1、定义自己的loss

driving\models\losses\shuai_loss.py

import torch
from torch import nn
from mmseg.models import LOSSES@LOSSES.register_module()
class ShuaiLoss(nn.Module):def __init__(self,loss_weight=1.0):super().__init__()self.ce_loss = nn.CrossEntropyLoss()self.loss_weight = loss_weightdef forward(self,input,target,device_id='cpu',sample_ratio=1.0):loss = {}if len(target)==0:loss["cls_cost"] = torch.tensor(0.0,dtype=torch.float32,device=device_id)else:loss["cls_cost"] = self.ce_loss(input,target)loss["total_road_cls_loss"] = loss["cls_cost"] * self.loss_weight * sample_ratio # + other losses, if havereturn loss

看下LOSSES注册表(@LOSSES.register_module())

• 可以看到ShuaiLoss可以被注册到LOSSES中
• 其实，这里的LOSSES是BACKBONES NECKS HEADS LOSSES SEGMENTORS的总和

2、调用Shuai_loss

if __name__ == "__main__":print("call shuai_loss:")from mmseg.models import build_loss# 1.配置 dictloss = dict(type='ShuaiLoss',loss_weight=1.0,loss_name='loss_shuai')# 从注册器中构建shuai_loss = build_loss(loss)# 使用shuai losspred = torch.Tensor([[0, 2, 3, 0], [0,2,3,0]])   # [2,4]target = torch.Tensor([[1, 1, 1, 0], [1,1,1,1]]) # [2,4]loss = shuai_loss(pred, target)print("loss:",loss)