如何简便计时运行在GPU上的PyTorch代码？

当然有啦！因为GPU操作大多是异步执行的，直接用Python内置的time模块计时会严重不准——你记录的可能只是GPU任务提交的时间，而非实际运行时间。PyTorch本身就提供了专门针对GPU的计时工具，下面给你介绍几个实用又简便的方法：

方法一：使用torch.cuda.Event（最常用） #

这是PyTorch官方推荐的GPU计时方式，能精准记录GPU操作的开始和结束时间，还能自动处理同步问题。示例代码如下：

import torch

# 初始化两个CUDA事件，开启计时功能
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)

# 记录开始事件（此时GPU开始执行后续任务）
start_event.record()

# 👇这里替换成你的GPU代码
input_tensor = torch.randn(1024, 1024).cuda()
output = torch.matmul(input_tensor, input_tensor)

# 记录结束事件
end_event.record()

# 等待GPU完成所有操作（同步GPU，确保计时准确）
torch.cuda.synchronize()

# 计算耗时，单位为毫秒
elapsed_ms = start_event.elapsed_time(end_event)
print(f"GPU代码实际耗时: {elapsed_ms:.2f} ms")

注意：torch.cuda.synchronize()是关键——它会阻塞CPU，直到GPU完成所有pending任务，这样你拿到的才是真实的GPU运行时间，而不是任务提交的时间差。

方法二：使用torch.utils.benchmark（适合基准测试） #

如果你需要多次运行代码并统计性能（比如均值、标准差），这个工具会更省心，它能自动处理同步和重复测试：

from torch.utils.benchmark import Timer
import torch

# 定义你的GPU操作函数
def run_gpu_task():
    input_tensor = torch.randn(1024, 1024).cuda()
    output = torch.matmul(input_tensor, input_tensor)
    torch.cuda.synchronize()  # 内部同步确保任务完成

# 创建Timer对象，指定要测试的语句和环境
timer = Timer(
    stmt="run_gpu_task()",
    setup="from __main__ import run_gpu_task",
    num_threads=1
)

# 运行100次测试，输出统计结果
result = timer.timeit(100)
print(result)

运行后会输出类似这样的结果，包含耗时的均值、标准差等，非常适合做性能对比：

<torch.utils.benchmark.utils.common.Measurement object at 0x7f8b1c0b5d30>
run_gpu_task()
  1.23 ms ± 0.05 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

补充：Jupyter Notebook专属技巧 #

如果你在Notebook里调试代码，可以用%timeit魔法命令，但一定要加上同步操作，否则结果会不准：

# 确保tensor在GPU上
input_tensor = torch.randn(1024, 1024).cuda()

# 带同步的计时
%timeit -n 10 -r 3 (torch.matmul(input_tensor, input_tensor), torch.cuda.synchronize())

这样就能快速得到GPU操作的平均耗时啦！

内容的提问来源于stack exchange，提问作者macharya