近日,微軟亞洲研究院和華為舉辦了 ONNX 合作伙伴研討會,這是 ONNX 開源社區(qū)成立以來首次在中國舉辦的活動。在研討會中,微軟、Facebook、華為和英特爾等的開發(fā)者介紹了他們在 ONNX 上的開源貢獻及思考。
在過去的一年多中,ONNX 這種「通用」的神經網絡交換格式已經有了很長遠的發(fā)展,用不同框架編寫的模型可以在不同的平臺中流通。在這次研討會中,我們確切地感受到了這一點,因為開源社區(qū)圍繞著 ONNX 介紹了很多優(yōu)化工具和資源庫。
微軟上個月開源了 ONNX Runtime,其專為 ONNX 格式的模型設計了高性能推理引擎。Facebook 早兩個月前開源了 ONNXIFI,其為 ONNX 提供了用于框架集成的接口,即一組用于加載和執(zhí)行 ONNX 計算圖的跨平臺 API。更早一些,英特爾在今年 3 月份就開源 nGraph,它能編譯 ONNX 格式的模型,并在 CPU 或 GPU 等硬件加速模型的運行。
而到了昨天,微軟又開源了 ONNX.JS,它是一種在瀏覽器和 Node.js 上運行 ONNX 模型的 JavaScript 庫。它部署的模型效率非常高,且能實現交互式的直觀推理。該開源項目給出了圖像分類的交互式演示,且在 Chrome 瀏覽器和 CPU 下比 TensorFlow.JS 快了近 8 倍,后文將詳細介紹這一開源庫。
當然除了這些開源工作,ONNX 社區(qū)還有更多的實踐,例如如何部署 ONNX 模型到邊緣設備、如何維護一個包羅萬象的 ONNX Model Zoo 等。本文主要從什么是 ONNX、怎樣用 ONNX,以及如何優(yōu)化 ONNX 三方面看看 ONNX 是不是已經引領「框架間的江湖」了。
什么是 ONNX
很多開發(fā)者在玩 GitHub 的時候都有這樣「悲痛」的經歷,好不容易找到令人眼前一亮的項目,然而發(fā)現它使用我們不熟悉的框架寫成。其實我們會發(fā)現很多優(yōu)秀的視覺模型是用 Caffe 寫的,很多新的研究論文是用 PyTorch 寫的,而更多的模型用 TensorFlow 寫成。因此如果我們要測試它們就必須擁有對應的框架環(huán)境,但 ONNX 交換格式令我們在同一環(huán)境下測試不同模型有了依靠。
簡而言之 ONNX 就是一種框架間的轉換格式,例如我們用 TensorFlow 寫的模型可以轉換為 ONNX 格式,并在 Caffe2 環(huán)境下運行該模型。
機器之心SyncedONNX小程序
項目地址:https://github.com/onnx/onnx
ONNX 定義了一種可擴展的計算圖模型、一系列內置的運算單元(OP)和標準數據類型。每一個計算流圖都定義為由節(jié)點組成的列表,并構建有向無環(huán)圖。其中每一個節(jié)點都有一個或多個輸入與輸出,每一個節(jié)點稱之為一個 OP。這相當于一種通用的計算圖,不同深度學習框架構建的計算圖都能轉化為它。
如下所示,目前 ONNX 已經支持大多數框架,使用這些框架構建的模型可以轉換為通用的 ONNX 計算圖和 OP?,F階段 ONNX 只支持推理,所以導入的模型都需要在原框架完成訓練。
其中 Frameworks 下的框架表示它們已經內嵌了 ONNX,開發(fā)者可以直接通過這些框架的內置 API 將模型導出為 ONNX 格式,或采用它們作為推理后端。而 Converters 下的框架并不直接支持 ONNX 格式,但是可以通過轉換工具導入或導出這些框架的模型。
其實并不是所有框架都支持導入和導出 ONNX 格式的模型,有一些并不支持導入 ONNX 格式的模型,例如 PyTorch 和 Chainer 等,TensorFlow 的 ONNX 導入同樣也正處于實驗階段。下圖展示了各框架對 ONNX 格式的支持情況:
怎樣使用 ONNX
對于內建了 ONNX 的框架而言,使用非常簡單,只需要調用 API 導出或導入已訓練模型就可以了。例如對 PyTorch 而言,只需要幾個簡單的步驟就能完成模型的導出和導入。簡單而言,首先加載 torch.onnx 模塊,然后導出預訓練模型并查看模型結構信息,最后再將導出的 ONNX 模型加載到另外的框架就能執(zhí)行推理了。
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: AlexNet/Sequential[features]/Conv2d[0]
%18 : Float(10, 64, 55, 55) = onnx::Relu(%17), scope: AlexNet/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: AlexNet/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: AlexNet/Sequential[classifier]/Linear[6]
return (%output1);
}
其實我們也可以借助 ONNX 檢查中間表征,不過這里并不介紹。后面加載另外一個框架并執(zhí)行推理同樣非常簡單。如下所示,我們可以從 caffe2 中加載 ONNX 的后端,并將前面保存的模型加載到該后端,從而在新框架下進行推理。這里我們能選擇執(zhí)行推理的硬件,并直接推理得出輸出結果。
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 可能不一樣,但主要過程也就這簡單的幾步。
怎樣優(yōu)化 ONNX
前面就已經介紹了 Model Zoo、ONNX Runtime 和 ONNX.JS,現在,我們可以具體看看它們都是什么,它們怎樣才能幫助我們優(yōu)化 ONNX 模型的選擇與推理速度。
Model Zoo
ONNX Model Zoo 包含了一系列預訓練模型,它們都是 ONNX 格式,且能獲得當前最優(yōu)的性能。因此只要下載這樣的模型,我們本地不論是 TensorFlow 還是 MXNet,只要是只是能加載模型的框架,就能運行這些預訓練模型。
項目地址:https://github.com/onnx/models
更重要的是,這個 Model Zoo 不僅有調用預訓練模型的代碼,它還為每個預訓練模型開放了對應的訓練代碼。訓練和推理代碼都是用 Jupyter Notebook 寫的,數據和模型等都有對應的鏈接。
目前該 Model Zoo 主要從圖像分類、檢測與分割、圖像超分辨、機器翻譯和語音識別等 14 個方向包含 19 種模型,還有更多的模型還在開發(fā)中。如下展示了圖像分類中已經完成的模型,它們都是通用的 ONNX 格式。
此外在這次的研討會中,Model Zoo 的維護者還和大家討論了目前面臨的問題及解決方法,例如目前的預訓練模型主要集中在計算機視覺方面、ONNX 缺少一些特定的 OP、權重計算圖下載慢等。因此 Model Zoo 接下來也會更關注其它語音和語言等模型,優(yōu)化整個 GitHub 項目的下載結構。
ONNX Runtime
微軟開源的 ONNX Runtime 推理引擎支持 ONNX 中定義的所有運算單元,它非常關注靈活性和推理性能。因此不論我們的開發(fā)環(huán)境是什么,Runtime 都會基于各種平臺與硬件選擇不同的自定義加速器,并希望以最小的計算延遲和資源占用完成推理。
文檔地址:https://docs.microsoft.com/en-us/python/api/overview/azure/onnx/intro
ONNX Runtime 可以自動調用各種硬件加速器,例如英偉達的 CUDA、TensorRT 和英特爾的 MKL-DNN、nGraph。如下所示,ONNX 格式的模型可以傳入到藍色部分的 Runtime,并自動完成計算圖分割及并行化處理,最后我們只需要如橙色所示的輸入數據和輸出結果就行了。
其實在實際使用的時候,開發(fā)者根本不需要考慮藍色的部分,不論是編譯還是推理,代碼都賊簡單。如下所示,導入 onnxruntime 模塊后,調用 InferenceSession() 方法就能導入 ONNX 格式的模型,并完成上圖一系列復雜的優(yōu)化。最后只需要 session.run() 就可以進行推理了,所有的優(yōu)化過程都隱藏了細節(jié)。
import onnxruntime
session = onnxruntime.InferenceSession("your_model.onnx")
prediction = session.run(None, {"input1": value})
在研討會中,開發(fā)者表示 Runtime 的目標是構建高性能推理引擎,它需要利用最好的加速器和完整的平臺支持。只需要幾行代碼就能把計算圖優(yōu)化一遍,這對 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,開發(fā)者可以直接將預訓練的 ONNX 模型部署到瀏覽器,這些預訓練模型可以是 Model Zoo 中的,也可以是自行轉換的。部署到瀏覽器有很大的優(yōu)勢,它能減少服務器與客戶端之間的信息交流,并獲得免安裝和跨平臺的機器學習模型體驗。如下所示為部署到網頁端的 SqueezeNet:
如上若是選擇 GPU,它會采用 WebGL 訪問 GPU。如果選擇 CPU,那么其不僅會采用 WebAssembly 以接近原生的速度執(zhí)行模型,同時也會采用 Web Workers 提供的「多線程」環(huán)境來并行化數據處理。該項目表明,通過充分利用 WebAssembly 和 Web Workers,CPU 可以獲得很大的性能提升。這一點在項目提供的 Benchmarks 中也有相應的展示:
以上微軟在 Chrome 和 Edge 瀏覽器中測試了 ResNet-50 的推理速度,其中比較顯著的是 CPU 的推理速度。這主要是因為 Keras.js 和 TensorFlow.js 在任何瀏覽器中都不支持 WebAssembly。
最后,從 ONNXIFI 到 ONNX.js,開源社區(qū)已經為 ONNX 格式構建出眾多的優(yōu)化庫、轉換器和資源。很多需要支持多框架的場景也都將其作為默認的神經網絡格式,也許以后,ONNX 真的能統(tǒng)一神經網絡之間的江湖。