forked from ccf-ai-infra/GPUCodeForces
54 lines
1.4 KiB
Python
54 lines
1.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
# -------------------------------------------------------------
|
|
# 常量定义
|
|
# -------------------------------------------------------------
|
|
N_BATCH = 100
|
|
D_VECTOR = 128
|
|
|
|
# 假设 dim=1
|
|
DIM = 1
|
|
|
|
|
|
# -------------------------------------------------------------
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
汉明距离 (Hamming Distance) 的纯 PyTorch 基准实现
|
|
"""
|
|
|
|
def __init__(self, dim=1):
|
|
super().__init__()
|
|
# 假设 dim=1
|
|
self.dim = dim
|
|
|
|
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
|
|
# x1, x2: (N, D), 假设为 long/int/bool 类型
|
|
|
|
# 1. 比较 (x1 != x2)
|
|
# 如果 x1=[1, 0, 1], x2=[1, 1, 1]
|
|
# diff=[False, True, False]
|
|
diff = (x1 != x2)
|
|
|
|
# 2. 求和 (Sum)
|
|
# False.sum() = 0, True.sum() = 1
|
|
# sum([0, 1, 0]) = 1
|
|
# 我们转换为 float 以匹配 CUDA 版本的输出类型
|
|
return torch.sum(diff, dim=self.dim).to(torch.float32)
|
|
|
|
|
|
def get_inputs():
|
|
"""
|
|
生成两个 (N, D) 形状的 *整数* 输入
|
|
(汉明距离的标准输入)
|
|
"""
|
|
# 随机生成 0 或 1
|
|
input1 = torch.randint(0, 2, (N_BATCH, D_VECTOR), dtype=torch.long)
|
|
input2 = torch.randint(0, 2, (N_BATCH, D_VECTOR), dtype=torch.long)
|
|
return [input1, input2]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [DIM] |