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

背景

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

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

隨著 OpenVINO 2023.0 版本的發(fā)布，OpenVINO 工具庫中預(yù)置了全新的 PyTorch 前端，為開發(fā)者們提供了一條全新的 PyTorch 模型支持路徑，帶來更友好的用戶體驗——OpenVINO 的 mo 工具可以直接將 PyTorch 模型對象轉(zhuǎn)化為 OpenVINO 的模型對象，開發(fā)者可以不需要將 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 內(nèi)部，PyTorch 前端基于 TorchScript 進(jìn)行模型導(dǎo)出，而 TorchScript 支持兩種模型導(dǎo)出模式，一種稱為 Tracing，一種稱為 Scripting。其中 Tracing 指的是 PyTorch 在模型運行時，追蹤運行經(jīng)過的模塊算子，實時構(gòu)建計算流圖，并最終總結(jié)為一種中間表示，Trace 是個雙刃劍，好處是用戶無需了解 Python 代碼個中細(xì)節(jié)，無論是 Function、Module 還是 Generators、Coroutines，Tracing 都會忠實地記錄下經(jīng)過的 Tensor 以及 Tensor Function，非常適用于不涉及數(shù)據(jù)相關(guān)控制流的簡單模塊和功能，例如標(biāo)準(zhǔn)卷積神經(jīng)網(wǎng)絡(luò)，壞處就在于 Tracing 不能感知控制流和計算圖的動態(tài)，如 if 語句或循環(huán)。比如他會把循環(huán)展開，一方面可能可以增加編譯優(yōu)化的空間，另一方面如果該循環(huán)在不同 infer 的時候是動態(tài)變長的，那么 Tracing 不能感知到這一點，只會將 Tracing 時候的循環(huán)記錄下來。為了轉(zhuǎn)換包含依賴于數(shù)據(jù)的控制流的模塊和函數(shù)，提供了一種 Scripting 機制，Scripting 從 Python 源代碼級別進(jìn)行解析，而非在運行時構(gòu)建。Scripting 會去理解所有的 code，真正像一個編譯器一樣去進(jìn)行語法分析等操作。Scripting 相當(dāng)于一個嵌入在 Python/Pytorch 的DSL，其語法只是 PyTorch 語法的子集，這意味著存在一些 op 和語法 Scripting 不支持，這樣在編譯的時候就會遇到問題。

在剛剛的例子中 PyTorch 前端使用 Scripting 進(jìn)行模型導(dǎo)出，如果想使用 Tracing 的方式，可以在接口中新增一個 example_input 參數(shù)，此時 PyTorch 前端會優(yōu)先調(diào)用 Tracing 的方式，當(dāng) Tracing 的方式失敗后，再調(diào)用 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 支持的數(shù)據(jù)格式有：

· 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)

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

最后，如果開發(fā)者希望導(dǎo)出靜態(tài) IR 文件以便后續(xù)使用，也可以調(diào)用以下接口，將 OpenVINO 的模型對象進(jìn)行序列化：

serialize(ov_model, str(ir_model_xml))

二、BERT 模型案例實戰(zhàn)

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

1. 獲取 PyTorch 模型對象

torch_model = BertForSequenceClassification.from_pretrained(PRETRAINED_MODEL_DIR)

2. 設(shè)置模型參數(shù)并轉(zhuǎn)化為 OpenVINO 模型對象

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

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. 準(zhǔn)備校驗數(shù)據(jù)集，并啟動量化

上一步中獲得的 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. 編譯量化后的模型對象，并進(jìn)行推理

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'}")

最終結(jié)果如下：

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

三、總結(jié)

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