《解析 MMDetection 中的损失函数辅助工具》
在 MMDetection 中,损失函数的计算通常需要一些辅助工具来进行灵活的处理。本文将深入解析 mmdet/models/losses/utils.py
文件中的函数。
一、reduce_loss
函数
def reduce_loss(loss: Tensor, reduction: str) -> Tensor:
"""Reduce loss as specified.
Args:
loss (Tensor): Elementwise loss tensor.
reduction (str): Options are "none", "mean" and "sum".
Return:
Tensor: Reduced loss tensor.
"""
reduction_enum = F._Reduction.get_enum(reduction)
# none: 0, elementwise_mean:1, sum: 2
if reduction_enum == 0:
return loss
elif reduction_enum == 1:
return loss.mean()
elif reduction_enum == 2:
return loss.sum()
这个函数用于根据指定的reduction
策略来减少损失张量。它接受一个损失张量和一个表示减少策略的字符串参数。通过使用 PyTorch 的F._Reduction.get_enum
方法,将字符串转换为对应的枚举值。如果reduction
是"none",则直接返回原始损失张量;如果是"mean",则计算损失张量的平均值;如果是"sum",则计算损失张量的总和。
二、weight_reduce_loss
函数
def weight_reduce_loss(loss: Tensor,
weight: Optional[Tensor] = None,
reduction: str = 'mean',
avg_factor: Optional[float] = None) -> Tensor:
"""Apply element-wise weight and reduce loss.
Args:
loss (Tensor): Element-wise loss.
weight (Optional[Tensor], optional): Element-wise weights.
Defaults to None.
reduction (str, optional): Same as built-in losses of PyTorch.
Defaults to 'mean'.
avg_factor (Optional[float], optional): Average factor when
computing the mean of losses. Defaults to None.
Returns:
Tensor: Processed loss values.
"""
# if weight is specified, apply element-wise weight
if weight is not None:
loss = loss * weight
# if avg_factor is not specified, just reduce the loss
if avg_factor is None:
loss = reduce_loss(loss, reduction)
else:
# if reduction is mean, then average the loss by avg_factor
if reduction == 'mean':
# Avoid causing ZeroDivisionError when avg_factor is 0.0,
# i.e., all labels of an image belong to ignore index.
eps = torch.finfo(torch.float32).eps
loss = loss.sum() / (avg_factor + eps)
# if reduction is 'none', then do nothing, otherwise raise an error
elif reduction!= 'none':
raise ValueError('avg_factor can not be used with reduction="sum"')
return loss
这个函数首先根据是否提供了权重张量,对损失进行逐元素加权。然后,根据是否提供了平均因子和指定的减少策略来进一步处理损失。如果没有提供平均因子,直接使用reduce_loss
函数进行减少。如果减少策略是"mean"且提供了平均因子,则通过将损失总和除以平均因子加上一个极小值来避免除零错误。如果减少策略不是"none"且提供了平均因子,同时减少策略是"sum",则会抛出一个值错误。
三、weighted_loss
装饰器
def weighted_loss(loss_func: Callable) -> Callable:
"""Create a weighted version of a given loss function.
To use this decorator, the loss function must have the signature like
`loss_func(pred, target, **kwargs)`. The function only needs to compute
element-wise loss without any reduction. This decorator will add weight
and reduction arguments to the function. The decorated function will have
the signature like `loss_func(pred, target, weight=None, reduction='mean',
avg_factor=None, **kwargs)`.
:Example:
>>> import torch
>>> @weighted_loss
>>> def l1_loss(pred, target):
>>> return (pred - target).abs()
>>> pred = torch.Tensor([0, 2, 3])
>>> target = torch.Tensor([1, 1, 1])
>>> weight = torch.Tensor([1, 0, 1])
>>> l1_loss(pred, target)
tensor(1.3333)
>>> l1_loss(pred, target, weight)
tensor(1.)
>>> l1_loss(pred, target, reduction='none')
tensor([1., 1., 2.])
>>> l1_loss(pred, target, weight, avg_factor=2)
tensor(1.5000)
"""
@functools.wraps(loss_func)
def wrapper(pred: Tensor,
target: Tensor,
weight: Optional[Tensor] = None,
reduction: str = 'mean',
avg_factor: Optional[int] = None,
**kwargs) -> Tensor:
"""
Args:
pred (Tensor): The prediction.
target (Tensor): Target bboxes.
weight (Optional[Tensor], optional): The weight of loss for each
prediction. Defaults to None.
reduction (str, optional): Options are "none", "mean" and "sum".
Defaults to 'mean'.
avg_factor (Optional[int], optional): Average factor that is used
to average the loss. Defaults to None.
Returns:
Tensor: Loss tensor.
"""
# get element-wise loss
loss = loss_func(pred, target, **kwargs)
loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
return loss
return wrapper
这个装饰器用于创建一个带权重的损失函数版本。被装饰的损失函数应该只计算逐元素的损失,不进行任何减少操作。装饰器会为函数添加权重和减少策略的参数。在装饰后的函数中,首先调用原始的损失函数计算逐元素损失,然后使用weight_reduce_loss
函数对损失进行加权和减少处理,并返回最终的损失张量。