沒(méi)有“中間商賺差價(jià)”，OpenVINO直接支持 PyTorch 模型對(duì)象

背景

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

一、OpenVINO 直接支持 PyTorch 模型對(duì)象

隨著 OpenVINO 2023.0 版本的發(fā)布，OpenVINO 工具庫(kù)中預(yù)置了全新的 PyTorch 前端，為開(kāi)發(fā)者們提供了一條全新的 PyTorch 模型支持路徑，帶來(lái)更友好的用戶體驗(yàn)——OpenVINO 的 mo 工具可以直接將 PyTorch 模型對(duì)象轉(zhuǎn)化為 OpenVINO 的模型對(duì)象，開(kāi)發(fā)者可以不需要將 ONNX 模型作為中間過(guò)渡。

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

對(duì)比以 ONNX 作為中間過(guò)度的方式，新 PyTorch 前端有以下特點(diǎn)：

目前支持的 PyTorch 模型對(duì)象有：

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

在剛剛的例子中 PyTorch 前端使用 Scripting 進(jìn)行模型導(dǎo)出，如果想使用 Tracing 的方式，可以在接口中新增一個(gè) example_input 參數(shù)，此時(shí) PyTorch 前端會(huì)優(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)

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

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

serialize(ov_model, str(ir_model_xml))

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

接下來(lái)我們通過(guò)一個(gè)實(shí)例來(lái)看下如何完成從 BERT 模型轉(zhuǎn)化到量化的全過(guò)程。

1. 獲取 PyTorch 模型對(duì)象

torch_model = BertForSequenceClassification.from_pretrained(PRETRAINED_MODEL_DIR)

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

由于 BERT 是一個(gè)多輸入模型，這里額外添加了一個(gè) input=input_info 參數(shù)，可以用來(lái)指定多輸入模型中每一個(gè) 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)備校驗(yàn)數(shù)據(jù)集，并啟動(dòng)量化

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

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

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

三、總結(jié)

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