使用 ONNX Runtime 在不同硬件目标上推理 PyTorch 模型

作为想要部署 PyTorch 或 ONNX 模型并最大化性能和硬件灵活性的开发人员，您可以利用 ONNX Runtime 在您的硬件平台上以最佳方式执行您的模型。

在本教程中，您将学习

  1. 如何使用 PyTorch ResNet-50 模型进行图像分类
  2. 转换为 ONNX，以及
  3. 使用 ONNX Runtime 部署到默认 CPU、NVIDIA CUDA (GPU) 和 Intel OpenVINO – 使用相同的应用程序代码在不同硬件平台上加载和执行推理。

ONNX 由微软、Meta、亚马逊和其他科技公司开发为开源 ML 模型格式，旨在标准化并简化在各种硬件类型上部署机器学习模型的过程。ONNX Runtime 由微软贡献和维护，旨在优化 ONNX 模型在 PyTorch、Tensorflow 等框架上的性能。ResNet-50 模型使用 ImageNet 数据集进行训练后，常用于图像分类。

本教程演示了如何使用 Microsoft Azure 机器学习，通过 ONNX Runtime 在 CPU、GPU 和 Intel 硬件（使用 OpenVINO）上运行 ONNX 模型。

设置

操作系统先决条件

您的环境应安装 curl。

设备先决条件

onnxruntime-gpu 库需要访问您设备或计算集群中的 NVIDIA CUDA 加速器，但仅在 CPU 上运行适用于 CPU 和 OpenVINO-CPU 演示。

推理先决条件

确保您有一张用于推理的图像。在本教程中，我们有一个 “cat.jpg” 图像，它与 Notebook 文件位于同一目录中。

环境先决条件

在 Azure Notebook 终端或 AnaConda 提示符窗口中，运行以下命令来创建您的 CPU、GPU 和/或 OpenVINO 的 3 个环境（差异以粗体显示）。

CPU

conda create -n cpu_env_demo python=3.8
conda activate cpu_env_demo
conda install -c anaconda ipykernel
conda install -c conda-forge ipywidgets
python -m ipykernel install --user --name=cpu_env_demo
jupyter notebook

GPU

conda create -n gpu_env_demo python=3.8
conda activate gpu_env_demo 
conda install -c anaconda ipykernel
conda install -c conda-forge ipywidgets
python -m ipykernel install --user --name=gpu_env_demo 
jupyter notebook

OpenVINO

conda create -n openvino_env_demo python=3.8
conda activate openvino_env_demo 
conda install -c anaconda ipykernel
conda install -c conda-forge ipywidgets
python -m ipykernel install --user --name=openvino_env_demo
python -m pip install --upgrade pip
pip install openvino

库要求

在第一个代码单元格中，使用以下代码片段安装必要的库（差异以粗体显示）。

CPU + GPU

import sys

if sys.platform in ['linux', 'win32']: # Linux or Windows
    !{sys.executable} -m pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio===0.9.0 -f https://download.pytorch.org/whl/torch_stable.html
else: # Mac
    print("PyTorch 1.9 MacOS Binaries do not support CUDA, install from source instead")

!{sys.executable} -m pip install onnxruntime-gpu onnx onnxconverter_common==1.8.1 pillow

OpenVINO

import sys

if sys.platform in ['linux', 'win32']: # Linux or Windows
    !{sys.executable} -m pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio===0.9.0 -f https://download.pytorch.org/whl/torch_stable.html
else: # Mac
    print("PyTorch 1.9 MacOS Binaries do not support CUDA, install from source instead")

!{sys.executable} -m pip install onnxruntime-openvino onnx onnxconverter_common==1.8.1 pillow

import openvino.utils as utils
utils.add_openvino_libs_to_path()

ResNet-50 演示

环境设置

导入必要的库以获取模型并运行推理。

from torchvision import models, datasets, transforms as T
import torch
from PIL import Image
import numpy as np

加载预训练的 ResNet-50 模型并导出到 ONNX

从 PyTorch 下载预训练的 ResNet-50 模型并导出为 ONNX 格式。

resnet50 = models.resnet50(pretrained=True)

# Download ImageNet labels
!curl -o imagenet_classes.txt https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt

# Read the categories
with open("imagenet_classes.txt", "r") as f:
    categories = [s.strip() for s in f.readlines()]

# Export the model to ONNX
image_height = 224
image_width = 224
x = torch.randn(1, 3, image_height, image_width, requires_grad=True)
torch_out = resnet50(x)
torch.onnx.export(resnet50,                     # model being run
                  x,                            # model input (or a tuple for multiple inputs)
                  "resnet50.onnx",              # where to save the model (can be a file or file-like object)
                  export_params=True,           # store the trained parameter weights inside the model file
                  opset_version=12,             # the ONNX version to export the model to
                  do_constant_folding=True,     # whether to execute constant folding for optimization
                  input_names = ['input'],      # the model's input names
                  output_names = ['output'])    # the model's output names

示例输出

% Total % Received % Xferd Average Speed Time Time Time Current
Dload Upload Total Spent Left Speed
100 10472 100 10472 0 0 50581 0 --:--:-- --:--:-- --:--:-- 50834

设置推理的预处理

为您想要使用模型进行推理的图像（例如 cat.jpg）创建预处理。

# Pre-processing for ResNet-50 Inferencing, from https://pytorch.ac.cn/hub/pytorch_vision_resnet/
resnet50.eval()  
filename = 'cat.jpg' # change to your filename

input_image = Image.open(filename)
preprocess = T.Compose([
    T.Resize(256),
    T.CenterCrop(224),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model

# move the input and model to GPU for speed if available
print("GPU Availability: ", torch.cuda.is_available())
if torch.cuda.is_available():
    input_batch = input_batch.to('cuda')
    resnet50.to('cuda')

示例输出

GPU Availability: False

使用 ONNX Runtime 推理 ResNet-50 ONNX 模型

通过为环境选择合适的执行提供程序，使用 ONNX Runtime 推理模型。如果您的环境使用 CPU，请取消注释 CPUExecutionProvider；如果环境使用 NVIDIA CUDA，请取消注释 CUDAExecutionProvider；如果环境使用 OpenVINOExecutionProvider，请取消注释 OpenVINOExecutionProvider – 注释掉其他 onnxruntime.InferenceSession 代码行。

# Inference with ONNX Runtime
import onnxruntime
from onnx import numpy_helper
import time

session_fp32 = onnxruntime.InferenceSession("resnet50.onnx", providers=['CPUExecutionProvider'])
# session_fp32 = onnxruntime.InferenceSession("resnet50.onnx", providers=['CUDAExecutionProvider'])
# session_fp32 = onnxruntime.InferenceSession("resnet50.onnx", providers=['OpenVINOExecutionProvider'])

def softmax(x):
    """Compute softmax values for each sets of scores in x."""
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum()

latency = []
def run_sample(session, image_file, categories, inputs):
    start = time.time()
    input_arr = inputs.cpu().detach().numpy()
    ort_outputs = session.run([], {'input':input_arr})[0]
    latency.append(time.time() - start)
    output = ort_outputs.flatten()
    output = softmax(output) # this is optional
    top5_catid = np.argsort(-output)[:5]
    for catid in top5_catid:
        print(categories[catid], output[catid])
    return ort_outputs

ort_output = run_sample(session_fp32, 'cat.jpg', categories, input_batch)
print("ONNX Runtime CPU/GPU/OpenVINO Inference time = {} ms".format(format(sum(latency) * 1000 / len(latency), '.2f')))

示例输出

Egyptian cat 0.78605634
tabby 0.117310025
tiger cat 0.020089425
Siamese cat 0.011728076
plastic bag 0.0052174763
ONNX Runtime CPU Inference time = 32.34 ms

与 PyTorch 的比较

使用 PyTorch 对比 ONNX Runtime CPU 和 GPU 的准确性和延迟。

# Inference with OpenVINO
from openvino.runtime import Core

ie = Core()
onnx_model_path = "./resnet50.onnx"
model_onnx = ie.read_model(model=onnx_model_path)
compiled_model_onnx = ie.compile_model(model=model_onnx, device_name="CPU")

# inference
output_layer = next(iter(compiled_model_onnx.outputs))

latency = []
input_arr = input_batch.detach().numpy()
inputs = {'input':input_arr}
start = time.time()
request = compiled_model_onnx.create_infer_request()
output = request.infer(inputs=inputs)

outputs = request.get_output_tensor(output_layer.index).data
latency.append(time.time() - start)

print("OpenVINO CPU Inference time = {} ms".format(format(sum(latency) * 1000 / len(latency), '.2f')))

print("***** Verifying correctness *****")
for i in range(2):
    print('OpenVINO and ONNX Runtime output {} are close:'.format(i), np.allclose(ort_output, outputs, rtol=1e-05, atol=1e-04))

示例输出

Egyptian cat 0.7820879
tabby 0.113261245
tiger cat 0.020114701
Siamese cat 0.012514038
plastic bag 0.0056432663
OpenVINO CPU Inference time = 31.83 ms
***** Verifying correctness *****
PyTorch and ONNX Runtime output 0 are close: True
PyTorch and ONNX Runtime output 1 are close: True

与 OpenVINO 的比较

使用 OpenVINO 对比 ONNX Runtime OpenVINO 的准确性和延迟。

# Inference with OpenVINO
from openvino.runtime import Core

ie = Core()
onnx_model_path = "./resnet50.onnx"
model_onnx = ie.read_model(model=onnx_model_path)
compiled_model_onnx = ie.compile_model(model=model_onnx, device_name="CPU")

# inference
output_layer = next(iter(compiled_model_onnx.outputs))

latency = []
input_arr = input_batch.detach().numpy()
inputs = {'input':input_arr}
start = time.time()
request = compiled_model_onnx.create_infer_request()
output = request.infer(inputs=inputs)

outputs = request.get_output_tensor(output_layer.index).data
latency.append(time.time() - start)

print("OpenVINO CPU Inference time = {} ms".format(format(sum(latency) * 1000 / len(latency), '.2f')))

print("***** Verifying correctness *****")
for i in range(2):
    print('OpenVINO and ONNX Runtime output {} are close:'.format(i), np.allclose(ort_output, outputs, rtol=1e-05, atol=1e-04))

示例输出

Egyptian cat 0.7820879
tabby 0.113261245
tiger cat 0.020114701
Siamese cat 0.012514038
plastic bag 0.0056432663
OpenVINO CPU Inference time = 31.83 ms
***** Verifying correctness *****
PyTorch and ONNX Runtime output 0 are close: True
PyTorch and ONNX Runtime output 1 are close: True

结论

我们已经证明，ONNX Runtime 是在 CPU、NVIDIA CUDA (GPU) 和 Intel OpenVINO (移动设备) 上运行 PyTorch 或 ONNX 模型的有效方法。 ONNX Runtime 支持部署到更多类型的硬件，这些硬件可以在执行提供程序中找到。我们很乐意通过参与我们的 ONNX Runtime Github 仓库来听取您的反馈。

视频演示

观看此处的视频，详细了解 ResNet-50 部署和灵活推理的分步指南。