importtorchimporttime# 1. 设置设备iftorch.cuda.is_available(): device=torch.device("cuda")print(f"检测到设备: {torch.cuda.get_device_name(0)}")else: device=torch.device("cpu")print("使用 CPU")# 2. 增加矩阵尺寸，让 GPU 有事可做 (比如 4000x4000 或 8000x8000)# 1000x1000 对高端卡来说太小了size=4000print(f"正在初始化矩阵 (Size: {size}x{size})...")A=torch.randn(size, size,device=device)B=torch.randn(size, size,device=device)# --- 关键步骤：GPU 预热 (Warm-up) ---print("正在进行 GPU 预热 (第一次运行通常很慢)...")# 跑几次空运算，触发库加载和初始化for_inrange(5): _=torch.mm(A, B)torch.cuda.synchronize()print("预热完成，开始正式计时！")# ------------------------------------# 3. 正式计时 GPUtorch.cuda.synchronize()start_time=time.time()# 循环多次取平均值更准确loop_times=10for_inrange(loop_times): C=torch.mm(A, B)torch.cuda.synchronize()# 等待所有计算完成end_time=time.time()avg_gpu_time=(end_time - start_time)/ loop_times print(f"GPU 平均运行时间: {avg_gpu_time:.5f} 秒")# 4. 对比 CPU (如果尺寸太大，CPU可能会非常慢，建议 CPU 测试时如果太慢就减少 loop)iftorch.cuda.is_available(): print("正在准备 CPU 数据 (这可能需要一点时间)...")A_cpu=A.cpu()B_cpu=B.cpu()start_time_cpu=time.time()# CPU 一般不需要像 GPU 那样严格预热，但跑一次也没坏处# 注意：如果 size 是 4000，CPU 跑这 10 次可能会很慢cpu_loops=2# 减少 CPU 循环次数以免等太久for_inrange(cpu_loops): C_cpu=torch.mm(A_cpu, B_cpu)end_time_cpu=time.time()avg_cpu_time=(end_time_cpu - start_time_cpu)/ cpu_loops print(f"CPU 平均运行时间: {avg_cpu_time:.5f} 秒")# 修正加速比公式：CPU时间 / GPU时间speedup=avg_cpu_time / avg_gpu_time print(f"---")print(f"真实性能对比: GPU 比 CPU 快 {speedup:.2f} 倍")
