File size: 1,729 Bytes
26225c5 | 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 | import torch
from torch import Tensor
from torch.nn import BCEWithLogitsLoss as TorchBCEWithLogitsLoss
from src.loss.weighted import WeightedLossMixIn
__all__ = ['WeightedBCEWithLogitsLoss', 'BCEWithLogitsLoss']
class WeightedBCEWithLogitsLoss(WeightedLossMixIn, TorchBCEWithLogitsLoss):
"""Weighted BCE loss between predicted and target offsets. This is
basically the BCEWithLogitsLoss except that positive weights must be
passed at forward time to give more importance to some items.
Besides, we remove the constraint of passing `pos_weight` as a
Tensor. This simplifies instantiation with hydra.
"""
def __init__(self, *args, pos_weight=None, **kwargs):
if pos_weight is not None and not isinstance(pos_weight, Tensor):
pos_weight = torch.as_tensor(pos_weight)
super().__init__(
*args, pos_weight=pos_weight, reduction='none', **kwargs)
def load_state_dict(self, state_dict, strict=True):
"""Normal `load_state_dict` behavior, except for the shared
`pos_weight`.
"""
# Get the weight from the state_dict
pos_weight = state_dict.get('pos_weight')
state_dict.pop('pos_weight')
# Normal load_state_dict, ignoring pos_weight
out = super().load_state_dict(state_dict, strict=strict)
# Set the pos_weight
self.pos_weight = pos_weight
return out
class BCEWithLogitsLoss(WeightedBCEWithLogitsLoss):
"""BCE loss between predicted and target offsets.
The forward signature allows using this loss as a weighted loss,
with input weights ignored.
"""
def forward(self, input, target, weight):
return super().forward(input, target, None)
|