forked from ccf-ai-infra/GPUCodeForces
44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
|
|||
|
|
class Model(nn.Module):
|
|||
|
|
"""
|
|||
|
|
合理优化的PyTorch BCE Loss实现
|
|||
|
|
使用PyTorch内置函数，避免重复的数值处理
|
|||
|
|
"""
|
|||
|
|
def __init__(self):
|
|||
|
|
super(Model, self).__init__()
|
|||
|
|
|
|||
|
|
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|||
|
|
"""
|
|||
|
|
使用PyTorch内置的binary_cross_entropy函数
|
|||
|
|
让PyTorch自己处理数值稳定性
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
input (torch.Tensor): 预测概率 (0-1之间)
|
|||
|
|
target (torch.Tensor): 真实标签 (0或1)
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
torch.Tensor: BCE Loss标量值
|
|||
|
|
"""
|
|||
|
|
# 直接使用内置函数，让PyTorch处理数值稳定性
|
|||
|
|
return torch.nn.functional.binary_cross_entropy(
|
|||
|
|
input,
|
|||
|
|
target,
|
|||
|
|
reduction='sum'
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
batch_size = 256
|
|||
|
|
num_features = 2000
|
|||
|
|
|
|||
|
|
def get_inputs():
|
|||
|
|
"""
|
|||
|
|
生成合理的测试数据
|
|||
|
|
"""
|
|||
|
|
input_probs = torch.sigmoid(torch.randn(batch_size, num_features))
|
|||
|
|
target_labels = torch.randint(0, 2, (batch_size, num_features)).float()
|
|||
|
|
return [input_probs, target_labels]
|
|||
|
|
|
|||
|
|
def get_init_inputs():
|
|||
|
|
return [] # 没有特殊的初始化输入需求
|