forked from ccf-ai-infra/GPUCodeForces
69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
class Model(nn.Module):
|
||
"""
|
||
Focal Loss implementation for binary classification.
|
||
Focal Loss = -α * (1-pt)^γ * log(pt)
|
||
where pt = p if target=1, else (1-p)
|
||
"""
|
||
def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
|
||
super(Model, self).__init__()
|
||
self.alpha = alpha
|
||
self.gamma = gamma
|
||
self.reduction = reduction
|
||
|
||
def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Compute Focal Loss between inputs and targets.
|
||
|
||
Args:
|
||
inputs (torch.Tensor): Predicted logits of shape (batch_size, num_classes)
|
||
targets (torch.Tensor): Ground truth labels of shape (batch_size,)
|
||
|
||
Returns:
|
||
torch.Tensor: Computed focal loss
|
||
"""
|
||
# 确保输入输出类型完全一致
|
||
inputs = inputs.to(torch.float32)
|
||
targets = targets.to(torch.float32)
|
||
|
||
# 确保targets的形状与inputs匹配
|
||
if inputs.dim() == 2 and inputs.size(1) == 1:
|
||
targets = targets.unsqueeze(1) # 将targets从[batch_size]变为[batch_size, 1]
|
||
|
||
# Convert logits to probabilities
|
||
probs = torch.sigmoid(inputs)
|
||
|
||
# Compute pt
|
||
pt = torch.where(targets == 1, probs, 1 - probs)
|
||
|
||
# Compute focal weight
|
||
focal_weight = self.alpha * torch.pow(1 - pt, self.gamma)
|
||
|
||
# Compute binary cross entropy
|
||
bce = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
|
||
|
||
# Apply focal weight
|
||
focal_loss = focal_weight * bce
|
||
|
||
if self.reduction == 'mean':
|
||
return focal_loss.mean()
|
||
elif self.reduction == 'sum':
|
||
return focal_loss.sum()
|
||
else:
|
||
return focal_loss
|
||
|
||
batch_size = 32
|
||
num_classes = 1 # Binary classification
|
||
|
||
def get_inputs():
|
||
# Generate random logits - 明确指定float32
|
||
inputs = torch.randn(batch_size, num_classes, dtype=torch.float32)
|
||
# Generate random binary targets (0 or 1) - 明确指定float32
|
||
targets = torch.randint(0, 2, (batch_size,), dtype=torch.float32)
|
||
return [inputs, targets]
|
||
|
||
def get_init_inputs():
|
||
return [] # No special initialization inputs needed |