（可选）将模型从 PyTorch 导出到 ONNX 并使用 ONNX 运行时运行它

创建时间： 2019年7月17日 |上次更新时间：2024 年 7 月 17 日 |上次验证： Nov 05， 2024

注意

从 PyTorch 2.1 开始，ONNX 导出器有两个版本。

• torch.onnx.dynamo_export是基于 PyTorch 2.0 发布的 TorchDynamo 技术的最新（仍处于测试阶段）导出器。

• torch.onnx.export基于 TorchScript 后端，自 PyTorch 1.2.0 起可用。

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

导出的模型将使用 ONNX 运行时执行。 ONNX 运行时是一个以性能为中心的 ONNX 模型引擎， 跨多个平台和硬件高效推理 （Windows、Linux 和 Mac 以及 CPU 和 GPU）。 事实证明，ONNX 运行时的性能大大超过 多个模型，如此处

在本教程中，您需要安装 ONNX 和 ONNX 运行时。 您可以使用

%%bash
pip install onnx onnxruntime

ONNX 运行时建议为 PyTorch 使用最新的稳定运行时。

# Some standard imports
import numpy as np

from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.onnx

超分辨率是一种提高图像、视频分辨率的方法 并广泛用于图像处理或视频编辑。对于这个 教程中，我们将使用一个小的超分辨率模型。

首先，让我们在 PyTorch 中创建一个模型。 该模型使用“使用高效 Sub-Pixel Convolutional Neural Network“ - Shi 等人将图像的分辨率提高了一个放大因子。 模型需要图像的 Y 分量作为输入，并且 以超分辨率输出放大的 Y 分量。SuperResolutionYCbCr

这 model 直接来自 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)

通常，您现在将训练此模型;但是，对于本教程， 我们将下载一些预训练的权重。请注意，此模型 没有经过充分的训练以获得良好的准确性，此处用于 仅用于演示目的。

请务必调用 或在导出模型之前将模型转换为推理模式。 这是必需的，因为 dropout 或 batchnorm 等运算符的行为 在推理和训练模式下有所不同。torch_model.eval()torch_model.train(False)

# 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 中导出模型是通过跟踪或脚本编写的。这 教程将使用 Tracing 导出的模型作为示例。 要导出模型，我们调用 function。 这将执行模型，记录 what 运算符的跟踪 用于计算输出。 因为运行模型，所以我们需要提供一个输入 张肌。this 中的值可以是随机的，只要它是 正确的类型和大小。 请注意，输入大小将在导出的 ONNX 图中固定 所有输入的维度，除非指定为 Dynamic Axes（动态轴）。 在此示例中，我们导出输入为 batch_size 1 的模型。 但随后在 的参数中将第一个维度指定为 dynamic。 因此，导出的模型将接受大小为 [batch_size， 1， 224， 224] 的输入 其中 batch_size 可以是可变的。torch.onnx.export()exportxdynamic_axestorch.onnx.export()

要了解有关 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'}})

我们还计算了 ，模型之后的输出， 我们将使用它来验证我们导出的模型是否计算 在 ONNX 运行时运行时的值相同。torch_out

但在使用 ONNX 运行时验证模型的输出之前，我们将检查 使用 ONNX API 的 ONNX 模型。 首先，将加载保存的模型并 将输出一个结构（用于捆绑 ML 模型的顶级文件/容器格式。 有关更多信息，请访问 nx.proto 文档。 然后，将验证模型的结构 并确认模型具有有效的 schema。 通过检查模型的 version、图形的结构以及节点及其输入 和输出。onnx.load("super_resolution.onnx")onnx.ModelProtoonnx.checker.check_model(onnx_model)

import onnx

onnx_model = onnx.load("super_resolution.onnx")
onnx.checker.check_model(onnx_model)

现在，让我们使用 ONNX Runtime 的 Python API 计算输出。 这部分通常可以在单独的流程中完成，也可以在另一个流程中完成 机器，但我们将继续相同的过程，以便我们可以 验证 ONNX 运行时和 PyTorch 是否正在计算相同的值 对于网络。

为了使用 ONNX 运行时运行模型，我们需要创建一个 具有所选配置的模型的推理会话 parameters （这里我们使用 default config）。 创建会话后，我们使用 run（） API 评估模型。 此调用的输出是一个包含模型输出的列表 由 ONNX 运行时计算。

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 运行时运行的输出匹配 以给定的精度 （ 和 ） 进行数值运算。 顺便说一句，如果它们不匹配，则 ONNX 导出器，因此在这种情况下请联系我们。rtol=1e-03atol=1e-05

模型之间的时序比较

由于 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 运行时在图像上运行模型

到目前为止，我们已经从 PyTorch 导出了一个模型，并展示了如何加载它 并在 ONNX 运行时中运行它，并将虚拟张量作为输入。

在本教程中，我们将使用一个被广泛使用的著名猫图像，它 如下所示

首先，让我们加载图像，使用标准 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)

现在，作为下一步，让我们以表示 灰度调整 CAT 图像的大小并在 如前所述，ONNX 运行时。

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]

此时，模型的输出是一个张量。 现在，我们将处理模型的输出以构造回 output tensor 的最终输出图像，然后保存图像。 后处理步骤已采用 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 运行时是一个跨平台引擎，您可以在 多个平台以及 CPU 和 GPU 上。

ONNX 运行时也可以部署到云中进行模型推理 使用 Azure 机器学习服务。更多信息请点击这里。

有关 ONNX 运行时性能的更多信息，请点击此处。

有关 ONNX 运行时的更多信息，请点击此处。