| 398 | |
| 399 | |
| 400 | class DINOLoss(nn.Module): |
| 401 | def __init__(self, out_dim, out_dim_selfpatch, ncrops, warmup_teacher_temp, teacher_temp, |
| 402 | warmup_teacher_temp_epochs, nepochs, student_temp=0.1, |
| 403 | center_momentum=0.9): |
| 404 | super().__init__() |
| 405 | self.student_temp = student_temp |
| 406 | self.center_momentum = center_momentum |
| 407 | self.ncrops = ncrops |
| 408 | self.register_buffer("center", torch.zeros(1, 1, out_dim)) |
| 409 | self.register_buffer("patch_center", torch.zeros(1, out_dim_selfpatch)) |
| 410 | |
| 411 | # we apply a warm up for the teacher temperature because |
| 412 | # a too high temperature makes the training instable at the beginning |
| 413 | self.teacher_temp_schedule = np.concatenate(( |
| 414 | np.linspace(warmup_teacher_temp, |
| 415 | teacher_temp, warmup_teacher_temp_epochs), |
| 416 | np.ones(nepochs - warmup_teacher_temp_epochs) * teacher_temp |
| 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 | |