forked from ccf-ai-infra/GPUCodeForces
92 lines
3.0 KiB
Python
92 lines
3.0 KiB
Python
|
###########################################################
|
|||
|
|
# 性能和精度验证程序
|
|||
|
|
###########################################################
|
|||
|
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
import time
|
|||
|
|
from conv2d_torch import Model, get_inputs, get_init_inputs
|
|||
|
|
from conv2d_cuda import ModelNew
|
|||
|
|
|
|||
|
|
# 禁用 TF32，确保与自定义 FP32 核精度对齐
|
|||
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|||
|
|
torch.backends.cudnn.allow_tf32 = False
|
|||
|
|
torch.backends.cudnn.deterministic = True
|
|||
|
|
|
|||
|
|
def _time_cuda_model(fn, inputs, iters=300, warmup=50):
|
|||
|
|
torch.cuda.synchronize()
|
|||
|
|
for _ in range(warmup):
|
|||
|
|
_ = fn(*inputs)
|
|||
|
|
torch.cuda.synchronize()
|
|||
|
|
start = torch.cuda.Event(enable_timing=True)
|
|||
|
|
end = torch.cuda.Event(enable_timing=True)
|
|||
|
|
start.record()
|
|||
|
|
for _ in range(iters):
|
|||
|
|
_ = fn(*inputs)
|
|||
|
|
end.record()
|
|||
|
|
torch.cuda.synchronize()
|
|||
|
|
ms = start.elapsed_time(end) / iters # 平均每次毫秒
|
|||
|
|
return ms / 1000.0 # 转为秒
|
|||
|
|
|
|||
|
|
def run_benchmark():
|
|||
|
|
# 检查 CUDA 是否可用
|
|||
|
|
if not torch.cuda.is_available():
|
|||
|
|
print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
|
|||
|
|
return
|
|||
|
|
else:
|
|||
|
|
device = torch.device("cuda")
|
|||
|
|
|
|||
|
|
torch.backends.cudnn.benchmark = True
|
|||
|
|
|
|||
|
|
# 初始化模型
|
|||
|
|
init_inputs = get_init_inputs()
|
|||
|
|
init_inputs = [
|
|||
|
|
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
|
|||
|
|
]
|
|||
|
|
inputs = get_inputs()
|
|||
|
|
inputs = [
|
|||
|
|
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
torch_model = Model(*init_inputs).cuda()
|
|||
|
|
cuda_model = ModelNew(*init_inputs).cuda()
|
|||
|
|
|
|||
|
|
torch_model.eval()
|
|||
|
|
cuda_model.eval()
|
|||
|
|
|
|||
|
|
print("-------------------- 精度对齐验证 --------------------")
|
|||
|
|
with torch.no_grad():
|
|||
|
|
output_torch = torch_model(*inputs)
|
|||
|
|
output_cuda = cuda_model(*inputs)
|
|||
|
|
|
|||
|
|
abs_diff = (output_torch - output_cuda).abs()
|
|||
|
|
max_diff = abs_diff.max().item()
|
|||
|
|
mean_diff = abs_diff.mean().item()
|
|||
|
|
|
|||
|
|
print(f"最大差异: {max_diff:.6f}")
|
|||
|
|
print(f"平均差异: {mean_diff:.6f}")
|
|||
|
|
|
|||
|
|
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-5, atol=1e-5)
|
|||
|
|
if precision_flag:
|
|||
|
|
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
|||
|
|
else:
|
|||
|
|
print("❌ 精度不一致！")
|
|||
|
|
|
|||
|
|
print("\n-------------------- 性能加速比测试 --------------------")
|
|||
|
|
num_iterations = 300
|
|||
|
|
|
|||
|
|
# 计时
|
|||
|
|
torch_time = _time_cuda_model(torch_model, inputs, iters=num_iterations, warmup=50)
|
|||
|
|
cuda_time = _time_cuda_model(cuda_model, inputs, iters=num_iterations, warmup=50)
|
|||
|
|
|
|||
|
|
print(f"PyTorch (conv2d) 平均执行时间: {torch_time:.6f} 秒")
|
|||
|
|
print(f"自定义 CUDA conv2d 平均执行时间: {cuda_time:.6f} 秒")
|
|||
|
|
speedup = 0.0
|
|||
|
|
if cuda_time > 0:
|
|||
|
|
speedup = torch_time / cuda_time
|
|||
|
|
print(f"相对 PyTorch 加速比: {speedup:.2f}x")
|
|||
|
|
else:
|
|||
|
|
print("CUDA 内核执行时间为0，无法计算加速比。")
|
|||
|
|
return precision_flag, speedup
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
precision_flag, speedup = run_benchmark()
|