File size: 1,162 Bytes
07423df |
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 |
import logging
from typing import Any, KeysView
from torch import nn
__all__ = ["Losses"]
logger = logging.getLogger(__name__)
class CrossEntropyLoss(nn.Module):
def __init__(self, cfg: Any):
super().__init__()
self.cfg = cfg
self.loss_fn = nn.CrossEntropyLoss()
def forward(self, logits, labels):
return self.loss_fn(logits, labels.reshape(-1).long())
class BinaryCrossEntropyLoss(nn.Module):
def __init__(self, cfg: Any):
super().__init__()
self.cfg = cfg
self.loss_fn = nn.BCEWithLogitsLoss()
def forward(self, logits, labels):
return self.loss_fn(logits, labels)
class Losses:
"""Losses factory."""
_losses = {
"CrossEntropyLoss": CrossEntropyLoss,
"BinaryCrossEntropyLoss": BinaryCrossEntropyLoss,
}
@classmethod
def names(cls) -> KeysView:
return cls._losses.keys()
@classmethod
def get(cls, name: str) -> Any:
"""Access to Losses.
Args:
name: losses name
Returns:
A class to build the Losses
"""
return cls._losses.get(name, CrossEntropyLoss)
|