|
import torch |
|
import torch.nn as nn |
|
from uniperceiver.config import configurable |
|
from .build import LOSSES_REGISTRY |
|
|
|
@LOSSES_REGISTRY.register() |
|
class CrossEntropy(nn.Module): |
|
@configurable |
|
def __init__(self, loss_weight=1.0, reduction='mean', loss_fp32=False): |
|
super(CrossEntropy, self).__init__() |
|
if reduction is None: |
|
reduction = 'mean' |
|
self.criterion_func = nn.CrossEntropyLoss(ignore_index=-1, reduction=reduction) |
|
if not isinstance(loss_weight, float): |
|
self.loss_weight = 1.0 |
|
else: |
|
self.loss_weight = loss_weight |
|
self.reduction = reduction |
|
self.loss_fp32 = loss_fp32 |
|
|
|
def criterion(self, x, target): |
|
if self.loss_fp32 and x.dtype != torch.float32: |
|
loss = self.criterion_func(x.to(torch.float32), target).to(x.dtype) |
|
else: |
|
loss = self.criterion_func(x, target) |
|
return loss.mean() |
|
|
|
@classmethod |
|
def from_config(cls, cfg): |
|
return { |
|
'loss_weight': getattr(cfg.LOSSES, 'LOSS_WEIGHT', None), |
|
'reduction': getattr(cfg.LOSSES, 'REDUCTION', 'mean'), |
|
'loss_fp32': getattr(cfg.LOSSES, 'LOSS_FP32', False), |
|
} |
|
|
|
@classmethod |
|
def add_config(cls, cfg): |
|
cfg.LOSSES.LOSS_WEIGHT = None |
|
cfg.LOSSES.REDUCTION = 'mean' |
|
|
|
def forward(self, outputs_dict): |
|
ret = {} |
|
|
|
for logit, target, loss_identification in zip(outputs_dict['logits'], |
|
outputs_dict['targets'], |
|
outputs_dict['loss_names']): |
|
|
|
loss = self.criterion(logit, target) |
|
if self.loss_weight != 1.0: |
|
loss *= self.loss_weight |
|
loss_name = 'CrossEntropy_Loss' |
|
if len(loss_identification) > 0: |
|
loss_name = loss_name+ f' ({loss_identification})' |
|
ret.update({ loss_name: loss }) |
|
|
|
|
|
return ret |
|
|
