基于ONNX結(jié)構(gòu)分析

Q

什么是ONNX？

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

Q

為什么要用ONNX？

深度學(xué)習(xí)算法大多通過(guò)計(jì)算數(shù)據(jù)流圖來(lái)完成神經(jīng)網(wǎng)絡(luò)的深度學(xué)習(xí)過(guò)程。一些框架（例如CNTK，Caffe2，Theano和TensorFlow）使用靜態(tài)圖形，而其他框架（例如PyTorch和Chainer）使用動(dòng)態(tài)圖形。但是這些框架都提供了接口，使開(kāi)發(fā)人員可以輕松構(gòu)建計(jì)算圖和運(yùn)行時(shí)，以優(yōu)化的方式處理圖。這些圖用作中間表示（IR），捕獲開(kāi)發(fā)人員源代碼的特定意圖，有助于優(yōu)化和轉(zhuǎn)換在特定設(shè)備（CPU，GPU，FPGA等）上運(yùn)行。假設(shè)一個(gè)場(chǎng)景：現(xiàn)在某組織因?yàn)橹饕_(kāi)發(fā)用TensorFlow為基礎(chǔ)的框架，現(xiàn)在有一個(gè)深度算法，需要將其部署在移動(dòng)設(shè)備上，以觀測(cè)變現(xiàn)。傳統(tǒng)地我們需要用Caffe2重新將模型寫(xiě)好，然后再訓(xùn)練參數(shù)；試想下這將是一個(gè)多么耗時(shí)耗力的過(guò)程。此時(shí)，ONNX便應(yīng)運(yùn)而生，Caffe2，PyTorch，Microsoft Cognitive Toolkit，Apache MXNet等主流框架都對(duì)ONNX有著不同程度的支持。這就便于我們的算法及模型在不同框架之間的遷移。

ONNX結(jié)構(gòu)分析

ONNX將每一個(gè)網(wǎng)絡(luò)的每一層或者說(shuō)是每一個(gè)算子當(dāng)作節(jié)點(diǎn)Node，再由這些Node去構(gòu)建一個(gè)Graph，相當(dāng)于是一個(gè)網(wǎng)絡(luò)。最后將Graph和這個(gè)ONNX模型的其他信息結(jié)合在一起，生成一個(gè)Model，也就是最終的.onnx的模型。構(gòu)建一個(gè)簡(jiǎn)單的ONNX模型，實(shí)質(zhì)上，只要構(gòu)建好每一個(gè)node，然后將它們和輸入輸出超參數(shù)一起塞到Graph，最后轉(zhuǎn)成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網(wǎng)絡(luò)結(jié)構(gòu)和參數(shù)（查看網(wǎng)址：https://netron.app/）

ONNX安裝、使用

安裝ONNX環(huán)境，在終端中執(zhí)行以下命令，環(huán)境中需要提前準(zhǔn)本 python3.6. 以下流程以u(píng)bunt 20.04 為例。

模型轉(zhuǎn)換流程

超分辨率是一種提高圖像、視頻分辨率的算法，廣泛用于圖像處理或視頻編輯。首先，讓我們?cè)赑yTorch中創(chuàng)建一個(gè)SuperResolution 模型。該模型使用描述的高效子像素卷積層將圖像的分辨率提高了一個(gè)放大因子。該模型將圖像的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

模型下載

由于本教程以演示為目的，因此采用下載預(yù)先訓(xùn)練好的權(quán)重。在導(dǎo)出模型之前調(diào)用torch_model.eval()或torch_model.train(False)將模型轉(zhuǎn)換為推理模式很重要。因?yàn)閐ropout或batchnorm等運(yùn)算符在推理和訓(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()

2

模型導(dǎo)出

要導(dǎo)出模型，我們調(diào)用該torch.onnx.export()函數(shù)。這將執(zhí)行模型，記錄用于計(jì)算輸出的運(yùn)算符。因?yàn)閑xport運(yùn)行模型，我們需要提供一個(gè)輸入張量x。只要它是正確的類型和大小，其中的值可以是隨機(jī)的。請(qǐng)注意，除非指定為動(dòng)態(tài)軸，否則所有輸入維度的導(dǎo)出ONNX圖中的輸入大小將是固定的。在此示例中，我們使用batch_size 1的輸入導(dǎo)出模型，但隨后在dynamic_axes參數(shù)中將第一個(gè)維度指定為動(dòng)態(tài) torch.onnx.export() . 因此，導(dǎo)出的模型將接受大小為[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

導(dǎo)出模型測(cè)試

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

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.加載處理前圖片，使用標(biāo)準(zhǔn)PIL python庫(kù)對(duì)其進(jìn)行預(yù)處理。 2.調(diào)整圖像大小以適應(yīng)模型輸入的大小 (224x224)。