Mountchicken's picture
Upload 704 files
9bf4bd7
raw
history blame
3.43 kB
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional
import torch
import torch.nn as nn
from mmocr.registry import MODELS
@MODELS.register_module()
class MaskedDiceLoss(nn.Module):
"""Masked dice loss.
Args:
eps (float, optional): Eps to avoid zero-divison error. Defaults to
1e-6.
"""
def __init__(self, eps: float = 1e-6) -> None:
super().__init__()
assert isinstance(eps, float)
self.eps = eps
def forward(self,
pred: torch.Tensor,
gt: torch.Tensor,
mask: Optional[torch.Tensor] = None) -> torch.Tensor:
"""Forward function.
Args:
pred (torch.Tensor): The prediction in any shape.
gt (torch.Tensor): The learning target of the prediction in the
same shape as pred.
mask (torch.Tensor, optional): Binary mask in the same shape of
pred, indicating positive regions to calculate the loss. Whole
region will be taken into account if not provided. Defaults to
None.
Returns:
torch.Tensor: The loss value.
"""
assert pred.size() == gt.size() and gt.numel() > 0
if mask is None:
mask = torch.ones_like(gt)
assert mask.size() == gt.size()
pred = pred.contiguous().view(pred.size(0), -1)
gt = gt.contiguous().view(gt.size(0), -1)
mask = mask.contiguous().view(mask.size(0), -1)
pred = pred * mask
gt = gt * mask
dice_coeff = (2 * (pred * gt).sum()) / (
pred.sum() + gt.sum() + self.eps)
return 1 - dice_coeff
@MODELS.register_module()
class MaskedSquareDiceLoss(nn.Module):
"""Masked square dice loss.
Args:
eps (float, optional): Eps to avoid zero-divison error. Defaults to
1e-3.
"""
def __init__(self, eps: float = 1e-3) -> None:
super().__init__()
assert isinstance(eps, float)
self.eps = eps
def forward(self,
pred: torch.Tensor,
gt: torch.Tensor,
mask: Optional[torch.Tensor] = None) -> torch.Tensor:
"""Forward function.
Args:
pred (torch.Tensor): The prediction in any shape.
gt (torch.Tensor): The learning target of the prediction in the
same shape as pred.
mask (torch.Tensor, optional): Binary mask in the same shape of
pred, indicating positive regions to calculate the loss. Whole
region will be taken into account if not provided. Defaults to
None.
Returns:
torch.Tensor: The loss value.
"""
assert pred.size() == gt.size() and gt.numel() > 0
if mask is None:
mask = torch.ones_like(gt)
assert mask.size() == gt.size()
batch_size = pred.size(0)
pred = pred.contiguous().view(batch_size, -1)
gt = gt.contiguous().view(batch_size, -1).float()
mask = mask.contiguous().view(batch_size, -1).float()
pred = pred * mask
gt = gt * mask
a = torch.sum(pred * gt, dim=1)
b = torch.sum(pred * pred, dim=1) + self.eps
c = torch.sum(gt * gt, dim=1) + self.eps
d = (2 * a) / (b + c)
loss = 1 - d
loss = torch.mean(loss)
return loss