forked from ccf-ai-infra/GPUCodeForces
76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
|
|||
|
|
class Model(nn.Module):
|
|||
|
|
def __init__(self, dim=-1):
|
|||
|
|
super(Model, self).__init__()
|
|||
|
|
self.dim = dim
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
# 简化实现：直接使用PyTorch内置Softmax
|
|||
|
|
return torch.softmax(x, dim=self.dim)
|
|||
|
|
|
|||
|
|
def get_inputs():
|
|||
|
|
"""生成适合Softmax测试的输入数据"""
|
|||
|
|
torch.manual_seed(42)
|
|||
|
|
|
|||
|
|
# 使用更适合CUDA优化的数据规模
|
|||
|
|
batch_size = 128 # 减少批量大小，避免内存限制
|
|||
|
|
seq_length = 512 # 减少序列长度，更适合CUDA优化
|
|||
|
|
|
|||
|
|
# 生成简单的测试数据
|
|||
|
|
input_data = torch.randn(batch_size, seq_length) * 0.1
|
|||
|
|
|
|||
|
|
return [input_data]
|
|||
|
|
|
|||
|
|
def get_init_inputs():
|
|||
|
|
"""获取模型初始化参数"""
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
def generate_test_cases():
|
|||
|
|
"""生成多种测试用例，用于全面测试Softmax性能"""
|
|||
|
|
test_cases = []
|
|||
|
|
|
|||
|
|
# 小规模测试用例 - 测试warp级优化
|
|||
|
|
test_cases.append({
|
|||
|
|
'name': 'small_batch_small_seq',
|
|||
|
|
'input': torch.randn(64, 128) * 0.1, # 适合warp级处理
|
|||
|
|
'description': '小批量小序列测试（warp优化）'
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 中等规模测试用例 - 测试块级优化
|
|||
|
|
test_cases.append({
|
|||
|
|
'name': 'medium_batch_medium_seq',
|
|||
|
|
'input': torch.randn(256, 512) * 0.15, # 适合块级处理
|
|||
|
|
'description': '中等批量中等序列测试（块优化）'
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 大规模测试用例 - 测试分层归约优化
|
|||
|
|
test_cases.append({
|
|||
|
|
'name': 'large_batch_large_seq',
|
|||
|
|
'input': torch.randn(512, 2048) * 0.2, # 适合分层归约
|
|||
|
|
'description': '大批量大序列测试（分层归约）'
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 超大规模测试用例 - 测试多遍归约优化
|
|||
|
|
test_cases.append({
|
|||
|
|
'name': 'huge_batch_huge_seq',
|
|||
|
|
'input': torch.randn(1024, 4096) * 0.25, # 适合多遍归约
|
|||
|
|
'description': '超大批量超长序列测试（多遍归约）'
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 极端情况测试用例 - 测试数值稳定性
|
|||
|
|
test_cases.append({
|
|||
|
|
'name': 'extreme_values',
|
|||
|
|
'input': torch.tensor([[100.0, -100.0, 50.0, -50.0, 200.0, -200.0]]),
|
|||
|
|
'description': '极端数值稳定性测试'
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 边界情况测试用例 - 测试边界处理
|
|||
|
|
test_cases.append({
|
|||
|
|
'name': 'boundary_cases',
|
|||
|
|
'input': torch.tensor([[1e-10, 1e10, 0.0, -1e10, 1e-5]]),
|
|||
|
|
'description': '边界数值处理测试'
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
return test_cases
|