pytorch计算运行时间

PyTorch提供了多种计算代码运行时间的方法，其中包括timeit和profile两种常用的工具。timeit可以测量代码开始时刻和结束时刻，然后求差，但是在GPU上的运行是异步的，因此需要使用PyTorch内置的计时工具和同步工具。而profile则是PyTorch自带的用于计算模型每个部分耗时的工具，可以计算CPU和GPU的耗时。下面是两种方法的示例代码：

使用timeit：

import torch

x = torch.randn((1, 1)).cuda()
y = torch.randn((1, 1)).cuda()

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：

import torch

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"))