BCE损失解析

 BCE完整代码分析

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class BKDLoss(nn.Module):
    def __init__(self, cls_num_list, T, weight=None):
        super(BKDLoss, self).__init__()
        self.T = T
        self.weight = weight
        # 计算类别的频率，并打印出来
        self.class_freq = torch.cuda.FloatTensor(cls_num_list / np.sum(cls_num_list))
        print(f"Class Frequencies (归一化类别频率): {self.class_freq}")
        
        # 初始化交叉熵损失函数
        self.CELoss = nn.CrossEntropyLoss().cuda()

    def forward(self, out_s, out_t, target, alpha):
        # 教师模型输出的 softmax，并打印结果
        pred_t = F.softmax(out_t / self.T, dim=1)
        print(f"Teacher Model Softmax Output (教师模型的 softmax 输出): {pred_t}")
        
        # 如果有 weight，加权教师模型的输出，并归一化
        if self.weight is not None:
            pred_t = pred_t * self.weight
            print(f"Weighted Teacher Model Output (加权的教师模型输出): {pred_t}")
            
            # 归一化处理，使概率和为1
            pred_t = pred_t / pred_t.sum(1)[:, None]
            print(f"Normalized Weighted Teacher Model Output (归一化后的加权教师模型输出): {pred_t}")
        
        # 学生模型的 KL 散度损失，并打印结果
        kd = F.kl_div(F.log_softmax(out_s / self.T, dim=1), pred_t, reduction='none').mean(dim=0)
        print(f"KL Divergence Loss (KL 散度损失): {kd}")
        
        # 计算缩放后的蒸馏损失
        kd_loss = kd.sum() * self.T * self.T
        print(f"Scaled KD Loss (缩放后的 KD 损失): {kd_loss}")
        
        # 计算交叉熵损失，并打印
        ce_loss = self.CELoss(out_s, target)
        print(f"Cross Entropy Loss (交叉熵损失): {ce_loss}")
        
        # 计算总损失，并打印
        loss = alpha * kd_loss + ce_loss
        print(f"Total Loss (总损失): {loss}")
        
        return loss, kd


# 示例输入
cls_num_list = np.array([100, 200, 300, 400])  # 每个类别的样本数
T = 2  # 温度
weight = torch.FloatTensor([1, 1.5, 2, 2.5]).cuda()  # 类别权重

# 初始化 BKDLoss
bkd_loss_fn = BKDLoss(cls_num_list, T, weight)

# 假设学生模型和教师模型的输出 (batch_size=3, num_classes=4)
out_s = torch.randn(3, 4).cuda()  # 学生模型输出
out_t = torch.randn(3, 4).cuda()  # 教师模型输出
target = torch.tensor([0, 1, 2], dtype=torch.long).cuda()  # 真实标签
alpha = 0.5  # 调节蒸馏损失与交叉熵损失的系数

# 计算损失
loss, kd = bkd_loss_fn(out_s, out_t, target, alpha)

pred_t = pred_t / pred_t.sum(1)[:, None]

[:, None] 的作用是通过增加一个维度来改变张量的形状。下面详细解释这个操作的含义。

1. pred_t.sum(1)

pred_t.sum(1) 是对 pred_t 张量的第 1 维（即每一行）进行求和操作。假设 pred_t 的形状是 (batch_size, num_classes)，那么 pred_t.sum(1) 会对每个样本（即每行）的所有类别的预测概率进行求和，生成一个形状为 (batch_size,) 的向量，其中每个元素是对应行的所有类别概率的和。

2. [:, None] 的作用

[:, None] 是一种利用切片操作增加维度的方式。在 pred_t.sum(1) 后得到的向量形状是 (batch_size,)，这是一个 1D 张量。通过 [:, None]，我们在它的末尾增加了一个维度，使其变为 (batch_size, 1)。

具体地，[:, None] 等价于 unsqueeze(1)，它将 1D 张量扩展为 2D 张量。

3. 为什么需要增加维度？

在 pred_t = pred_t / pred_t.sum(1)[:, None] 这行代码中，我们需要对 pred_t 的每一行除以其所有元素的和。由于 pred_t 是一个形状为 (batch_size, num_classes) 的 2D 张量，而 pred_t.sum(1) 是一个形状为 (batch_size,) 的 1D 张量，无法直接进行逐元素相除。

通过 [:, None] 将 pred_t.sum(1) 扩展为形状为 (batch_size, 1) 的 2D 张量，这样就可以广播（broadcast）到 pred_t 的每一行上，从而实现逐行相除的操作。

举个例子：

假设 pred_t 是一个形状为 (2, 4) 的张量（即有两个样本，每个样本有四个类别的预测概率）：

pred_t = torch.tensor([[0.1, 0.2, 0.3, 0.4],
                       [0.3, 0.3, 0.2, 0.2]])

pred_t.sum(1) 的结果是：

tensor([1.0, 1.0])  # 对每一行的概率求和

通过 [:, None] 操作，结果变成：

tensor([[1.0],
        [1.0]])  # 扩展为 (batch_size, 1) 形状

现在可以将 pred_t 的每一行分别除以 1.0，保持维度一致，执行按元素相除操作。

总结：

[:, None] 的作用是将 1D 张量转换为 2D 张量，以便进行广播，使得我们可以在元素级别上进行逐行相除的操作。它通过增加一个维度来确保张量形状匹配。