File size: 624 Bytes
482ab8a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
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()
|