BCE损失解析
BCE完整代码分析
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as npclass BKDLoss(nn.Module):def __init__(self, cls_num_list, T, weight=None):super(BKDLoss, self).__init__()self.T = Tself.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.weightprint(f"Weighted Teacher Model Output (加权的教师模型输出): {pred_t}")# 归一化处理,使概率和为1pred_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.Tprint(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_lossprint(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 张量,以便进行广播,使得我们可以在元素级别上进行逐行相除的操作。它通过增加一个维度来确保张量形状匹配。