如何正确地计算神经网络模型的推理时间【含代码，以pytorch为例】

文章目录

- 前言
- 原理
- 代码
- 参考资料

前言

无论是做学术还是做工程项目，实际推理时间(inference time)总是我们衡量一个模型好坏的重要参照指标。目前已经也有很多博客在介绍如何计算一个神经网络模型的推理时间，但是写得都比较"粗糙"，在看了一些国内外的博客后，对这个问题有了一些总结性的经验，故写下这篇博客。

原理

某些博客直接将python程序的测速方式用到神经网络模型的测速中来（使用time.time计时），这忽略了深度学习问题的特殊性（在GPU上进行推理）。直接使用这种方式进行测速会忽略两个关键问题，导致测出来的速度不准：① 异步执行；② GPU预热。

下文中，我们先介绍不同的问题以及其原理，然后提出解决办法。

异步执行

在深度学习模型推理过程中，我们通常会将模型部署在GPU上，运行过程会涉及到cpu与gpu两种不同设备，而目前深度学习基本都是采用异步执行的方式（具体来说，就是在GPU调用某个函数时，它会在指定的设备上进行排队，不会占用其他的设备资源，从而完成并行计算）。

在这里插入图片描述

图片来自参考链接[1]

这种机制能够加快模型的推理速度，但这会使得计算推理时间变得麻烦，因为通常python程序使用time进行计时，这是在CPU端运行的，有可能存在GPU端的程序还没有运行完，但这边已经停止这一轮次的计时了，导致最终得到的结果不准确。

GPU预热

目前的GPU设备在不同的状态下会启用不同的工作模式，如果一台GPU在未使用的情况下，会处于低功耗状态（会关闭许多硬件，包括内存子系统、内部子系统、计算核等）。

如果此时将模型部署在上面并进行推理时，会触发GPU的一系列初始化，这会浪费较多的时间在GPU预热阶段，导致时间测量不准确。计算出来的推理速度也无法反映真正使用场景下的速度(持续模式)。

解决方法

为了解决异步执行的问题，使用pytorch自带的torch.cuda.synchronize()，它会将异步模式转换成同步模式，只有在当前轮次的GPU推理完成之后才停止计时，并且使用torch.cuda.Event来跟踪GPU程序并进行更新；

torch.cuda.Event和torch.cuda.synchronize()的官方文档说明

为了解决GPU预热的问题，我们可以在进行实际测量前运行一些示例，完成GPU的预热。

代码

完整测速代码如下所示，以torchvision自带的resnet101为例。

import torch
from torchvision.models.resnet import resnet101

iterations = 300   # 重复计算的轮次

model = resnet101()
device = torch.device("cuda:0")
model.to(device)

random_input = torch.randn(1, 3, 224, 224).to(device)
starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)

# GPU预热
for _ in range(50):
    _ = model(random_input)

# 测速
times = torch.zeros(iterations)     # 存储每轮iteration的时间
with torch.no_grad():
    for iter in range(iterations):
        starter.record()
        _ = model(random_input)
        ender.record()
        # 同步GPU时间
        torch.cuda.synchronize()
        curr_time = starter.elapsed_time(ender) # 计算时间
        times[iter] = curr_time
        # print(curr_time)

mean_time = times.mean().item()
print("Inference time: {:.6f}, FPS: {} ".format(mean_time, 1000/mean_time))

参考资料

[1] The Correct Way to Measure Inference Time of Deep Neural Networks
[2] 如何正确测量 PyTorch 模型的推理时间