Update center used for teacher output.
(self, teacher_output, it)
| 457 | |
| 458 | @torch.no_grad() |
| 459 | def update_center(self, teacher_output, it): |
| 460 | """ |
| 461 | Update center used for teacher output. |
| 462 | """ |
| 463 | batch_center = torch.sum(teacher_output, dim=0, keepdim=True) |
| 464 | dist.all_reduce(batch_center) |
| 465 | batch_center = batch_center / (len(teacher_output) * dist.get_world_size()) |
| 466 | |
| 467 | # ema update |
| 468 | self.center = self.center * self.center_momentum + batch_center * (1 - self.center_momentum) |
| 469 | |
| 470 | @torch.no_grad() |
| 471 | def update_patch_center(self, teacher_output, it): |