Spaces:
Sleeping
Sleeping
File size: 1,627 Bytes
6cc8ae1 | 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 55 56 | """
Loss function factory.
Supports CrossEntropy with optional label smoothing and focal loss.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
class FocalLoss(nn.Module):
"""
Focal Loss for handling class imbalance.
FL(pt) = -alpha * (1 - pt)^gamma * log(pt)
"""
def __init__(self, alpha: float = 1.0, gamma: float = 2.0, reduction: str = 'mean'):
super().__init__()
self.alpha = alpha
self.gamma = gamma
self.reduction = reduction
def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
ce_loss = F.cross_entropy(inputs, targets, reduction='none')
pt = torch.exp(-ce_loss)
focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss
if self.reduction == 'mean':
return focal_loss.mean()
elif self.reduction == 'sum':
return focal_loss.sum()
return focal_loss
def create_loss(
loss_type: str = 'cross_entropy',
label_smoothing: float = 0.0,
focal_alpha: float = 1.0,
focal_gamma: float = 2.0,
) -> nn.Module:
"""
Factory function for loss functions.
Args:
loss_type: 'cross_entropy' or 'focal'
label_smoothing: smoothing factor (0 = none)
focal_alpha: alpha for focal loss
focal_gamma: gamma for focal loss
"""
if loss_type == 'cross_entropy':
return nn.CrossEntropyLoss(label_smoothing=label_smoothing)
elif loss_type == 'focal':
return FocalLoss(alpha=focal_alpha, gamma=focal_gamma)
else:
raise ValueError(f"Unknown loss type: {loss_type}")
|