【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_weight
    def 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 have

        return 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.配置 dict
    loss = dict(type='ShuaiLoss',
                 loss_weight=1.0,
                 loss_name='loss_shuai')
    # 从注册器中构建
    shuai_loss = build_loss(loss)

    # 使用shuai loss
    pred = 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)