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 |