ONNXRuntime与PyTorch哪个更快？

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

        系统：Ubuntu22.04

        CPU:    Intel8750H

        显卡：3060

        模型选择最常用的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基本是一致的。 

文章转载自一个机器视觉工程师的自我修养_C#,RealSense+CSharp,onnxruntime-CSDN博客

 作为学习记录