forked from ccf-ai-infra/GPUCodeForces
52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
class Model(nn.Module):
|
||
"""
|
||
Baseline Focal Loss using PyTorch's most numerically stable API.
|
||
"""
|
||
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:
|
||
# --- 关键修复1：强制转换数据类型 ---
|
||
# 确保所有计算都在float32下进行，避免类型错误
|
||
inputs = inputs.to(torch.float32)
|
||
targets = targets.to(torch.float32)
|
||
|
||
# --- 关键修复2：统一输入维度 ---
|
||
if inputs.dim() == 2 and inputs.size(1) == 1:
|
||
inputs = inputs.squeeze(1)
|
||
|
||
# --- 使用PyTorch官方推荐的稳定API计算BCE ---
|
||
# 现在inputs和targets的维度和类型都正确了
|
||
bce = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
|
||
|
||
# 后续的Focal Loss计算保持不变
|
||
probs = torch.sigmoid(inputs)
|
||
pt = torch.where(targets == 1, probs, 1 - probs)
|
||
focal_weight = self.alpha * torch.pow(1 - pt, self.gamma)
|
||
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 = 1024
|
||
num_classes = 1
|
||
|
||
def get_inputs():
|
||
inputs = torch.randn(batch_size, num_classes)
|
||
targets = torch.randint(0, 2, (batch_size,))
|
||
return [inputs, targets]
|
||
|
||
def get_init_inputs():
|
||
return []
|