pytorch框架下计算模型推理时间

pytorch框架下计算推理时间

我们需要使用torch.cuda.synchronize()函数进行推理时间开销计算。
关于torch.cuda.synchronize()，参照torch官方文档torch.cuda.synchronize

其主要作用是等待 CUDA 设备上所有流中的所有内核完成。

一般进程按照如下操作：

torch.cuda.synchronize()
start = time.time()
XXX（进程）
XXX（进程）
XXX(进程)
torch.cuda.synchronize()
end = time.time()
total_time = end - start
print('total_time:{:.2f}'.format(total_time))

图像FPS计算：

total_time = 0
torch.cuda.synchronize()
start = time.time()
XXX（进程）
torch.cuda.synchronize()
end = time.time()
single_fps = 1 / (end - start)
total_time += end - start
fps = (i + 1) / total_time
print(' ({:.2f} fps total_time:{:.2f} single_fps:{})'.format(fps, total_time,single_fps))