PyTorch 自定义损失函数解析：Loss 模块详解及输出形状分析

这段代码定义了一个名为 Loss 的自定义模块，用于计算深度学习模型的损失值。它接受两个输入参数 output 和 target，并根据 self.mode 选择不同的指标（metric）来计算损失。

如果 self.mode 为 'Dice'，则使用 DSC 类的实例作为 metric；如果 self.mode 为 'CE'，则使用 CE 类的实例作为 metric；否则，打印出错误信息并使用 DSC 类的实例作为 metric。

class Loss(nn.Module):
    def __init__(self, mode):
        super(Loss, self).__init__()
        self.mode = mode

    def forward(self, output, target):
        if self.mode == 'Dice':
            metric = DSC()
        elif self.mode == 'CE':
            metric = CE()
        else:
            print('The chosen loss function is not provided!!!! We use Dice Loss instead!!!')
            metric = DSC()

        batchSize = target.size(0)
        classNum = target.size(1)
        loss = torch.empty(batchSize, classNum)
        for ii in range(classNum):
            loss[:, ii] = metric(output[:, ii], target[:, ii])
        return loss

接下来，代码获取了 target 的 batch 大小 (batchSize) 和类别数量 (classNum)，并创建了一个大小为 (batchSize, classNum) 的空张量 loss。

然后，通过 for 循环遍历类别数量，并将 metric 计算得到的损失值存储在 loss 张量的相应位置。

最后，返回 loss 张量作为输出。

输出形状分析：

根据代码逻辑和给定的输入参数，最终的输出的 shape 将是 (batchSize, classNum)，其中 batchSize 是输入 target 的第一个维度的大小，classNum 是输入 target 的第二个维度的大小。也就是说，输出的张量将具有与 target 相同的 batch 大小和类别数量，每个位置存储了对应类别的损失值。