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

2019-09-29  本文已影响0人  几时见得清梦

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

# 方式一:
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()
上一篇下一篇

猜你喜欢

热点阅读