PyTorch测试模型执行计算耗费的时间

一般我们都会使用这种方式一测试时间：

# 方式一：
start = time.time()
result = model(input)
end = time.time()

但是正确的应该是下边这种方式二：

# 方式二：
torch.cuda.synchronize()
start = time.time()
result = model(input)
torch.cuda.synchronize()
end = time.time()

为什么是这样呢？
在pytorch里面，程序的执行都是异步的。如果采用第一种方式，测试的时间会很短，因为执行完end=time.time()程序就退出了，后台的cu也因为python的退出退出了；
如果采用第二种方式，代码会同步cu的操作，等待gpu上的操作都完成了再继续计算end = time.time()

如果将方式一代码改为方式三：

# 方式三：
start = time.time()
result = model(input)
print(result) # 或 result.cpu()
end = time.time()

这时候会发现方式三和方式二的时间是类似的，因为方式三会等待gpu上的结果执行完传给print函数，所以此时间就和方式二同步操作的时间基本上是一致的了。将print(result)换成result.cpu()也可以得到相同的结果。

torch.cuda.synchronize()
等待当前设备上所有流中的所有核心完成。

如果不加syn, forward会马上返回。但加上syn后，cpu会等待模型实际运行完再获取数据.实际使用时，如果不逐段统计时间，可以不加这个sync

t3= time.time()
output = model(torch_img)
torch.cuda.synchronize()
t4 = time.time()
print(t4-t3)

参考

1. PyTorch官方文档：torch.cuda
2. CSDN：pytorch 正确的测试时间的代码 torch.cuda.synchronize()