Spaces:
Running
Running
| # Copyright (C) 2022-present Naver Corporation. All rights reserved. | |
| # Licensed under CC BY-NC-SA 4.0 (non-commercial use only). | |
| # | |
| # -------------------------------------------------------- | |
| # Criterion to train CroCo | |
| # -------------------------------------------------------- | |
| # References: | |
| # MAE: https://github.com/facebookresearch/mae | |
| # -------------------------------------------------------- | |
| import torch | |
| class MaskedMSE(torch.nn.Module): | |
| def __init__(self, norm_pix_loss=False, masked=True): | |
| """ | |
| norm_pix_loss: normalize each patch by their pixel mean and variance | |
| masked: compute loss over the masked patches only | |
| """ | |
| super().__init__() | |
| self.norm_pix_loss = norm_pix_loss | |
| self.masked = masked | |
| def forward(self, pred, mask, target): | |
| if self.norm_pix_loss: | |
| mean = target.mean(dim=-1, keepdim=True) | |
| var = target.var(dim=-1, keepdim=True) | |
| target = (target - mean) / (var + 1.e-6)**.5 | |
| loss = (pred - target) ** 2 | |
| loss = loss.mean(dim=-1) # [N, L], mean loss per patch | |
| if self.masked: | |
| loss = (loss * mask).sum() / mask.sum() # mean loss on masked patches | |
| else: | |
| loss = loss.mean() # mean loss | |
| return loss | |