|
import torch |
|
import torch.nn as nn |
|
from uniperceiver.config import configurable |
|
from .build import LOSSES_REGISTRY |
|
|
|
@LOSSES_REGISTRY.register() |
|
class BCEWithLogits(nn.Module): |
|
@configurable |
|
def __init__(self, loss_weight=1.0): |
|
super(BCEWithLogits, self).__init__() |
|
self.criterion = nn.BCEWithLogitsLoss(reduction="mean") |
|
if not isinstance(loss_weight, float): |
|
self.loss_weight = 1.0 |
|
else: |
|
self.loss_weight = loss_weight |
|
|
|
@classmethod |
|
def from_config(cls, cfg): |
|
return { |
|
'loss_weight' : getattr(cfg.LOSSES, 'LOSS_WEIGHT', 1.0) |
|
} |
|
|
|
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 = 'BCEWithLogits_Loss' |
|
if len(loss_identification) > 0: |
|
loss_name = loss_name+ f' ({loss_identification})' |
|
ret.update({ loss_name: loss }) |
|
|
|
return ret |
|
|
