Q

什么是ONNX？

ONNX（Open Neural Network Exchange）- 開放神經網絡交換格式，作為框架共用的一種模型交換格式，使用protobuf 二進制格式來序列化模型，可以提供更好的傳輸性能我們可能會在某一任務中Pytorch或者TensorFlow模型轉化為ONNX模型(ONNX模型一般用于中間部署階段)，然后再拿轉化后的ONNX模型進而轉化為我們使用不同框架部署需要的類型，ONNX相當于一個翻譯的作用。

Q

為什么要用ONNX？

深度學習算法大多通過計算數據流圖來完成神經網絡的深度學習過程。一些框架（例如CNTK，Caffe2，Theano和TensorFlow）使用靜態圖形，而其他框架（例如PyTorch和Chainer）使用動態圖形。但是這些框架都提供了接口，使開發人員可以輕松構建計算圖和運行時，以優化的方式處理圖。這些圖用作中間表示（IR），捕獲開發人員源代碼的特定意圖，有助于優化和轉換在特定設備（CPU，GPU，FPGA等）上運行。假設一個場景：現在某組織因為主要開發用TensorFlow為基礎的框架，現在有一個深度算法，需要將其部署在移動設備上，以觀測變現。傳統地我們需要用Caffe2重新將模型寫好，然后再訓練參數；試想下這將是一個多么耗時耗力的過程。此時，ONNX便應運而生，Caffe2，PyTorch，Microsoft Cognitive Toolkit，Apache MXNet等主流框架都對ONNX有著不同程度的支持。這就便于我們的算法及模型在不同框架之間的遷移。

ONNX結構分析

ONNX將每一個網絡的每一層或者說是每一個算子當作節點Node，再由這些Node去構建一個Graph，相當于是一個網絡。最后將Graph和這個ONNX模型的其他信息結合在一起，生成一個Model，也就是最終的.onnx的模型。構建一個簡單的ONNX模型，實質上，只要構建好每一個node，然后將它們和輸入輸出超參數一起塞到Graph，最后轉成Model就可以了。

graph{
    node{
        input: "1"
        input: "2"
        output: "12"
        op_type: "Conv"
    }
    attribute{
        name: "strides"
        ints: 1
        ints: 1
    }
    attribute{
        name: "pads"
        ints: 2
        ints: 2
    }
    ...
}

我們查看ONNX網絡結構和參數（查看網址：https://netron.app/）

ONNX安裝、使用

安裝ONNX環境，在終端中執行以下命令，環境中需要提前準本 python3.6. 以下流程以ubunt 20.04 為例。

模型轉換流程

超分辨率是一種提高圖像、視頻分辨率的算法，廣泛用于圖像處理或視頻編輯。首先，讓我們在PyTorch中創建一個SuperResolution 模型。該模型使用描述的高效子像素卷積層將圖像的分辨率提高了一個放大因子。該模型將圖像的YCbCr的Y分量作為輸入，并以超分辨率輸出放大的Y分量。

# Some standard imports
import io
import numpy as np


from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.onnx
# 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)

1

模型下載

由于本教程以演示為目的，因此采用下載預先訓練好的權重。在導出模型之前調用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()

2

模型導出

要導出模型，我們調用該torch.onnx.export()函數。這將執行模型，記錄用于計算輸出的運算符。因為export運行模型，我們需要提供一個輸入張量x。只要它是正確的類型和大小，其中的值可以是隨機的。請注意，除非指定為動態軸，否則所有輸入維度的導出ONNX圖中的輸入大小將是固定的。在此示例中，我們使用batch_size 1的輸入導出模型，但隨后在dynamic_axes參數中將第一個維度指定為動態 torch.onnx.export() . 因此，導出的模型將接受大小為[batch_size, 1, 224, 224]的輸入，其中batch_size可以是可變的。

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

3

導出模型測試

在使用ONNX Runtime驗證模型的輸出之前，我們將使用ONNX的 API檢查ONNX 模型。首先，onnx.load("super_resolution.onnx")將加載保存的模型并輸出 onnx.ModelProto結構（用于捆綁 ML 模型的頂級文件/容器格式）。然后，onnx.checker.check_model(onnx_model)將驗證模型的結構并確認模型具有有效的架構。ONNX 圖的有效性通過檢查模型的版本、圖的結構以及節點及其輸入和輸出來驗證。

import onnx


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


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!")

1.加載處理前圖片，使用標準PIL python庫對其進行預處理。 2.調整圖像大小以適應模型輸入的大小 (224x224)。