PyTorch (可選）將模型從 PyTorch 導(dǎo)出到 ONNX 并使用 ONNX Runtime 運行

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

在本教程中，我們描述了如何將 PyTorch 中定義的模型轉(zhuǎn)換為 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 分支，可以按照此處的說明進行安裝

# 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 中創(chuàng)建一個 SuperResolution 模型。 該模型使用了“使用高效的子像素卷積神經(jīng)網(wǎng)絡(luò)的實時單圖像和視頻超分辨率”(Shi 等人）中描述的有效的子像素卷積層來提高圖像的分辨率 受高檔因素的影響。 該模型期望輸入圖像的 YCbCr 的 Y 分量作為輸入，并以超分辨率輸出放大的 Y 分量。

模型直接來自 PyTorch 的示例，未經(jīng)修改：

# 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)

通常，您現(xiàn)在將訓(xùn)練該模型。 但是，在本教程中，我們將下載一些預(yù)訓(xùn)練的權(quán)重。 請注意，此模型未經(jīng)過充分訓(xùn)練以提供良好的準確性，在此僅用于演示目的。

在導(dǎo)出模型之前，請先調(diào)用torch_model.eval()或torch_model.train(False)，以將模型轉(zhuǎn)換為推理模式，這一點很重要。 這是必需的，因為像 dropout 或 batchnorm 這樣的運算符在推斷和訓(xùn)練模式下的行為會有所不同。

# 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 中導(dǎo)出模型是通過跟蹤或腳本編寫的。 本教程將以通過跟蹤導(dǎo)出的模型為例。 要導(dǎo)出模型，我們調(diào)用torch.onnx.export()函數(shù)。 這將執(zhí)行模型，并記錄使用什么運算符計算輸出的軌跡。 因為export運行模型，所以我們需要提供輸入張量x。 只要是正確的類型和大小，其中的值就可以是隨機的。 請注意，除非指定為動態(tài)軸，否則輸入尺寸將在導(dǎo)出的 ONNX 圖形中固定為所有輸入尺寸。 在此示例中，我們使用輸入 batch_size 1 導(dǎo)出模型，但隨后在torch.onnx.export()的dynamic_axes參數(shù)中將第一維指定為動態(tài)。 因此，導(dǎo)出的模型將接受大小為[batch_size，1、224、224]的輸入，其中 batch_size 可以是可變的。

要了解有關(guān) PyTorch 導(dǎo)出界面的更多詳細信息，請查看 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，即模型之后的輸出，我們將用來驗證導(dǎo)出的模型在 ONNX Runtime 中運行時是否計算出相同的值。

但是，在通過 ONNX Runtime 驗證模型的輸出之前，我們將使用 ONNX 的 API 檢查 ONNX 模型。 首先，onnx.load("super_resolution.onnx")將加載保存的模型并輸出 onnx.ModelProto 結(jié)構(gòu)(用于捆綁 ML 模型的頂級文件/容器格式。有關(guān)更多信息，請參見 onnx.proto 文檔。）。 然后，onnx.checker.check_model(onnx_model)將驗證模型的結(jié)構(gòu)并確認模型具有有效的架構(gòu)。 通過檢查模型的版本，圖形的結(jié)構(gòu)以及節(jié)點及其輸入和輸出，可以驗證 ONNX 圖的有效性。

import onnx


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

現(xiàn)在，我們使用 ONNX 運行時的 Python API 計算輸出。 這部分通?？梢栽趩为毜倪^程中或在另一臺機器上完成，但是我們將繼續(xù)同一過程，以便我們可以驗證 ONNX Runtime 和 PyTorch 正在為網(wǎng)絡(luò)計算相同的值。

為了使用 ONNX Runtime 運行模型，我們需要使用所選的配置參數(shù)為模型創(chuàng)建一個推理會話(此處使用默認配置）。 創(chuàng)建會話后，我們將使用 run(）API 評估模型。 此調(diào)用的輸出是一個列表，其中包含由 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!")

我們應(yīng)該看到 PyTorch 和 ONNX Runtime 的輸出在數(shù)值上與給定的精度匹配(rtol = 1e-03 和 atol = 1e-05）。 附帶說明一下，如果它們不匹配，則說明 ONNX 導(dǎo)出器中存在問題，因此請與我們聯(lián)系。

使用 ONNX Runtime 在圖像上運行模型

到目前為止，我們已經(jīng)從 PyTorch 導(dǎo)出了一個模型，并演示了如何使用虛擬張量作為輸入在 ONNX Runtime 中加載和運行該模型。

在本教程中，我們將使用廣泛使用的著名貓圖像，如下圖所示

首先，讓我們加載圖片，使用標準 PIL python 庫對其進行預(yù)處理。 請注意，此預(yù)處理是處理數(shù)據(jù)以訓(xùn)練/測試神經(jīng)網(wǎng)絡(luò)的標準做法。

我們首先調(diào)整圖像大小以適合模型輸入的大小(224x224）。 然后，我們將圖像分為 Y，Cb 和 Cr 分量。 這些分量代表灰度圖像(Y），藍差(Cb）和紅差(Cr）色度分量。 Y 分量對人眼更敏感，我們對將要轉(zhuǎn)換的 Y 分量感興趣。 提取 Y 分量后，我們將其轉(zhuǎn)換為張量，這將是模型的輸入。

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)

現(xiàn)在，作為下一步，讓我們使用代表灰度尺寸調(diào)整后的貓圖像的張量，并按照先前的說明在 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]

此時，模型的輸出為張量。 現(xiàn)在，我們將處理模型的輸出，以根據(jù)輸出張量構(gòu)造最終的輸出圖像，并保存圖像。 后處理步驟已從此處的超分辨率模型的 PyTorch 實現(xiàn)中采用。

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 機器學(xué)習(xí)服務(wù)將 ONNX Runtime 部署到云中以進行模型推斷。 更多信息此處。

在上了解有關(guān) ONNX 運行時性能的更多信息。

有關(guān) ONNX 運行時的更多信息，請參見。

腳本的總運行時間：(0 分鐘 0.000 秒）

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