ONNXRuntime与PyTorch运行时间对比

        在网上看到一些对比ONNXRuntime与PyTorch运行效率的文章，很多人运行的结果展示出ONNX可以提升几倍的运行效率，那么究竟有没有这么神奇呢，我来试一试。

        系统：Ubuntu18.04

        CPU:    Intel8750H

        显卡：2070

        模型选择最常用的ResNet50

import torch
import torchvision.models as models

# import
model = models.resnet50(pretrained=True)

        首先将模型导出为onnx

# PyTorch model
torch.save(model, 'resnet.pth')

# random input
data = torch.rand(1,3,224,224)
# ONNX needs data example
torch.onnx.export(model, data, 'resnet.onnx')

        之后加载这两个模型

import onnxruntime

# PyTorch model
torch_model = torch.load('resnet.pth')
# ONNX model
onnx_model = onnxruntime.InferenceSession('resnet.onnx')
#torch_model.to("cuda:0")

        将模型多次运行，计算平均用时。

from timeit import timeit
import numpy as np

data = np.random.rand(1,3,224,224).astype(np.float32)
torch_data = torch.from_numpy(data)


def torch_inf():
    torch_model(torch_data)

def onnx_inf():
    onnx_model.run(None,{
                onnx_model.get_inputs()[0].name: data
           })
n = 200

#warmup
#for i in range(1,100):
#    torch_inf()
torch_t = timeit(lambda : torch_inf(), number=n)/n
onnx_t = timeit(lambda : onnx_inf(), number=n)/n

print(f"PyTorch {torch_t} VS ONNX {onnx_t}")

        我得到的效果是这样的

PyTorch 0.12086693297999773 VS ONNX 0.005529450080002789

        到这里，我们可以看到，ONNX的运行时间远远小于PyTorch。

        但是，这里有一个重要的问题。很多文章在这样测试的时候，忽略了model.eval()。这会让模型中一些inference时用不到的层也在执行。所以，我们在代码中加上

torch_model.eval()

        之后的输出就变成了

PyTorch 0.09768170792000092 VS ONNX 0.006109018884999386

        可以看到，效率差异没有那么的巨大了。那么如果使用GPU进行对比呢?

        我把onnxruntime换成onnxruntime-gpu，之后代码稍加改动

import onnxruntime

# PyTorch model
torch_model = torch.load('resnet.pth')
# ONNX model
onnx_model = onnxruntime.InferenceSession('resnet.onnx')
torch_model.to("cuda:0")
torch_data = torch.from_numpy(data).to("cuda:0")
torch_model.eval()

        此时的输出为

PyTorch 0.0061957750350029525 VS ONNX 0.006347028439995484

        那么，在GPU环境下，使用onnxruntime的时间和PyTorch基本是一致的。