| 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 | """ |