forked from ccf-ai-infra/GPUCodeForces
47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
class Model(nn.Module):
|
||
"""
|
||
BCE with Sigmoid implementation for binary classification.
|
||
使用binary_cross_entropy_with_logits作为基准
|
||
"""
|
||
def __init__(self, reduction='sum'):
|
||
super(Model, self).__init__()
|
||
self.reduction = reduction
|
||
|
||
def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Compute BCE Loss with logits.
|
||
|
||
Args:
|
||
inputs (torch.Tensor): Predicted logits
|
||
targets (torch.Tensor): Ground truth labels (0 or 1)
|
||
|
||
Returns:
|
||
torch.Tensor: Computed BCE loss
|
||
"""
|
||
# 确保类型一致
|
||
inputs = inputs.to(torch.float32)
|
||
targets = targets.to(torch.float32)
|
||
|
||
# 使用PyTorch的标准实现
|
||
return F.binary_cross_entropy_with_logits(
|
||
inputs,
|
||
targets,
|
||
reduction=self.reduction
|
||
)
|
||
|
||
batch_size = 256
|
||
num_features = 2000
|
||
|
||
def get_inputs():
|
||
# 生成logits（不是概率）
|
||
input_logits = torch.randn(batch_size, num_features, dtype=torch.float32)
|
||
target_labels = torch.randint(0, 2, (batch_size, num_features), dtype=torch.float32)
|
||
return [input_logits, target_labels]
|
||
|
||
def get_init_inputs():
|
||
return []
|