forked from ccf-ai-infra/GPUCodeForces
36 lines
1.0 KiB
Python
36 lines
1.0 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
|
||
BATCH_SIZE = 512
|
||
DIM = 4096
|
||
SHAPE = (BATCH_SIZE, DIM)
|
||
REDUCTION = 'mean'
|
||
|
||
class Model(nn.Module):
|
||
"""
|
||
使用 PyTorch 内置的 torch.nn.SoftMarginLoss 作为基准模型。
|
||
"""
|
||
def __init__(self, reduction='mean'):
|
||
super(Model, self).__init__()
|
||
self.loss_fn = nn.SoftMarginLoss(reduction=reduction)
|
||
|
||
def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor:
|
||
return self.loss_fn(input_tensor, target_tensor)
|
||
|
||
def get_inputs():
|
||
|
||
"""
|
||
生成用于测试的输入张量。
|
||
"""
|
||
input_tensor = torch.randn(SHAPE, dtype=torch.float32)
|
||
# target 张量必须只包含 1 和 -1
|
||
# 使用 randint 生成 0 或 1，然后映射到 -1 或 1
|
||
target_tensor = torch.randint(0, 2, SHAPE, dtype=torch.float32) * 2 - 1
|
||
|
||
return [input_tensor.contiguous(), target_tensor.contiguous()]
|
||
|
||
def get_init_inputs():
|
||
"""
|
||
提供模型初始化所需的参数。
|
||
"""
|
||
return [REDUCTION] |