如何转换 PyTorch 模型并使用 OpenVINO™ 工具包运行它





5.00/5 (3投票s)
一个非常简单的指南,适用于所有希望提高推理性能的 PyTorch 开发人员。
注意:本文档是基于 OpenVINO 2022.1 创建的。如果您想了解如何使用 OpenVINO 2021.4 的旧 API,请查看 此笔记本。
虽然 PyTorch 是一个优秀的 AI 训练框架,并且可以用于推理,但 OpenVINO™ 工具包 在推理性能方面可以提供额外的优势,因为它针对此任务进行了高度优化。要使用它,您只需要 3 个简单的步骤:安装 OpenVINO、转换和优化您的模型以及运行推理。为了向您展示整个过程,我们决定使用 FastSeg 模型,这是一个针对 Cityscapes 数据集进行预训练的语义分割网络。
OpenVINO 能够运行 中间表示 (IR) 格式的网络进行推理。因此,您必须使用开发包中的命令行工具 模型优化器 将您的网络转换为 IR 格式。最简单的安装方法是使用 PyPi。
pip install openvino-dev[pytorch,onnx]
第一步是将模型导出为 ONNX 格式。您需要使用 Opset 版本 11,因为该版本受 OpenVINO 支持。此外,对于此特定模型,不允许进行常量折叠。
from fastseg import MobileV3Large
model = MobileV3Large.from_pretrained().cpu().eval()
dummy_input = torch.randn(1, 3, 512, 1024)
torch.onnx.export(model, dummy_input, "fastseg1024.onnx", opset_version=11, do_constant_folding=False)
OpenVINO 直接支持 ONNX,因此您可以将导出的模型加载到运行时并开始处理。但是,为了获得更好的性能,建议将模型转换为 IR 并进行优化。在终端中使用以下命令:
mo --input_model fastseg1024.onnx --input_shape "[1,3,512,1024]"
这意味着您正在将 model.onnx 转换为输入大小为 1024x512x3 (W x H x C) 的模型。当然,您可以提供其他参数,例如预处理或模型精度 (FP32 或 FP16)。
mo --input_model fastseg1024.onnx --input_shape "[1,3,512,1024]" --mean_values="[123.675,116.28,103.53]" --scale_values="[58.395,57.12,57.375]" --data_type FP16
均值的减去和除以标准差将直接构建到处理图中,并且推理将使用 FP16 运行。执行后,您应该会看到类似以下内容,其中包含所有显式和隐式参数,例如模型路径、输入形状、精度、均值和缩放值、转换参数等等。
Model Optimizer arguments:
Common parameters:
- Path to the Input Model: /home/adrian/repos/openvino_notebooks/notebooks/102-pytorch-onnx-to-openvino/fastseg.onnx
- Path for generated IR: /home/adrian/repos/openvino_notebooks/notebooks/102-pytorch-onnx-to-openvino/
- IR output name: fastseg
- Log level: ERROR
- Batch: Not specified, inherited from the model
- Input layers: Not specified, inherited from the model
- Output layers: Not specified, inherited from the model
- Input shapes: [1,3,512,1024]
- Source layout: Not specified
- Target layout: Not specified
- Layout: Not specified
- Mean values: [123.675,116.28,103.53]
- Scale values: [58.395,57.12,57.375]
- Scale factor: Not specified
- Precision of IR: FP16
- Enable fusing: True
- User transformations: Not specified
- Reverse input channels: False
- Enable IR generation for fixed input shape: False
- Use the transformations config file: None
Advanced parameters:
- Force the usage of legacy Frontend of Model Optimizer for model conversion into IR: False
- Force the usage of new Frontend of Model Optimizer for model conversion into IR: False
OpenVINO runtime found in: /home/adrian/repos/openvino_notebooks/venv/lib/python3.9/site-packages/openvino
OpenVINO runtime version: 2022.1.0-7019-cdb9bec7210-releases/2022/1
Model Optimizer version: 2022.1.0-7019-cdb9bec7210-releases/2022/1
[ SUCCESS ] Generated IR version 11 model.
[ SUCCESS ] XML file: /home/adrian/repos/openvino_notebooks/notebooks/102-pytorch-onnx-to-openvino/fastseg.xml
[ SUCCESS ] BIN file: /home/adrian/repos/openvino_notebooks/notebooks/102-pytorch-onnx-to-openvino/fastseg.bin
[ SUCCESS ] Total execution time: 0.55 seconds.
[ SUCCESS ] Memory consumed: 112 MB.
[ INFO ] The model was converted to IR v11, the latest model format that corresponds to the source DL framework input/output format. While IR v11 is backwards compatible with OpenVINO Inference Engine API v1.0, please use API v2.0 (as of 2022.1) to take advantage of the latest improvements in IR v11.
Find more information about API v2.0 and IR v11 at https://docs.openvino.ai
单词 SUCCESS 几乎在结尾处表示一切都已成功转换。您已获得 IR,它包含两个文件:.xml 和 .bin。现在您可以将此网络放入 OpenVINO™ 运行时并执行推理了。
import cv2
import numpy as np
from openvino.runtime import Coreimage_filename = "data/street.jpg"
image = cv2.cvtColor(cv2.imread(image_filename), cv2.COLOR_BGR2RGB)resized_image = cv2.resize(image, (1024, 512))
# Convert the resized images to network input shape
input_image = np.expand_dims(np.transpose(resized_image, (2, 0, 1)), 0)
# Load the network in Inference Engine
core = Core()
model_ir = core.read_model(model="fastseg.xml")
compiled_model_ir = core.compile_model(model=model_ir, device_name="CPU")
# Get output layer
output_layer_ir = compiled_model_ir.output(0)
# Run inference on the input image
res_ir = compiled_model_ir([input_image])[output_layer_ir]
result_mask_ir = np.squeeze(np.argmax(res_ir, axis=1)).astype(np.uint8)
成功了!下方的街道已被分割。您可以使用此 演示 亲自尝试。
或者,您可以使用此 工具 来查找针对您的环境下载和安装 OpenVINO™ 工具包的最佳方法。