File size: 1,753 Bytes
3b96cb1 |
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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import torch
import torch.nn as nn
from opencd.registry import MODELS
def bcl_loss(
pred,
target,
margin=2.0,
eps=1e-4,
ignore_index=255,
**kwargs):
pred = pred.squeeze()
target = target.squeeze()
assert pred.size() == target.size() and target.numel() > 0
mask = (target != ignore_index).float()
target = target * mask
utarget = 1 - target
n_u = utarget.sum() + eps
n_c = target.sum() + eps
loss = torch.sum(utarget * torch.pow(pred, 2) * mask) / n_u + \
torch.sum(target * torch.pow(torch.clamp(margin - pred, min=0.), 2)) / n_c
return loss
@MODELS.register_module()
class BCLLoss(nn.Module):
"""Batch-balanced Contrastive Loss"""
def __init__(
self,
margin=2.0,
loss_weight=1.0,
ignore_index=255,
loss_name='bcl_loss',
**kwargs):
super().__init__()
self.margin = margin
self.loss_weight = loss_weight
self.ignore_index = ignore_index
self._loss_name = loss_name
def forward(self,
pred,
target,
**kwargs):
loss = self.loss_weight * bcl_loss(
pred, target, self.margin, self.ignore_index)
return loss
@property
def loss_name(self):
"""Loss Name.
This function must be implemented and will return the name of this
loss function. This name will be used to combine different loss items
by simple sum operation. In addition, if you want this loss item to be
included into the backward graph, `loss_` must be the prefix of the
name.
Returns:
str: The name of this loss item.
"""
return self._loss_name |