(可选）将模型从 PyTorch 导出到 ONNX 并使用 ONNX Runtime 运行

原文： https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html

注意

单击此处的下载完整的示例代码

在本教程中，我们描述了如何将 PyTorch 中定义的模型转换为 ONNX 格式，然后在 ONNX Runtime 中运行它。

ONNX Runtime 是针对 ONNX 模型的以性能为中心的引擎，可在多个平台和硬件(Windows，Linux 和 Mac 以及 CPU 和 GPU 上）高效地进行推理。 如此处所述，事实证明，ONNX Runtime 可大大提高多个模型的性能。

对于本教程，您将需要安装 ONNX 和 ONNX Runtime 。 您可以使用pip install onnx onnxruntime获得 ONNX 和 ONNX Runtime 的二进制版本。 请注意，ONNX 运行时与 Python 3.5 至 3.7 版本兼容。

NOTE：本教程需要 PyTorch master 分支，可以按照此处的说明进行安装

  1. # Some standard imports
  2. import io
  3. import numpy as np
  5. from torch import nn
  6. import torch.utils.model_zoo as model_zoo
  7. import torch.onnx

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

首先，让我们在 PyTorch 中创建一个 SuperResolution 模型。 该模型使用了“使用高效的子像素卷积神经网络的实时单图像和视频超分辨率”(Shi 等人）中描述的有效的子像素卷积层来提高图像的分辨率 受高档因素的影响。 该模型期望输入图像的 YCbCr 的 Y 分量作为输入，并以超分辨率输出放大的 Y 分量。

模型直接来自 PyTorch 的示例，未经修改：

  1. # Super Resolution model definition in PyTorch
  2. import torch.nn as nn
  3. import torch.nn.init as init
  5. class SuperResolutionNet(nn.Module):
  6.     def __init__(self, upscale_factor, inplace=False):
  7.         super(SuperResolutionNet, self).__init__()
  9.         self.relu = nn.ReLU(inplace=inplace)
  10.         self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
  11.         self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
  12.         self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
  13.         self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
  14.         self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
  16.         self._initialize_weights()
  18.     def forward(self, x):
  19.         x = self.relu(self.conv1(x))
  20.         x = self.relu(self.conv2(x))
  21.         x = self.relu(self.conv3(x))
  22.         x = self.pixel_shuffle(self.conv4(x))
  23.         return x
  25.     def _initialize_weights(self):
  26.         init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
  27.         init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
  28.         init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
  29.         init.orthogonal_(self.conv4.weight)
  31. # Create the super-resolution model by using the above model definition.
  32. torch_model = SuperResolutionNet(upscale_factor=3)

通常，您现在将训练该模型。 但是，在本教程中，我们将下载一些预训练的权重。 请注意，此模型未经过充分训练以提供良好的准确性，在此仅用于演示目的。

在导出模型之前，请先调用torch_model.eval()或torch_model.train(False)，以将模型转换为推理模式，这一点很重要。 这是必需的，因为像 dropout 或 batchnorm 这样的运算符在推断和训练模式下的行为会有所不同。

  1. # Load pretrained model weights
  2. model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'
  3. batch_size = 1    # just a random number
  5. # Initialize model with the pretrained weights
  6. map_location = lambda storage, loc: storage
  7. if torch.cuda.is_available():
  8.     map_location = None
  9. torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))
  11. # set the model to inference mode
  12. 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 文档。

  1. # Input to the model
  2. x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)
  3. torch_out = torch_model(x)
  5. # Export the model
  6. torch.onnx.export(torch_model,               # model being run
  7.                   x,                         # model input (or a tuple for multiple inputs)
  8.                   "super_resolution.onnx",   # where to save the model (can be a file or file-like object)
  9.                   export_params=True,        # store the trained parameter weights inside the model file
  10.                   opset_version=10,          # the ONNX version to export the model to
  11.                   do_constant_folding=True,  # whether to execute constant folding for optimization
  12.                   input_names = ['input'],   # the model's input names
  13.                   output_names = ['output'], # the model's output names
  14.                   dynamic_axes={'input' : {0 : 'batch_size'},    # variable lenght axes
  15.                                 'output' : {0 : 'batch_size'}})

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

但是，在通过 ONNX Runtime 验证模型的输出之前，我们将使用 ONNX 的 API 检查 ONNX 模型。 首先，onnx.load("super_resolution.onnx")将加载保存的模型并输出 onnx.ModelProto 结构(用于捆绑 ML 模型的顶级文件/容器格式。有关更多信息，请参见 onnx.proto 文档。）。 然后，onnx.checker.check_model(onnx_model)将验证模型的结构并确认模型具有有效的架构。 通过检查模型的版本，图形的结构以及节点及其输入和输出，可以验证 ONNX 图的有效性。

  1. import onnx
  3. onnx_model = onnx.load("super_resolution.onnx")
  4. onnx.checker.check_model(onnx_model)

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

为了使用 ONNX Runtime 运行模型，我们需要使用所选的配置参数为模型创建一个推理会话(此处使用默认配置）。 创建会话后，我们将使用 run(）API 评估模型。 此调用的输出是一个列表，其中包含由 ONNX Runtime 计算的模型的输出。

  1. import onnxruntime
  3. ort_session = onnxruntime.InferenceSession("super_resolution.onnx")
  5. def to_numpy(tensor):
  6.     return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
  8. # compute ONNX Runtime output prediction
  9. ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
  10. ort_outs = ort_session.run(None, ort_inputs)
  12. # compare ONNX Runtime and PyTorch results
  13. np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)
  15. print("Exported model has been tested with ONNXRuntime, and the result looks good!")

我们应该看到 PyTorch 和 ONNX Runtime 的输出在数值上与给定的精度匹配(rtol = 1e-03 和 atol = 1e-05）。 附带说明一下，如果它们不匹配，则说明 ONNX 导出器中存在问题，因此请与我们联系。

使用 ONNX Runtime 在图像上运行模型

到目前为止，我们已经从 PyTorch 导出了一个模型，并演示了如何使用虚拟张量作为输入在 ONNX Runtime 中加载和运行该模型。

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

cat

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

我们首先调整图像大小以适合模型输入的大小(224x224）。 然后，我们将图像分为 Y，Cb 和 Cr 分量。 这些分量代表灰度图像(Y），蓝差(Cb）和红差(Cr）色度分量。 Y 分量对人眼更敏感，我们对将要转换的 Y 分量感兴趣。 提取 Y 分量后，我们将其转换为张量，这将是模型的输入。

  1. from PIL import Image
  2. import torchvision.transforms as transforms
  4. img = Image.open("./_static/img/cat.jpg")
  6. resize = transforms.Resize([224, 224])
  7. img = resize(img)
  9. img_ycbcr = img.convert('YCbCr')
  10. img_y, img_cb, img_cr = img_ycbcr.split()
  12. to_tensor = transforms.ToTensor()
  13. img_y = to_tensor(img_y)
  14. img_y.unsqueeze_(0)

现在，作为下一步，让我们使用代表灰度尺寸调整后的猫图像的张量，并按照先前的说明在 ONNX Runtime 中运行超分辨率模型。

  1. ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(img_y)}
  2. ort_outs = ort_session.run(None, ort_inputs)
  3. img_out_y = ort_outs[0]

此时，模型的输出为张量。 现在，我们将处理模型的输出，以根据输出张量构造最终的输出图像，并保存图像。 后处理步骤已从此处的超分辨率模型的 PyTorch 实现中采用。

  1. img_out_y = Image.fromarray(np.uint8((img_out_y[0] * 255.0).clip(0, 255)[0]), mode='L')
  3. # get the output image follow post-processing step from PyTorch implementation
  4. final_img = Image.merge(
  5.     "YCbCr", [
  6.         img_out_y,
  7.         img_cb.resize(img_out_y.size, Image.BICUBIC),
  8.         img_cr.resize(img_out_y.size, Image.BICUBIC),
  9.     ]).convert("RGB")
  11. # Save the image, we will compare this with the output image from mobile device
  12. final_img.save("./_static/img/cat_superres_with_ort.jpg")

output\_cat

ONNX Runtime 是跨平台引擎，您可以跨多个平台在 CPU 和 GPU 上运行它。

还可以使用 Azure 机器学习服务将 ONNX Runtime 部署到云中以进行模型推断。 更多信息此处。

在上了解有关 ONNX 运行时性能的更多信息。

有关 ONNX 运行时的更多信息，请参见。

脚本的总运行时间：(0 分钟 0.000 秒）

Download Python source code: super_resolution_with_onnxruntime.py Download Jupyter notebook: super_resolution_with_onnxruntime.ipynb

由狮身人面像画廊生成的画廊