经验笔记：使用 PyTorch 计算多分类问题中Dice Loss 的正确方法

经验笔记：使用 PyTorch 计算多分类问题中Dice Loss 的正确方法

概述

Dice Loss 是一种广泛应用于图像分割任务中的损失函数，它基于 Dice 系数（也称为 F1-score），用于衡量预测结果与真实标签之间的相似度。在 PyTorch 中，计算 Dice Loss 时，我们需要注意如何正确地对张量求和以保留类别信息，从而准确评估模型性能。

输入张量的结构

对于一个二元分类或多类分类问题，输入张量通常具有 [batch_size, num_class] 的形状，其中 batch_size 表示批次大小，num_class 表示类别数量。

正确的 Dice Loss 计算方法

使用 .sum(dim=0) 按类别求和

这是推荐的方法，因为它保持了类别维度，允许我们分别计算每个类别的 Dice Loss，并最终取平均值来获得整体损失。以下是详细的步骤：

1. 初始化平滑因子：

   smooth = 1e-5  # 防止除以零
   

2. 计算交集和并集：

   intersection = (predict * true).sum(dim=0)  # 按类别求和
   union = predict.sum(dim=0) + true.sum(dim=0)  # 按类别求和
   

3. 计算 Dice 系数：

   dice_score = (2. * intersection + smooth) / (union + smooth)
   

4. 计算 Dice Loss：

   dice_loss = 1 - dice_score
   

5. 计算平均 Dice Loss：

   mean_dice_loss = dice_loss.mean()
   

示例代码

import torch

# 模拟的真实标签和预测标签
true = torch.tensor([
    [1, 0, 1, 0],  # 样本1
    [0, 1, 0, 1],  # 样本2
    [1, 1, 0, 0]   # 样本3
], dtype=torch.float32)  # shape: (3, 4)

predict = torch.tensor([
    [1, 1, 0, 0],  # 样本1
    [0, 1, 1, 0],  # 样本2
    [1, 0, 0, 1]   # 样本3
], dtype=torch.float32)  # shape: (3, 4)

smooth = 0 # 1e-5
# 本应使用smooth来防止除以零， 但是这里为了演示计算过程，将smooth设置为0来简化结果

# 计算交集和并集
intersection = (predict * true).sum(dim=0)  # 按类别求和
# (predict * true) = torch.tensor([
#     [1, 0, 0, 0],  # 样本1
#     [0, 1, 0, 0],  # 样本2
#     [1, 0, 0, 0]   # 样本3
# ], dtype=torch.float32)  # shape: (3, 4)

# intersection = torch.tensor([
#     [2, 1, 0, 0]
# ], dtype=torch.float32)  # shape: (4)

union = predict.sum(dim=0) + true.sum(dim=0)  # 按类别求和
# predict.sum(dim=0) = torch.tensor([
#     [2, 2, 1, 1]
# ], dtype=torch.float32)  # shape: (4)

# true.sum(dim=0) = torch.tensor([
#     [2, 2, 1, 1]
# ], dtype=torch.float32)  # shape: (4)

# union = torch.tensor([
#     [4, 4, 2, 2]
# ], dtype=torch.float32)  # shape: (4)

# Dice系数计算公式
dice_score = (2. * intersection + smooth) / (union + smooth)
#  (2. * intersection + smooth) = torch.tensor([
#     [4, 2, 0, 0]
# ], dtype=torch.float32)  # shape: (4)

#  (union + smooth) = torch.tensor([
#     [4, 4, 2, 2]
# ], dtype=torch.float32)  # shape: (4)

#  dice_score = torch.tensor([
#     [1.0000, 0.5000, 0.0000, 0.0000]
# ], dtype=torch.float32)  # shape: (4)


# Dice损失等于1减去Dice系数
dice_loss = 1 - dice_score
#  dice_loss = torch.tensor([
#     [0.0000, 0.5000, 1.0000, 1.0000]
# ], dtype=torch.float32)  # shape: (4)


# 计算平均Dice Loss
mean_dice_loss = dice_loss.mean()
#  mean_dice_loss = torch.tensor(
#     0.6250
# , dtype=torch.float32)  # shape: (1)

print("Mean Dice Loss:", mean_dice_loss.item())

不推荐的方法（对于输入张量为[batch_size, num_class] 来说错误的方法）

使用 .sum(dim=(0, 1)) 全局求和

这种方法将所有类别和样本的信息合并成一个标量值，失去了类别层面的细节，因此不推荐用于 Dice Loss 的计算。虽然可以得到一个整体的 Dice Loss，但这无法区分模型在不同类别上的性能差异。

示例代码

import torch

# 模拟的真实标签和预测标签
true = torch.tensor([
    [1, 0, 1, 0],  # 样本1
    [0, 1, 0, 1],  # 样本2
    [1, 1, 0, 0]   # 样本3
], dtype=torch.float32)  # shape: (3, 4)

predict = torch.tensor([
    [1, 1, 0, 0],  # 样本1
    [0, 1, 1, 0],  # 样本2
    [1, 0, 0, 1]   # 样本3
], dtype=torch.float32)  # shape: (3, 4)

smooth = 0 # 1e-5
# 本应使用smooth来防止除以零， 但是这里为了演示计算过程，将smooth设置为0来简化结果

# 全局计算交集和并集
total_intersection = (predict * true).sum(dim=(0, 1))  # 所有元素求和
# (predict * true) = torch.tensor([
#     [1, 0, 0, 0],  # 样本1
#     [0, 1, 0, 0],  # 样本2
#     [1, 0, 0, 0]   # 样本3
# ], dtype=torch.float32)  # shape: (3, 4)

# total_intersection = torch.tensor(
#     3
# , dtype=torch.float32)  # shape: (1)

total_union = predict.sum(dim=(0, 1)) + true.sum(dim=(0, 1))  # 所有元素求和
# predict.sum(dim=(0, 1)) = torch.tensor(
#     6
# , dtype=torch.float32)  # shape: (1)

# true.sum(dim=(0, 1)) = torch.tensor(
#     6
# , dtype=torch.float32)  # shape: (1)

# total_union = torch.tensor(
#     12
# , dtype=torch.float32)  # shape: (1)

# 全局Dice系数计算公式
global_dice_score = (2. * total_intersection + smooth) / (total_union + smooth)
# (2. * total_intersection + smooth) = torch.tensor(
#     6
# , dtype=torch.float32)  # shape: (1)

# (total_union + smooth) = torch.tensor(
#     12
# , dtype=torch.float32)  # shape: (1)

# global_dice_score = torch.tensor(
#     0.5000
# , dtype=torch.float32)  # shape: (1)


# 全局Dice损失等于1减去全局Dice系数
global_dice_loss = 1 - global_dice_score
# global_dice_loss = torch.tensor(
#     0.5000
# , dtype=torch.float32)  # shape: (1)

print("Global Dice Loss:", global_dice_loss.item())

总结

使用.sum(dim=0)最终得出的Dice_Loss为0.6250，
而使用.sum(dim=0)最终得出的Dice_Loss为0.5000。
可见，对于输入张量为[batch_size, num_class] 来说，.sum(dim=0)与.sum(dim=(0, 1))是两种不同的算法，并且.sum(dim=0)才是保留种类区别的正确计算方法。

• .sum(dim=0)：保持类别维度，分别计算每个类别的 Dice Loss 并取平均值，能够更准确地反映模型在不同类别上的表现。
• .sum(dim=(0, 1))：将所有类别和样本的信息合并成一个标量值，失去类别层面的细节，2个维度的计算一般用于计算图像分割的二维像素点，并且对于图像分割任务也常常会对每个类别分别计算 Dice Loss 并取平均，以确保模型在所有类别上的良好泛化能力。

通过正确使用 .sum(dim=0) 来计算 Dice Loss，我们可以更好地评估模型在图像分割任务中的性能，并根据各个类别的表现进行针对性优化。