Confusion about the kd_ce_loss

[fengjiejiejiejie]

Hi, thanks for your impressive paper and codes, but I'm confused about the kd_ce_loss:

def FNKD(self, student_outputs, teacher_outputs, student_feature, teacher_feature): student_L2norm = torch.norm(student_feature) teacher_L2norm = torch.norm(teacher_feature) q_fn = F.log_softmax(teacher_outputs / teacher_L2norm, dim=1) to_kd = F.softmax(student_outputs / student_L2norm, dim=1) KD_ce_loss = self.ce( q_fn, to_kd[:, 0].long()) return KD_ce_loss

Why use self.ce after the softmax/log_softmax? Is to_kd[:, 0] just using the first channel of the student_outputs (background)? Maybe the following is right?

q_fn = F.log_softmax(teacher_outputs / T, dim=1) to_kd = F.softmax(student_outputs / T, dim=1) KD_ce_loss = -torch.mean(torch.sum(to_kd * q_fn, dim=1))

best, fj