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 [] # 没有特殊的初始化输入需求
|