沒有“中間商賺差價”，OpenVINO直接支持 PyTorch 模型對象

背景

作為最熱門的開源深度學習框架之一，PyTorch 的易用性和靈活性使其深受學術和研究界的喜愛。之前 OpenVINO 對于 PyTorch 模型的支持也僅僅停留在 ONNX 過渡階段，需要通過將 PyTorch 動態模型導出為 ONNX 靜態格式后，才可以直接被 OpenVINO runtime 離線加載，雖然 PyTorch 也提供了官方的 torch.onnx.export 接口幫助開發者導出 ONNX 模型，但畢竟有這么一個“中間商”在那里，其中很多額外的配置工作也為 OpenVINO 開發者帶來了不便，諸如動態/靜態輸入設定，以及 opset 版本設定等。

一、OpenVINO 直接支持 PyTorch 模型對象

隨著 OpenVINO 2023.0 版本的發布，OpenVINO 工具庫中預置了全新的 PyTorch 前端，為開發者們提供了一條全新的 PyTorch 模型支持路徑，帶來更友好的用戶體驗——OpenVINO 的 mo 工具可以直接將 PyTorch 模型對象轉化為 OpenVINO 的模型對象，開發者可以不需要將 ONNX 模型作為中間過渡。

import torchvision import torch from openvino.tools.mo import convert_model model = torchvision.models.resnet50(pretrained=True) ov_model=convert_model(model)

對比以 ONNX 作為中間過度的方式，新 PyTorch 前端有以下特點：

目前支持的 PyTorch 模型對象有：

· torch.nn.Module

· torch.jit.ScriptModule

· torch.jit.ScriptFunction

在 OpenVINO 內部，PyTorch 前端基于 TorchScript 進行模型導出，而 TorchScript 支持兩種模型導出模式，一種稱為 Tracing，一種稱為 Scripting。其中 Tracing 指的是 PyTorch 在模型運行時，追蹤運行經過的模塊算子，實時構建計算流圖，并最終總結為一種中間表示，Trace 是個雙刃劍，好處是用戶無需了解 Python 代碼個中細節，無論是 Function、Module 還是 Generators、Coroutines，Tracing 都會忠實地記錄下經過的 Tensor 以及 Tensor Function，非常適用于不涉及數據相關控制流的簡單模塊和功能，例如標準卷積神經網絡，壞處就在于 Tracing 不能感知控制流和計算圖的動態，如 if 語句或循環。比如他會把循環展開，一方面可能可以增加編譯優化的空間，另一方面如果該循環在不同 infer 的時候是動態變長的，那么 Tracing 不能感知到這一點，只會將 Tracing 時候的循環記錄下來。為了轉換包含依賴于數據的控制流的模塊和函數，提供了一種 Scripting 機制，Scripting 從 Python 源代碼級別進行解析，而非在運行時構建。Scripting 會去理解所有的 code，真正像一個編譯器一樣去進行語法分析等操作。Scripting 相當于一個嵌入在 Python/Pytorch 的DSL，其語法只是 PyTorch 語法的子集，這意味著存在一些 op 和語法 Scripting 不支持，這樣在編譯的時候就會遇到問題。

在剛剛的例子中 PyTorch 前端使用 Scripting 進行模型導出，如果想使用 Tracing 的方式，可以在接口中新增一個 example_input 參數，此時 PyTorch 前端會優先調用 Tracing 的方式，當 Tracing 的方式失敗后，再調用 Scripting 方式。

import torchvision import torch from openvino.tools.mo import convert_model model = torchvision.models.resnet50(pretrained=True) ov_model=convert_model(model,example_input=torch.zeros(1,3,100,100))

目前 examle_input 支持的數據格式有：

· openvino.runtime.Tensor

· torch.Tensor

· np.ndarray

· listortuplewith tensors (openvino.runtime.Tensor/torch.Tensor/np.ndarray)

· dictionarywhere key is the input name, value is the tensor (openvino.runtime.Tensor/torch.Tensor/np.ndarray)

值得注意的是，以上兩個例子導出的均為動態輸入模型對象，如果想指定模型的輸入 shape，可以再次添加額外的參數 input_shape/input, 將輸入 shape 作為參數傳入，選其一即可。案例可參考以下的實戰部分。

最后，如果開發者希望導出靜態 IR 文件以便后續使用，也可以調用以下接口，將 OpenVINO 的模型對象進行序列化：

serialize(ov_model, str(ir_model_xml))

二、BERT 模型案例實戰

接下來我們通過一個實例來看下如何完成從 BERT 模型轉化到量化的全過程。

1. 獲取 PyTorch 模型對象

torch_model = BertForSequenceClassification.from_pretrained(PRETRAINED_MODEL_DIR)

2. 設置模型參數并轉化為 OpenVINO 模型對象

由于 BERT 是一個多輸入模型，這里額外添加了一個 input=input_info 參數，可以用來指定多輸入模型中每一個 input 的 shape 以及數據類型。

input_shape = PartialShape([1, -1]) input_info = [("input_ids", input_shape, np.int64),("attention_mask", input_shape, np.int64),("token_type_ids", input_shape, np.int64)] default_input = torch.ones(1, MAX_SEQ_LENGTH, dtype=torch.int64) inputs = { "input_ids": default_input, "attention_mask": default_input, "token_type_ids": default_input, } model=convert_model(torch_model,example_input=inputs,input=input_info)

3. 準備校驗數據集，并啟動量化

上一步中獲得的 model 為 openvino.runtime.Model 類型，可以直接被 NNCF 工具加載

calibration_dataset = nncf.Dataset(data_source, transform_fn) # Quantize the model. By specifying model_type, we specify additional transformer patterns in the model. quantized_model = nncf.quantize(model, calibration_dataset, model_type=ModelType.TRANSFORMER)

4. 編譯量化后的模型對象，并進行推理

compiled_quantized_model = core.compile_model(model=quantized_model, device_name="CPU") output_layer = compiled_quantized_model.outputs[0] result = compiled_quantized_model(inputs)[output_layer] result = np.argmax(result) print(f"Text 1: {sample['sentence1']}") print(f"Text 2: {sample['sentence2']}") print(f"Thesamemeaning:{'yes'ifresult==1else'no'}")

最終結果如下：

Text 1: Wal-Mart said it would check all of its million-plus domestic workers to ensure they were legally employed . Text 2: It has also said it would review all of its domestic employees more than 1 million to ensure they have legal status . Thesamemeaning:yes

完整實例和性能精度比較，可以參考：

https://github.com/openvinotoolkit/openvino_notebooks/blob/main/notebooks/105-language-quantize-bert/105-language-quantize-bert.ipynb

三、總結

作為近期發布的最新版本，OpenVINO 2023.0 中的 mo 工具可以在不需要通過 ONNX 中間過渡的情況下，直接將 PyTorch 模型對象轉化為 OpenVINO 對象，免去開發者離線轉化和額外配置的過程，帶來更友好的用戶體驗。鑒于該功能是預發布狀態，可能存在部分算子不支持的情況，此時，開發者依舊可以使用之前的路徑，依托 ONNX 前端進行 PyTorch 模型的轉換。