forked from ccf-ai-infra/GPUCodeForces
30 lines
991 B
Python
30 lines
991 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
BATCH_SIZE = 512
|
|
NUM_CLASSES = 4096
|
|
REDUCTION = 'mean'
|
|
P = 1 # 1 for L1 hinge, 2 for L2
|
|
MARGIN = 1.0
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
使用 PyTorch 内置的 torch.nn.MultiMarginLoss 作为基准模型。
|
|
"""
|
|
def __init__(self, p=1, margin=1.0, reduction='mean'):
|
|
super(Model, self).__init__()
|
|
# weight is not benchmarked for simplicity, but the CUDA kernel supports it.
|
|
self.loss_fn = nn.MultiMarginLoss(p=p, margin=margin, 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(BATCH_SIZE, NUM_CLASSES, dtype=torch.float32)
|
|
|
|
target_tensor = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
|
|
|
|
return [input_tensor.contiguous(), target_tensor.contiguous()]
|
|
|
|
def get_init_inputs():
|
|
return [P, MARGIN, REDUCTION] |