| import torch | |
| import torch.nn as nn | |
| def get_map_mask_loss(opt): | |
| return MapMaskLoss() | |
| class MapMaskLoss(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.bce_loss = nn.BCELoss(reduction="mean") | |
| def forward(self, out_map, mask): | |
| mask_size = mask.shape[-2:] | |
| if out_map.shape[-2:] != mask_size: | |
| out_map = nn.functional.interpolate( | |
| out_map, size=mask_size, mode="bilinear", align_corners=False | |
| ) | |
| loss = self.bce_loss(out_map, mask) | |
| return {"loss": loss} | |
| if __name__ == "__main__": | |
| map_mask_loss = MapMaskLoss() | |