(可选) 将Pytorch模型导出到ONNX并使用ONNX Runtime运行

创建于: 2019年7月17日 | 最后更新于: 2024年7月17日 | 最后验证于: 2024年11月5日

注意

截至 PyTorch 2.1，有兩個版本的 ONNX 导出器。

• torch.onnx.dynamo_export 是基于 PyTorch 2.0 发布的 TorchDynamo 技术的新版本导出工具（仍在测试中）。

• torch.onnx.export 是基于 TorchScript 后端，并且自 PyTorch 1.2.0 版本起就可用。

在本教程中，我们将描述如何使用TorchScript torch.onnx.export ONNX导出器将定义在PyTorch中的模型转换为ONNX格式。

导出的模型将使用ONNX Runtime执行。 ONNX Runtime是一个专注于性能的ONNX模型引擎， 在多个平台和硬件（Windows、Linux、Mac以及CPU和GPU）上高效推理。 如这里所述，ONNX Runtime已被证明可以显著提高多个模型的性能。

对于本教程，您需要安装ONNX 和ONNX 运行时。 您可以从以下链接获取 ONNX 和 ONNX 运行时的二进制构建：

%%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等人在 “使用高效的子像素卷积神经网络进行实时单幅图像和视频超分辨率” 中描述的高效子像素卷积层来提高图像的分辨率。 该模型期望输入图像的YCbCr的Y分量，并输出超分辨率的放大Y分量。

模型直接来自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()

Exporting 一个模型在 PyTorch 中通过追踪或脚本化完成。本教程将使用通过追踪导出的模型作为示例。 要导出一个模型，我们调用 torch.onnx.export() 函数。 这将执行模型，并记录用于计算输出的操作符的踪迹。 由于 export 执行了模型，因此我们需要提供一个输入张量 x。只要这个张量的类型和大小正确，其值可以是随机的。 请注意，在导出的 ONNX 图中，除非指定了动态轴，否则所有输入的维度大小将固定不变。 在这个示例中，我们使用批量大小为 1 的输入导出模型，但在 dynamic_axes 参数中的 torch.onnx.export() 函数中指定了第一个维度为动态。 因此，导出的模型可以接受大小为 [批量大小, 1, 224, 224] 的输入，其中批量大小可以变化。

要了解更多关于PyTorch导出接口的详细信息，请查看 torch.onnx文档。

# 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 结构（一种用于捆绑机器学习模型的顶级文件/容器格式）。有关更多信息，请参阅 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倍的速度。随着批量大小的增大，这种改进会更加明显。

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中加载并运行该模型。

对于这个教程，我们将使用一张广泛使用的著名猫图片，看起来如下所示：

首先，让我们加载图像，并使用标准的PIL Python库对其进行预处理。请注意，这种预处理是训练/测试神经网络时处理数据的标准做法。

我们首先将图像调整为模型输入大小（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的超分辨率模型实现中采用的 这里。

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 Machine Learning Services。更多信息 这里。

关于ONNX Runtime性能的更多信息 在这里。

有关ONNX Runtime的更多信息，请点击这里。