forked from ccf-ai-infra/GPUCodeForces
34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import numpy as np
|
||
|
||
BATCH_SIZE = 512
|
||
NUM_CLASSES = 1024
|
||
REDUCTION = 'mean'
|
||
# 每个样本的正类标签数量范围
|
||
MIN_LABELS = 1
|
||
MAX_LABELS = 10
|
||
|
||
class Model(nn.Module):
|
||
def __init__(self, reduction='mean'):
|
||
super(Model, self).__init__()
|
||
self.loss_fn = nn.MultiLabelMarginLoss(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每行包含正类索引，并用 -1 填充
|
||
target_np = np.full((BATCH_SIZE, NUM_CLASSES), -1, dtype=np.int64)
|
||
for i in range(BATCH_SIZE):
|
||
num_labels = np.random.randint(MIN_LABELS, MAX_LABELS + 1)
|
||
labels = np.random.choice(NUM_CLASSES, num_labels, replace=False)
|
||
target_np[i, :num_labels] = labels
|
||
target_tensor = torch.from_numpy(target_np)
|
||
|
||
return [input_tensor.contiguous(), target_tensor.contiguous()]
|
||
|
||
def get_init_inputs():
|
||
return [REDUCTION] |