Cross-entropy between softmax outputs of the teacher and student networks.
(self, teacher, student, student_output, teacher_output, epoch, it)
| 417 | )) |
| 418 | |
| 419 | def forward(self, teacher, student, student_output, teacher_output, epoch, it): |
| 420 | """ |
| 421 | Cross-entropy between softmax outputs of the teacher and student networks. |
| 422 | """ |
| 423 | # teacher centering and sharpening |
| 424 | student_cls = student_output[0][0].chunk(2) + student_output[1][0].chunk(self.ncrops-2) |
| 425 | student_loc = student_output[0][1].chunk(2) + student_output[1][1].chunk(self.ncrops-2) |
| 426 | |
| 427 | teacher_cls = teacher_output[0][0].chunk(2) + teacher_output[1][0].chunk(self.ncrops-2) |
| 428 | teacher_loc = teacher_output[0][1].chunk(2) + teacher_output[1][1].chunk(self.ncrops-2) |
| 429 | temp = self.teacher_temp_schedule[epoch] |
| 430 | |
| 431 | c_loss = 0 |
| 432 | p_loss = 0 |
| 433 | n_loss_terms = 0 |
| 434 | m_loss_terms = 0 |
| 435 | assert len(teacher_cls) == self.ncrops |
| 436 | for iq in range(len(teacher_cls)): |
| 437 | q_cls = F.softmax((teacher_cls[iq] - self.center)/ temp, dim=-1).detach() |
| 438 | for v in range(self.ncrops): |
| 439 | if v == iq: |
| 440 | q_pat = F.softmax((teacher_loc[iq] - self.patch_center)/ temp, dim=-1).detach() |
| 441 | p_pat = student_loc[v] |
| 442 | patch_loss = torch.sum(-q_pat * F.log_softmax(p_pat / self.student_temp, dim=-1), dim=-1) |
| 443 | p_loss += patch_loss.mean() |
| 444 | m_loss_terms += 1 |
| 445 | else: |
| 446 | if iq > 1: |
| 447 | continue |
| 448 | cls_loss = torch.sum(-q_cls * F.log_softmax(student_cls[v] / self.student_temp, dim=-1), dim=-1) |
| 449 | c_loss += cls_loss.mean() |
| 450 | n_loss_terms += 1 |
| 451 | c_loss /= n_loss_terms |
| 452 | p_loss /= m_loss_terms |
| 453 | |
| 454 | self.update_center(torch.cat(teacher_cls), it) |
| 455 | self.update_patch_center(teacher_loc, it) |
| 456 | return (c_loss + p_loss*0.1), c_loss.item(), p_loss.item() |
| 457 | |
| 458 | @torch.no_grad() |
| 459 | def update_center(self, teacher_output, it): |
nothing calls this directly
no test coverage detected