备注
点击 这里 下载完整示例代码
(可选) 从PyTorch导出模型到ONNX并使用ONNX运行时运行¶
Created On: Jul 17, 2019 | Last Updated: Jul 17, 2024 | Last Verified: Nov 05, 2024
备注
截至PyTorch 2.1，有两个版本的ONNX导出器。
在本教程中，我们介绍了如何使用TorchScript``torch.onnx.export`` ONNX导出器将PyTorch定义的模型转换为ONNX格式。
导出的模型将使用ONNX Runtime执行。ONNX Runtime是一个针对ONNX模型的性能优化引擎，可以在多个平台和硬件上高效推断（Windows、Linux、Mac，以及CPU和GPU）。正如`这里 <https://cloudblogs.microsoft.com/opensource/2019/05/22/onnx-runtime-machine-learning-inferencing-0-4-release>`__所解释的，ONNX Runtime已证明在多个模型上大幅提高性能。
本教程需要安装`ONNX <https://github.com/onnx/onnx>`__和`ONNX Runtime <https://github.com/microsoft/onnxruntime>`__。您可以通过以下方式获取ONNX和ONNX Runtime的二进制版本：
%%bash
pip install onnx onnxruntime
ONNX Runtime建议为PyTorch使用最新稳定版运行时。
# Some standard imports
import numpy as np
from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.onnx
超分辨率是一种增加图像和视频分辨率的方法，广泛用于图像处理或视频编辑。在本教程中，我们将使用一个小型的超分辨率模型。
首先，让我们在PyTorch中创建一个``SuperResolution``模型。该模型使用一种高效的子像素卷积层（详见`”实时单图像和视频超分辨率使用高效子像素卷积神经网络” - Shi等人 <https://arxiv.org/abs/1609.05158>`__)，通过一个放大因子增加图像的分辨率。该模型以图像的``YCbCr``的Y分量为输入，并输出经过超分辨率处理的Y分量。
`该模型 <https://github.com/pytorch/examples/blob/master/super_resolution/model.py>`__直接来自于PyTorch的示例，未作修改：
# Super Resolution model definition in PyTorch
import torch.nn as nn
import torch.nn.init as init
class SuperResolutionNet(nn.Module):
def __init__(self, upscale_factor, inplace=False):
super(SuperResolutionNet, self).__init__()
self.relu = nn.ReLU(inplace=inplace)
self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
self._initialize_weights()
def forward(self, x):
x = self.relu(self.conv1(x))
x = self.relu(self.conv2(x))
x = self.relu(self.conv3(x))
x = self.pixel_shuffle(self.conv4(x))
return x
def _initialize_weights(self):
init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
init.orthogonal_(self.conv4.weight)
# Create the super-resolution model by using the above model definition.
torch_model = SuperResolutionNet(upscale_factor=3)
通常情况下，您现在会训练该模型；然而，为了本教程的目的，我们将改为下载一些预训练权重。请注意，该模型并未完全训练以获得良好的准确性，此处仅为演示使用。
导出模型前调用``torch_model.eval()``或``torch_model.train(False)``非常重要，以将模型设置为推断模式。这是因为诸如dropout或batchnorm等算子在推断模式和训练模式下的行为不同。
# Load pretrained model weights
model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'
batch_size = 64 # just a random number
# Initialize model with the pretrained weights
map_location = lambda storage, loc: storage
if torch.cuda.is_available():
map_location = None
torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))
# set the model to inference mode
torch_model.eval()
在PyTorch中可以通过追踪或脚本方式导出模型。本教程将以追踪方式导出的模型作为示例。要导出模型，我们调用``torch.onnx.export()``函数。该函数将执行模型，记录计算输出所使用的算子。由于``export``会运行模型，我们需要提供一个输入张量``x``。只要类型和尺寸正确，该张量的值可以是随机的。请注意，除非指定为动态轴，导出的ONNX图将为所有输入维度的尺寸固定输入大小。在本示例中，我们将使用batch_size为1的输入导出模型，但随后在``torch.onnx.export()``里的``dynamic_axes``参数中指定第一个维度为动态。因此，导出的模型将接受尺寸为[batch_size, 1, 224, 224]的输入，其中batch_size可以是可变的。
要了解有关PyTorch导出接口的更多细节，请查看`torch.onnx文档 <https://pytorch.org/docs/master/onnx.html>`__。
# Input to the model
x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)
torch_out = torch_model(x)
# Export the model
torch.onnx.export(torch_model, # model being run
x, # model input (or a tuple for multiple inputs)
"super_resolution.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=10, # 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
dynamic_axes={'input' : {0 : 'batch_size'}, # variable length axes
'output' : {0 : 'batch_size'}})
我们还计算了``torch_out``，模型的输出，我们将用它来验证导出的模型在ONNX Runtime中计算的是相同值。
在使用ONNX Runtime验证模型的输出之前，我们将使用ONNX API检查ONNX模型。首先，``onnx.load(“super_resolution.onnx”)``将加载保存的模型并输出一个``onnx.ModelProto``结构（一种用于打包ML模型的顶级文件/容器格式。有关更多信息，请参考`onnx.proto文档 <https://github.com/onnx/onnx/blob/master/onnx/onnx.proto>`__）。然后，``onnx.checker.check_model(onnx_model)``将验证模型的结构，并确认模型具有有效的模式。ONNX图的有效性通过检查模型版本、图结构以及节点及其输入和输出来验证。
import onnx
onnx_model = onnx.load("super_resolution.onnx")
onnx.checker.check_model(onnx_model)
现在，让我们使用ONNX Runtime的Python API计算输出。这部分通常可以在单独的进程或另一台机器上完成，但我们将继续在同一进程中进行，以验证ONNX Runtime和PyTorch正在为网络计算相同的值。
为了使用ONNX Runtime运行模型，我们需要根据选择的配置参数（此处使用默认配置）为模型创建推断会话。一旦会话被创建，我们使用run() API对模型进行评估。此调用的输出是一个列表，包含由ONNX Runtime计算的模型输出。
import onnxruntime
ort_session = onnxruntime.InferenceSession("super_resolution.onnx", providers=["CPUExecutionProvider"])
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
# compute ONNX Runtime output prediction
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)
# compare ONNX Runtime and PyTorch results
np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)
print("Exported model has been tested with ONNXRuntime, and the result looks good!")
我们应该能够看到PyTorch和ONNX Runtime运行的输出在给定精度下数值匹配（rtol=1e-03``和``atol=1e-05）。作为附注，如果它们不匹配，则表示ONNX导出器存在问题，请在这种情况下与我们联系。
模型之间的时间对比¶
由于ONNX模型针对推断速度进行了优化，使用ONNX模型而不是原生PyTorch模型运行相同数据应该会带来高达2倍的提升。提升在较大的batch_size中更为显著。
import time
x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)
start = time.time()
torch_out = torch_model(x)
end = time.time()
print(f"Inference of Pytorch model used {end - start} seconds")
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
start = time.time()
ort_outs = ort_session.run(None, ort_inputs)
end = time.time()
print(f"Inference of ONNX model used {end - start} seconds")
使用ONNX Runtime运行模型时加载图像¶
目前我们已经从PyTorch导出了一个模型，并展示了如何加载它以及使用一个虚拟张量作为输入在ONNX Runtime中运行它。
本教程将使用一张著名的猫图片，该图片被广泛使用，效果如下
首先，让我们加载图像，并使用标准的Python库PIL进行预处理。请注意，这种预处理是用于训练/测试神经网络的标准数据处理方法。
我们首先将图像调整为与模型输入尺寸（224x224）相符。然后，我们将图像分解为其Y、Cb和Cr分量。这些分量分别表示灰度图像（Y）以及蓝色差异（Cb）和红色差异（Cr）的色度分量。由于Y分量对人眼更为敏感，我们关注这个分量并对其进行转化。在提取Y分量后，我们将其转换为张量，它将作为模型的输入。
from PIL import Image
import torchvision.transforms as transforms
img = Image.open("./_static/img/cat.jpg")
resize = transforms.Resize([224, 224])
img = resize(img)
img_ycbcr = img.convert('YCbCr')
img_y, img_cb, img_cr = img_ycbcr.split()
to_tensor = transforms.ToTensor()
img_y = to_tensor(img_y)
img_y.unsqueeze_(0)
接下来，让我们将表示灰度调整后的猫图像的张量输入到ONNX Runtime中的超分辨率模型，并按照之前的解释运行。
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(img_y)}
ort_outs = ort_session.run(None, ort_inputs)
img_out_y = ort_outs[0]
此时，模型的输出是一个张量。现在，我们将处理模型的输出，以从输出张量构造最终的输出图像，并保存该图像。这些后处理步骤采用了PyTorch超分辨率模型的实现`这里 <https://github.com/pytorch/examples/blob/master/super_resolution/super_resolve.py>`__。
img_out_y = Image.fromarray(np.uint8((img_out_y[0] * 255.0).clip(0, 255)[0]), mode='L')
# get the output image follow post-processing step from PyTorch implementation
final_img = Image.merge(
"YCbCr", [
img_out_y,
img_cb.resize(img_out_y.size, Image.BICUBIC),
img_cr.resize(img_out_y.size, Image.BICUBIC),
]).convert("RGB")
# Save the image, we will compare this with the output image from mobile device
final_img.save("./_static/img/cat_superres_with_ort.jpg")
# Save resized original image (without super-resolution)
img = transforms.Resize([img_out_y.size[0], img_out_y.size[1]])(img)
img.save("cat_resized.jpg")
以下是两张图片的对比：
低分辨率图像
超分辨率处理后的图像
ONNX Runtime作为一个跨平台引擎，可以在多个平台上运行，并支持CPU和GPU。
ONNX Runtime还可以通过Azure机器学习服务部署到云端进行模型推断。更多信息请参考`这里 <https://docs.microsoft.com/en-us/azure/machine-learning/service/concept-onnx>`__。
关于ONNX Runtime性能的更多信息请参考`这里 <https://onnxruntime.ai/docs/performance>`__。
关于ONNX Runtime的更多信息请参考`这里 <https://github.com/microsoft/onnxruntime>`__。
