将模型从PYTORCH导出到ONNX并使用ONNX RUNTIME运行

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

ONNX Runtime是针对ONNX模型的以性能为中心的引擎，可在多个平台和硬件（Windows，Linux和Mac以及CPU和GPU上）高效地进行推理。ONNX运行时已被证明大大增加了多种型号的性能，说明这里

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

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

# Some standard imports
import io
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 = 1    # 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的输入导出模型，但随后dynamic_axes 在torch.onnx.export()。因此，导出的模型将接受大小为[batch_size，1、224、224]的输入，其中batch_size可以是可变的。

要了解有关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 lenght axes
                                '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图的有效性。

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")
​
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 Runtime在图像上运行模型

 

Running the model on an image using ONNX Runtime

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

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

 

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

我们首先调整图像大小以适合模型输入的大小（224x224）。然后，我们将图像分为Y，Cb和Cr分量。这些分量表示灰度图像（Y），以及蓝差（Cb）和红差（Cr）色度分量。Y分量对人眼更敏感，我们对将要转换的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")

 

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

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

关于ONNX运行时的性能的详细信息在这里。

在此处获取有关ONNX Runtime的更多信息。

​

 

接下来，给大家介绍一下租用GPU做实验的方法，我们是在智星云租用的GPU，使用体验很好。具体大家可以参考：智星云官网： http://www.ai-galaxy.cn/，淘宝店：https://shop36573300.taobao.com/公众号: 智星AI