PyTorch 测量代码段的运行时间

Contents

• timeit
• profile
• References

timeit

• timeit. 测量代码开始时刻和结束时刻，然后求差

---

• pytorch 的代码经常会运行在 GPU 上，而在 GPU 上的运行都是异步的，意味着采用一般的 timeit 操作不能准确地得到运行时总和，因此我们一般需要用 pytorch 内置的计时工具和同步工具 (单位：ms)

start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)

start.record()
z = x + y
end.record()

# Waits for everything to finish running
torch.cuda.synchronize()

print(start.elapsed_time(end))

profile

• profile. pytorch 自带或者第三方的代码耗时工具

---

• timeit 的方法测试一些小代码还勉强适用，但是在大规模的测试中显然会变得很麻烦，当然，你可以通过添加修饰器的方式去简化一行行重复人工添加这些时间测量代码的枯燥，但是这也并不是最好的解决方案
• 幸运的是，pytorch 自带了 profile 用于计算模型每个部分耗时 ，其既可以计算 cpu 耗时，也可以计算 gpu 耗时

x = torch.randn((1, 1), requires_grad=True)
with torch.autograd.profiler.profile(enabled=True) as prof:
	for _ in range(100):  # any normal python code, really!
    	y = x ** 2
print(prof.key_averages().table(sort_by="self_cpu_time_total"))

References

• 测量 pytorch 代码段的运行时间
• PyTorch Profiler Tutorial