如何自定义Dice Loss函数来替代pytorch中的nll_loss函数？

Question

我假设预测结果是

pred

，对应的标签变量是

label_face

。因为变量

label_face

在分割问题中包含了大量的数据不平衡。因此，我想用

Dice Loss

函数来代替PyTorch中的

nll_loss

。

pred = tensor([[-0.6813, -0.7052],
              [-0.6467, -0.7419], 
              [-0.7436, -0.6451],
              ...,
              [-0.5635, -0.8421],
              [-0.6089, -0.7852],
              [-0.7449, -0.6439]], device='cuda:0', grad_fn=<ViewBackward0>)

pred.shape --> torch.Size([7862, 2])

label_face = tensor([1, 1, 1,  ..., 1, 1, 1], device='cuda:0')
label_face.shape --> torch.Size([7862])


loss = F.nll_loss(pred, label_face)        
loss.backward()

我尝试了二进制分割函数的自定义 Dice Loss，如下所示：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.cuda.amp as amp


## Soft Dice Loss for binary segmentation
## pytorch autograd

class SoftDiceLoss(nn.Module):
    '''
    soft-dice loss, useful in binary segmentation
    '''
    def __init__(self,
                 p=1,
                 smooth=1):
        super(SoftDiceLossV1, self).__init__()
        self.p = p
        self.smooth = smooth

    def forward(self, logits, labels):
        '''
        inputs:
            logits: tensor of shape (N, H, W, ...)
            label: tensor of shape(N, H, W, ...)
        output:
            loss: tensor of shape(1, )
        '''
        probs = torch.sigmoid(logits)
        numer = (probs * labels).sum()
        denor = (probs.pow(self.p) + labels.pow(self.p)).sum()
        loss = 1. - (2 * numer + self.smooth) / (denor + self.smooth)
        return loss

custom_loss = SoftDiceLoss()
loss = custom_loss(pred, label_face)        
loss.backward()

我真的试过了但是做不到或者不知道错误在哪里。我是一个初学者，希望你能帮我写代码来自定义这个损失函数。任何评论都非常感谢。非常感谢。