PyTorch學習總結(三)——ONNX
1.什麼是ONNX
Open Neural Network Exchange (ONNX)是開放生態系統的第一步,它使人工智慧開發人員可以在專案的發展過程中選擇合適的工具;ONNX為AI models提供了一種開源格式。它定義了一個可以擴充套件的計算圖模型,同時也定義了內建操作符和標準資料型別。最初我們關注的是推理(評估)所需的能力。
Caffe2, PyTorch, Microsoft Cognitive Toolkit, Apache MXNet 和其他工具都在對ONNX進行支援。在不同的框架之間實現互操作性,並簡化從研究到產品化的過程,將提高人工智慧社群的創新速度。
2.torch.onnx
本文我們將主要介紹PyTorch中自帶的torch.onnx模組。該模組包含將模型匯出到ONNX IR格式的函式。這些模型可以被ONNX庫載入,然後將它們轉換成可在其他深度學習框架上執行的模型。
3 End-to-end AlexNet from PyTorch to Caffe2
這裡有一個簡單的指令碼,它將torchvision中預訓練的AlexNet匯出為ONNX。它執行一個簡單的推斷,然後將生成的跟蹤模型儲存到alexnet.proto
:
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()
torch.onnx.export(model, dummy_input, "alexnet.proto", verbose=True)
alexnet.proto
是一個二進位制的protobuf檔案,它包含您匯出的模型的網路結構和引數(在這裡,模型是AlexNet)。關鍵引數verbose=True
使exporter可以打印出一種人類可讀的網路表示:
# All parameters are encoded explicitly as inputs. By convention,
# learned parameters (ala nn.Module.state_dict) are first, and the
# actual inputs are last.
graph(%1 : Float(64, 3, 11, 11)
%2 : Float(64)
# The definition sites of all variables are annotated with type
# information, specifying the type and size of tensors.
# For example, %3 is a 192 x 64 x 5 x 5 tensor of floats.
%3 : Float(192, 64, 5, 5)
%4 : Float(192)
# ---- omitted for brevity ----
%15 : Float(1000, 4096)
%16 : Float(1000)
%17 : Float(10, 3, 224, 224)) { # the actual input!
# Every statement consists of some output tensors (and their types),
# the operator to be run (with its attributes, e.g., kernels, strides,
# etc.), its input tensors (%17, %1)
%19 : UNKNOWN_TYPE = Conv[kernels=[11, 11], strides=[4, 4], pads=[2, 2, 2, 2], dilations=[1, 1], group=1](%17, %1), uses = [[%20.i0]];
# UNKNOWN_TYPE: sometimes type information is not known. We hope to eliminate
# all such cases in a later release.
%20 : Float(10, 64, 55, 55) = Add[broadcast=1, axis=1](%19, %2), uses = [%21.i0];
%21 : Float(10, 64, 55, 55) = Relu(%20), uses = [%22.i0];
%22 : Float(10, 64, 27, 27) = MaxPool[kernels=[3, 3], pads=[0, 0, 0, 0], dilations=[1, 1], strides=[2, 2]](%21), uses = [%23.i0];
# ...
# Finally, a network returns some tensors
return (%58);
}
你還可以使用onnx庫來驗證protobuf。你可以用conda安裝onnx
:
conda install -c conda-forge onnx
然後,你可以執行:
import onnx
# Load the ONNX model
model = onnx.load("alexnet.proto")
# Check that the IR is well formed
onnx.checker.check_model(model)
# Print a human readable representation of the graph
onnx.helper.printable_graph(model.graph)
為了執行匯出的caffe2版本的指令碼,你需要以下兩項支援:
你需要安裝onnx-caffe2,一個純Python庫,它為ONNX提供了一個caffe2的編譯器。你可以用pip安裝onnx-caffe2:
pip install onnx-caffe2
安裝好以上依賴後,你可以使用針對Caffe2的編譯器了:
# ...continuing from above
import onnx_caffe2.backend as backend
import numpy as np
# or "CPU"
rep = backend.prepare(model, device="CUDA:0")
# For the Caffe2 backend:
# rep.predict_net is the Caffe2 protobuf for the network
# rep.workspace is the Caffe2 workspace for the network
# (see the class onnx_caffe2.backend.Workspace)
outputs = rep.run(np.random.randn(10, 3, 224, 224).astype(np.float32))
# To run networks with more than one input, pass a tuple
# rather than a single numpy ndarray.
print(outputs[0])
侷限性
ONNX exporter是一個基於跟蹤的exporter,這意味著它在執行您的模型時執行一次,並匯出在執行期間實際執行的操作。這意味著如果你的模型是動態的,例如,根據輸入資料改變操作行為,則export是不準確的。類似地,跟蹤可能只對特定的輸入大小有效(這就是為什麼我們需要對跟蹤進行顯式輸入的原因之一)。我們建議檢查模型跟蹤,並確保跟蹤的操作符看起來是合理的。
PyTorch和Caffe2通常有一些操作符的結果存在數值差異。根據模型結構,這些差異可以忽略不計,但也可能導致行為產生重大差異(特別對於那些未經訓練的模型)。在未來的版本中,我們打算讓Caffe2可以直接呼叫Torch中的一些操作,使得在一些關注精度的任務中,可以幫助研究人員緩和這些差異,同時也會記錄下這些差異。
支援的操作符
ONNX支援下面的操作符:
- add (nonzero alpha not supported)
- sub (nonzero alpha not supported)
- mul
- div
- cat
- mm
- addmm
- neg
- tanh
- sigmoid
- mean
- t
- expand (only when used before a broadcasting ONNX operator; e.g., add)
- transpose
- view
- split
- squeeze
- prelu (single weight shared among input channels not supported)
- threshold (non-zero threshold/non-zero value not supported)
- leaky_relu
- glu
- softmax
- avg_pool2d (ceil_mode not supported)
- log_softmax
- unfold (experimental support with ATen-Caffe2 integration)
- elu
- Conv
- BatchNorm
- MaxPool1d (ceil_mode not supported)
- MaxPool2d (ceil_mode not supported)
- MaxPool3d (ceil_mode not supported)
- Embedding (no optional arguments supported)
- RNN
- ConstantPadNd
- Dropout
- FeatureDropout (training mode not supported)
- Index (constant integer and tuple indices supported)
- Negate
上面操作符集合足以匯出以下模型:
- AlexNet
- DCGAN
- DenseNet
- Inception (warning: this model is highly sensitive to changes in operator implementation)
- ResNet
- SuperResolution
- VGG
- word_language_model
指定操作符定義的介面是高度實驗性和無文件的;喜歡嚐鮮的使用者請注意,APIs可能會在未來的介面中發生變化
函式
torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None)
將一個模型匯出到ONNX格式。該exporter會執行一次你的模型,以便於記錄模型的執行軌跡,並將其匯出;目前,exporter還不支援動態模型(例如,RNNs)。
另請參閱:onnx-export
引數:
- model(torch.nn.Module)-要被匯出的模型
- args(引數的集合)-模型的輸入,例如,這種model(*args)方式是對模型的有效呼叫。任何非Variable引數都將硬編碼到匯出的模型中;任何Variable引數都將成為匯出的模型的輸入,並按照他們在args中出現的順序輸入。如果args是一個Variable,這等價於用包含這個Variable的1-ary元組呼叫它。(注意:現在不支援向模型傳遞關鍵字引數。)
- f-一個類檔案的物件(必須實現檔案描述符的返回)或一個包含檔名字串。一個二進位制Protobuf將會寫入這個檔案中。
- export_params(bool,default True)-如果指定,所有引數都會被匯出。如果你只想匯出一個未訓練的模型,就將此引數設定為False。在這種情況下,匯出的模型將首先把所有parameters作為參arguments,順序由
model.state_dict().values()
指定。 - verbose(bool,default False)-如果指定,將會輸出被匯出的軌跡的除錯描述。
- training(bool,default False)-匯出訓練模型下的模型。目前,ONNX只面向推斷模型的匯出,所以一般不需要將該項設定為True。
- input_names(list of strings, default empty list)-按順序分配名稱到圖中的輸入節點。
- output_names(list of strings, default empty list)-按順序分配名稱到圖中的輸出節點。