開源一年多的模型交換格式ONNX，已經一統框架江湖了？

摘要： 在過去的一年多中，ONNX 這種「通用」的神經網路交換格式已經有了很長遠的發展，用不同框架編寫的模型可以在不同的平臺中流通。在這次研討會中，我們確切地感受到了這一點，因為開源社群圍繞著 ONNX 介紹了很多優化工具和資源庫。 微軟上個月開源了 ONNX Runtime，其專為 ONN...

在過去的一年多中，ONNX 這種「通用」的神經網路交換格式已經有了很長遠的發展，用不同框架編寫的模型可以在不同的平臺中流通。在這次研討會中，我們確切地感受到了這一點，因為開源社群圍繞著 ONNX 介紹了很多優化工具和資源庫。

微軟上個月開源了 ONNX Runtime，其專為 ONNX 格式的模型設計了高效能推理引擎。Facebook 早兩個月前開源了 ONNXIFI，其為 ONNX 提供了用於框架整合的介面，即一組用於載入和執行 ONNX 計算圖的跨平臺 API。更早一些，英特爾在今年 3 月份就開源 nGraph，它能編譯 ONNX 格式的模型，並在 CPU 或 GPU 等硬體加速模型的執行。

而到了昨天，微軟又開源了 ONNX.JS，它是一種在瀏覽器和 Node.js 上執行 ONNX 模型的 JavaScript 庫。它部署的模型效率非常高，且能實現互動式的直觀推理。該開源專案給出了影象分類的互動式演示，且在 Chrome 瀏覽器和 CPU 下比TensorFlow.JS 快了近 8 倍，後文將詳細介紹這一開源庫。

當然除了這些開源工作，ONNX 社群還有更多的實踐，例如如何部署 ONNX 模型到邊緣裝置、如何維護一個包羅永珍的 ONNX Model Zoo 等。本文主要從什麼是 ONNX、怎樣用 ONNX，以及如何優化 ONNX 三方面看看 ONNX 是不是已經引領「框架間的江湖」了。

什麼是 ONNX

很多開發者在玩 GitHub 的時候都有這樣「悲痛」的經歷，好不容易找到令人眼前一亮的專案，然而發現它使用我們不熟悉的框架寫成。其實我們會發現很多優秀的視覺模型是用 Caffe 寫的，很多新的研究論文是用 PyTorch 寫的，而更多的模型用TensorFlow 寫成。因此如果我們要測試它們就必須擁有對應的框架環境，但 ONNX 交換格式令我們在同一環境下測試不同模型有了依靠。

簡而言之 ONNX 就是一種框架間的轉換格式，例如我們用TensorFlow 寫的模型可以轉換為 ONNX 格式，並在 Caffe2 環境下執行該模型。

• 專案地址：https://github.com/onnx/onnx

ONNX 定義了一種可擴充套件的計算圖模型、一系列內建的運算單元（OP）和標準資料型別。每一個計算流圖都定義為由節點組成的列表，並構建有向無環圖。其中每一個節點都有一個或多個輸入與輸出，每一個節點稱之為一個 OP。這相當於一種通用的計算圖，不同深度學習框架構建的計算圖都能轉化為它。

如下所示，目前 ONNX 已經支援大多數框架，使用這些框架構建的模型可以轉換為通用的 ONNX 計算圖和 OP。現階段 ONNX 只支援推理，所以匯入的模型都需要在原框架完成訓練。

其中 Frameworks 下的框架表示它們已經內嵌了 ONNX，開發者可以直接通過這些框架的內建 API 將模型匯出為 ONNX 格式，或採用它們作為推理後端。而 Converters 下的框架並不直接支援 ONNX 格式，但是可以通過轉換工具匯入或匯出這些框架的模型。

其實並不是所有框架都支援匯入和匯出 ONNX 格式的模型，有一些並不支援匯入 ONNX 格式的模型，例如 PyTorch 和 Chainer 等，TensorFlow 的 ONNX 匯入同樣也正處於實驗階段。下圖展示了各框架對 ONNX 格式的支援情況：

怎樣使用 ONNX

對於內建了 ONNX 的框架而言，使用非常簡單，只需要呼叫 API 匯出或匯入已訓練模型就可以了。例如對 PyTorch 而言，只需要幾個簡單的步驟就能完成模型的匯出和匯入。簡單而言，首先載入 torch.onnx 模組，然後匯出預訓練模型並檢視模型結構資訊，最後再將匯出的 ONNX 模型載入到另外的框架就能執行推理了。

from torch.autograd import Variable
import torch.onnx
import torchvision

dummy_input = Variable(torch.randn(10, 3, 224, 224)).cuda()
model = torchvision.models.alexnet(pretrained=True).cuda()

input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
output_names = [ "output1" ]

torch.onnx.export(model, dummy_input, "alexnet.onnx", verbose=True, input_names=input_names, output_names=output_names)

如上所示將匯出 ONNX 格式的AlexNet模型，其中"alexnet.onnx"為儲存的模型，input_names、output_names 和 verbose=True 都是為了打印出模型結構資訊。同樣隨機產生的「影象」dummy_input 也是為了瞭解模型結構，因為我們可以通過它理解輸入與每一層具體的引數維度。以下展示了 ONNX 輸出的簡要模型資訊：

graph(%actual_input_1 : Float(10, 3, 224, 224)
%learned_0 : Float(64, 3, 11, 11)
%learned_1 : Float(64)
# ---- omitted for brevity ----
%learned_14 : Float(1000, 4096)
%learned_15 : Float(1000)) {
%17 : Float(10, 64, 55, 55) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[11, 11], pads=[2, 2, 2, 2], strides=[4, 4]](%actual_input_1, %learned_0, %learned_1), scope: <mark data-type="technologies" data-id="a0955638-affc-472d-80b0-f645a4fd4604">AlexNet</mark>/Sequential[features]/Conv2d[0]
%18 : Float(10, 64, 55, 55) = onnx::Relu(%17), scope: <mark data-type="technologies" data-id="a0955638-affc-472d-80b0-f645a4fd4604">AlexNet</mark>/Sequential[features]/ReLU[1]
%19 : Float(10, 64, 27, 27) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%18), scope: <mark data-type="technologies" data-id="a0955638-affc-472d-80b0-f645a4fd4604">AlexNet</mark>/Sequential[features]/MaxPool2d[2]
# ---- omitted for brevity ----
%output1 : Float(10, 1000) = onnx::Gemm[alpha=1, beta=1, broadcast=1, transB=1](%45, %learned_14, %learned_15), scope: <mark data-type="technologies" data-id="a0955638-affc-472d-80b0-f645a4fd4604">AlexNet</mark>/Sequential[classifier]/Linear[6]
return (%output1);
}

其實我們也可以藉助 ONNX 檢查中間表徵，不過這裡並不介紹。後面載入另外一個框架並執行推理同樣非常簡單。如下所示，我們可以從 caffe2 中載入 ONNX 的後端，並將前面儲存的模型載入到該後端，從而在新框架下進行推理。這裡我們能選擇執行推理的硬體，並直接推理得出輸出結果。

import caffe2.python.onnx.backend as backend
import numpy as np
import onnx

model = onnx.load("alexnet.onnx")
rep = backend.prepare(model, device="CUDA:0") # or "CPU"
outputs = rep.run(np.random.randn(10, 3, 224, 224).astype(np.float32))

其實也就兩三行程式碼涉及 ONNX 的核心操作，即匯出模型、載入模型和載入另一個框架的後端。TensorFlow 或 CNTK 等其它框架的具體 API 可能不一樣，但主要過程也就這簡單的幾步。

怎樣優化 ONNX

前面就已經介紹了 Model Zoo、ONNX Runtime 和 ONNX.JS，現在，我們可以具體看看它們都是什麼，它們怎樣才能幫助我們優化 ONNX 模型的選擇與推理速度。

Model Zoo

ONNX Model Zoo 包含了一系列預訓練模型，它們都是 ONNX 格式，且能獲得當前最優的效能。因此只要下載這樣的模型，我們本地不論是TensorFlow 還是MXNet，只要是只是能載入模型的框架，就能執行這些預訓練模型。

• 專案地址：https://github.com/onnx/models

更重要的是，這個 Model Zoo 不僅有呼叫預訓練模型的程式碼，它還為每個預訓練模型開放了對應的訓練程式碼。訓練和推理程式碼都是用JupyterNotebook 寫的，資料和模型等都有對應的連結。

目前該 Model Zoo 主要從影象分類、檢測與分割、影象超分辨、機器翻譯和語音識別等 14 個方向包含 19 種模型，還有更多的模型還在開發中。如下展示了影象分類中已經完成的模型，它們都是通用的 ONNX 格式。

此外在這次的研討會中，Model Zoo 的維護者還和大家討論了目前面臨的問題及解決方法，例如目前的預訓練模型主要集中在計算機視覺方面、ONNX 缺少一些特定的 OP、權重計算圖下載慢等。因此 Model Zoo 接下來也會更關注其它語音和語言等模型，優化整個 GitHub 專案的下載結構。

ONNX Runtime

微軟開源的 ONNX Runtime推理引擎支援 ONNX 中定義的所有運算單元，它非常關注靈活性和推理效能。因此不論我們的開發環境是什麼，Runtime 都會基於各種平臺與硬體選擇不同的自定義加速器，並希望以最小的計算延遲和資源佔用完成推理。

• 文件地址：https://docs.microsoft.com/en-us/python/api/overview/azure/onnx/intro

ONNX Runtime 可以自動呼叫各種硬體加速器，例如英偉達的 CUDA、TensorRT 和英特爾的 MKL-DNN、nGraph。如下所示，ONNX 格式的模型可以傳入到藍色部分的 Runtime，並自動完成計算圖分割及並行化處理，最後我們只需要如橙色所示的輸入資料和輸出結果就行了。

其實在實際使用的時候，開發者根本不需要考慮藍色的部分，不論是編譯還是推理，程式碼都賊簡單。如下所示，匯入 onnxruntime 模組後，呼叫 InferenceSession() 方法就能匯入 ONNX 格式的模型，並完成上圖一系列複雜的優化。最後只需要 session.run() 就可以進行推理了，所有的優化過程都隱藏了細節。

import onnxruntime

session = onnxruntime.InferenceSession("your_model.onnx") 
prediction = session.run(None, {"input1": value})

在研討會中，開發者表示 Runtime 的目標是構建高效能推理引擎，它需要利用最好的加速器和完整的平臺支援。只需要幾行程式碼就能把計算圖優化一遍，這對 ONNX 格式的模型是個大福利。

ONNX.JS

ONNX.js 是一個在瀏覽器上執行 ONNX 模型的庫，它採用了 WebAssembly 和 WebGL 技術，並在 CPU 或 GPU 上推理 ONNX 格式的預訓練模型。

• 專案地址：https://github.com/Microsoft/onnxjs

• Demo 展示地址：https://microsoft.github.io/onnxjs-demo

通過 ONNX.js，開發者可以直接將預訓練的 ONNX 模型部署到瀏覽器，這些預訓練模型可以是 Model Zoo 中的，也可以是自行轉換的。部署到瀏覽器有很大的優勢，它能減少伺服器與客戶端之間的資訊交流，並獲得免安裝和跨平臺的機器學習模型體驗。如下所示為部署到網頁端的 SqueezeNet：

如上若是選擇 GPU，它會採用 WebGL 訪問 GPU。如果選擇 CPU，那麼其不僅會採用 WebAssembly 以接近原生的速度執行模型，同時也會採用 Web Workers 提供的「多執行緒」環境來並行化資料處理。該專案表明，通過充分利用 WebAssembly 和 Web Workers，CPU 可以獲得很大的效能提升。這一點在專案提供的 Benchmarks 中也有相應的展示：

以上微軟在 Chrome 和 Edge 瀏覽器中測試了 ResNet-50 的推理速度，其中比較顯著的是 CPU 的推理速度。這主要是因為 Keras.js 和TensorFlow.js 在任何瀏覽器中都不支援 WebAssembly。

最後，從 ONNXIFI 到 ONNX.js，開源社群已經為 ONNX 格式構建出眾多的優化庫、轉換器和資源。很多需要支援多框架的場景也都將其作為預設的神經網路格式，也許以後，ONNX 真的能統一神經網路之間的江湖。