Spaces:
Runtime error
Runtime error
File size: 3,921 Bytes
412c852 |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional, Union
import torch
import torch.nn as nn
from torch import Tensor
from mmseg.registry import MODELS
from .utils import weight_reduce_loss
def silog_loss(pred: Tensor,
target: Tensor,
weight: Optional[Tensor] = None,
eps: float = 1e-4,
reduction: Union[str, None] = 'mean',
avg_factor: Optional[int] = None) -> Tensor:
"""Computes the Scale-Invariant Logarithmic (SI-Log) loss between
prediction and target.
Args:
pred (Tensor): Predicted output.
target (Tensor): Ground truth.
weight (Optional[Tensor]): Optional weight to apply on the loss.
eps (float): Epsilon value to avoid division and log(0).
reduction (Union[str, None]): Specifies the reduction to apply to the
output: 'mean', 'sum' or None.
avg_factor (Optional[int]): Optional average factor for the loss.
Returns:
Tensor: The calculated SI-Log loss.
"""
pred, target = pred.flatten(1), target.flatten(1)
valid_mask = (target > eps).detach().float()
diff_log = torch.log(target.clamp(min=eps)) - torch.log(
pred.clamp(min=eps))
valid_mask = (target > eps).detach() & (~torch.isnan(diff_log))
diff_log[~valid_mask] = 0.0
valid_mask = valid_mask.float()
diff_log_sq_mean = (diff_log.pow(2) * valid_mask).sum(
dim=1) / valid_mask.sum(dim=1).clamp(min=eps)
diff_log_mean = (diff_log * valid_mask).sum(dim=1) / valid_mask.sum(
dim=1).clamp(min=eps)
loss = torch.sqrt(diff_log_sq_mean - 0.5 * diff_log_mean.pow(2))
if weight is not None:
weight = weight.float()
loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
return loss
@MODELS.register_module()
class SiLogLoss(nn.Module):
"""Compute SiLog loss.
Args:
reduction (str, optional): The method used
to reduce the loss. Options are "none",
"mean" and "sum". Defaults to 'mean'.
loss_weight (float, optional): Weight of loss. Defaults to 1.0.
eps (float): Avoid dividing by zero. Defaults to 1e-3.
loss_name (str, optional): Name of the loss item. If you want this
loss item to be included into the backward graph, `loss_` must
be the prefix of the name. Defaults to 'loss_silog'.
"""
def __init__(self,
reduction='mean',
loss_weight=1.0,
eps=1e-6,
loss_name='loss_silog'):
super().__init__()
self.reduction = reduction
self.loss_weight = loss_weight
self.eps = eps
self._loss_name = loss_name
def forward(
self,
pred,
target,
weight=None,
avg_factor=None,
reduction_override=None,
):
assert pred.shape == target.shape, 'the shapes of pred ' \
f'({pred.shape}) and target ({target.shape}) are mismatch'
assert reduction_override in (None, 'none', 'mean', 'sum')
reduction = (
reduction_override if reduction_override else self.reduction)
loss = self.loss_weight * silog_loss(
pred,
target,
weight,
eps=self.eps,
reduction=reduction,
avg_factor=avg_factor,
)
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
|