ONNX 淺析：如何加速深度學習算法工程化？

以下文章來源于拍樂云Pano ，作者拍樂云視頻專家

AlphaGo擊敗圍棋世界冠軍李世石以來，關(guān)于人工智能和深度學習的研究呈現(xiàn)井噴之勢。

各種新的算法和網(wǎng)絡(luò)模型層出不窮，令人眼花繚亂。與之相隨的，深度學習的訓練和推理框架也在不斷的推陳出新，比較著名的有：微軟的CNTK、Google的TensorFlow、Facebook的PyTorch、Apple的CoreML、Intel 的OpenVINO、英偉達的cuDNN和TensorRT、騰訊的TNN和NCNN、阿里的MNN等等。

這些框架都有相似之處，他們的輸入是一個或者一組多維數(shù)據(jù)，數(shù)據(jù)經(jīng)過多層運算單元之后輸出運算結(jié)果。訓練框架支持BackPropogation等訓練方法，可以自動調(diào)整模型參數(shù)，用于算法設(shè)計。推理框架則是單純只能用于模型推理運算，不能調(diào)整模型本身，用于算法部署落地。

這些框架中，Google的TensorFlow的這個名字尤其具有美感。多維數(shù)據(jù)是為張量(Tensor)，數(shù)據(jù)在多層運算單元中的運算和傳遞是為流(FLow)，看到這個詞就仿佛看到了一個數(shù)據(jù)和運算的圖(Computation Graph)，真可謂妙手偶得之佳作。這些框架都需要構(gòu)建算子，并且將這些算子按照一定的次序連接起來，可以稱之為網(wǎng)絡(luò)模型。

01 Why ONNX？

每個深度學習框架都有自己獨有的格式來解釋和存儲網(wǎng)絡(luò)模型，并且這些框架的側(cè)重點不同，有些用于訓練學習，有些用于部署推理。在深度學習算法開發(fā)中，在不同的階段會選擇不同的框架，所以模型描述格式的不同，在客觀上造成了深度學習算法開發(fā)和落地的困難。

筆者之前曾開發(fā)深度神經(jīng)網(wǎng)絡(luò)算法，當時選擇的訓練框架是Caffe，需要落地部署到Linux、iOS、Android等多個平臺。Linux選擇的是Nvidia的cuDNN；iOS選擇的是CoreML；Android選擇的是NNAPI，Caffe的模型描述格式是caffemodel。

它使用自定義的Protobuf (https://github.com/BVLC/caffe/tree/master/src/caffe/proto)，但是顯然，無論是cuDNN、CoreML、NNAPI都無法直接使用caffemodel，CoreML的模型描述使用另一種定義 (https://apple.github.io/coremltools/mlmodel/index.html)，cuDNN和NNAPI都是low-level的推理引擎，需要使用者將這個模型組裝起來。

對于CoreML來說，我們需要把caffemodel轉(zhuǎn)為coremlmodel格式，對于 cuDNN和NNAPI，我們需要解析caffemodel，然后自己組裝出完整的網(wǎng)絡(luò)模型。這個過程繁瑣而且容易出錯，當時有強烈的沖動，希望定義一個統(tǒng)一的模型描述格式，所有的訓練框架訓練所得的網(wǎng)絡(luò)模型，都是用這個格式來描述，在設(shè)備上部署推理時，相應的推理引擎支持解析這個統(tǒng)一的描述格式，直接完成部署落地，豈不美哉。

當然此事并不容易，要定義個統(tǒng)一的模型描述格式，不僅僅需要對機器學習技術(shù)有深入的理解，而且將之推廣成為事實上的行業(yè)標準，更需要有很大的行業(yè)影響力，并不是如筆者這樣的無名小卒可以為之。所幸已經(jīng)有社區(qū)在做這個事情了，這就是Open Neural Network Exchange(ONNX)。

用ONNX自己的話來說，ONNX是一個對計算網(wǎng)絡(luò)(Computation Graph)的一個通用描述(Intermediate Representation)。它希望被設(shè)計成為開放的網(wǎng)絡(luò)描述格式，減少開發(fā)者在各個訓練和推理框架間切換的代價，可以讓開發(fā)者專注于算法和優(yōu)化。雖然ONNX還處于比較早期的階段，不過已經(jīng)有約來越多的人開始關(guān)注到它，未來會有機會得到更廣泛的應用。

02 計算模型描述

ONNX有兩個分類：基礎(chǔ)的ONNX主要用于描述神經(jīng)網(wǎng)絡(luò)、ONNX-ML是對基礎(chǔ)ONNX的擴展，增加了神經(jīng)網(wǎng)絡(luò)之外的其他機器學習算法的支持。本文不會涉及ONNX-ML，接下來的文字以一個簡單的ONNX模型為例，介紹一下 ONNX是如何來描述一個計算網(wǎng)絡(luò)的。該模型可以在ONNX的Github上下載(https://github.com/onnx/models/blob/master/vision/classification/mobilenet/model/mobilenetv2-7.onnx).

ONNX的模型描述采用了Google的Protocol Buffer語言。最外層的結(jié)構(gòu)是ModelProto，它的定義如下：

message ModelProto {
  int64 ir_version = 1;
  repeated OperatorSetIdProto opset_import = 8;
  string producer_name = 2;
  string producer_version = 3;
  string domain = 4;
  int64 model_version = 5;
  string doc_string = 6;
  GraphProto graph = 7;
  repeated StringStringEntryProto metadata_props = 14;
  repeated TrainingInfoProto training_info = 20;
  repeated FunctionProto functions = 25;
}

比較重要的字段有:

lr_version : 當前的ONNX模型文件的版本，目前發(fā)布的最新版本為IR_VERSION_2019_3_18 = 6. 發(fā)布于2019年，版本7還在制定中。

opset_import: 當前的模型文件所依賴的算子domain和版本。

graph: 這個模型執(zhí)行的運算圖，這個是最重要的字段。

GraphProto的定義如下：

message GraphProto {
  repeated NodeProto node = 1;
  string name = 2;   // namespace Graph
  repeated TensorProto initializer = 5;
  repeated SparseTensorProto sparse_initializer = 15;
  string doc_string = 10;
  repeated ValueInfoProto input = 11;
  repeated ValueInfoProto output = 12;
  repeated ValueInfoProto value_info = 13;
  repeated TensorAnnotation quantization_annotation = 14;
}

比較重要的字段有:

initializer : 模型的每一網(wǎng)絡(luò)層的參數(shù), 模型訓練完成之后參數(shù)就被固定下來。

input : 模型的輸入格式。

output : 模型的輸出格式。

nodes : 定義了模型的所有運算模塊， 依照推理的次序排布。

NodeProto的定義如下：

message NodeProto {
  repeated string input = 1;    // namespace Value
  repeated string output = 2;   // namespace Value
  string name = 3;     // namespace Node
  string op_type = 4;  // namespace Operator
  string domain = 7;   // namespace Domain
  repeated AttributeProto attribute = 5;
  string doc_string = 6;
}

比較重要的字段有:

input : 輸入?yún)?shù)的名字。

output : 輸出參數(shù)的名字，這里需要留意的是，每一個網(wǎng)絡(luò)層之間的連接使用輸入和輸出的名字來確立的。

op_type : 算子的類型。

attributes : 算子的屬性， 其解析取決于算子的類型。

ONNX中最復雜的部分就是關(guān)于各種算子的描述，這也可以理解，構(gòu)成神經(jīng)網(wǎng)絡(luò)的主體就是這些算子。attributes 就是算子的一組帶名字的屬性。

本文中，我們介紹一個在mobilenetv2-7.onnx使用最多的算子: conv。

卷積神經(jīng)網(wǎng)絡(luò)在語音，圖像，視頻等處理上獲得了巨大成功. ONNX關(guān)于卷積網(wǎng)絡(luò)層的屬性定義主要有:

dilations: 擴展卷積，默認為1，即普通卷積。其數(shù)學定義如下：

group: 分組卷積, 其定義見文獻14. 默認為1， 即不分組。

kernel_shape: 定義了卷積核的大小。

pads: 定義了上下左右填充的像素數(shù)。

strides: 定義了卷積運算的步長。

03 ONNX的支持情況

各家的訓練和推理框架還在繼續(xù)發(fā)展，ONNX想成為行業(yè)標準顯然還為時尚早，但是目前尚沒有看到其他更好的通用模型描述格式，我們簡單歸納一下現(xiàn)在的ONNX的支持情況(不完整)：

參考文獻

[1] ONNX: https://github.com/onnx/onnx

[2] TENSORFLOW: https://www.tensorflow.org/

[3] CNTK: https://github.com/Microsoft/CNTK

[4] PYTORCH: https://pytorch.org/

[5] TNN: https://github.com/Tencent/TNN

[6] MNN: https://github.com/alibaba/MNN

[7] CUDNN: https://developer.nvidia.com/zh-cn/cudnn

[8] TENSORRT: https://developer.nvidia.com/zh-cn/tensorrt

[9] COREML: https://developer.apple.com/documentation/coreml

[10] NCNN: https://github.com/Tencent/ncnn

[11] NNAPI: https://developer.android.com/ndk/guides/neuralnetworks

[12] Protocol Buffers: https://developers.google.com/protocol-buffers

[13] Dilated Convolutions https://arxiv.org/abs/1511.07122

[14] Dynamic Group Convolutions https://arxiv.org/abs/2007.04242