Pytorch如何精准记录函数运行时间

0. 引言

参考Pytorch官方文档对CUDA的描述，GPU的运算是异步执行的。一般来说，异步计算的效果对于调用者来说是不可见的，因为

1. 每个设备按照排队的顺序执行操作
2. Pytorch对于CPU和GPU的同步，GPU间的同步是自动执行的，不需要显示写在代码中

异步计算的后果是，没有同步的时间测量是不准确的。

1. 解决方案

参考引言中提到的帮助文档，Pytorch官方给出的解决方案是使用torch.cuda.Event记录时间，具体代码如下：

start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()# Run your code snippet hereend_event.record()
torch.cuda.synchronize()  # Wait for the events to be recorded!
elapsed_time_ms = start_event.elapsed_time(end_event)  # elapsed time (ms)

将你的代码插入start_event.record()和end_event.record()中间以测量时间（单位毫秒）。本人亲测，使用time.time()函数得到的函数运行时间为105ms，而使用该方法得到的运行时间为19ms！

有能力的读者也可以包装为装饰器或者with语句使用：

先书写一个自定义with类（ContextManager）

class CudaTimer:def __init__(self):self.start_event = torch.cuda.Event(enable_timing=True)self.end_event = torch.cuda.Event(enable_timing=True)def __enter__(self):self.start_event.record()return selfdef __exit__(self, exc_type, exc_value, traceback):self.end_event.record()torch.cuda.synchronize()self.elapsed_time = self.start_event.elapsed_time(self.end_event) / 1000 # ms -> s

再安装如下with语句返回：

with CudaTimer() as timer:# run your code here
dt = timer.elapsed_time  # s

这样保证了多个文件调用时语句的简单性。特别提醒：获取timer.elapsed_time操作不要写在with语句内部。在with语句未结束时，是无法获取timer的成员变量的。