没有“中间商赚差价”， OpenVINO™ 直接支持 PyTorch 模型对象

作者 | 杨亦诚

排版 | 李擎

一、OpenVINO™ 直接支持 PyTorch 模型对象

随着 OpenVINO™ 2023.0 版本的发布，OpenVINO™ 工具库中预置了全新的 PyTorch 前端，为开发者们提供了一条全新的 PyTorch 模型支持路径，带来更友好的用户体验—— OpenVINO™ 的 mo 工具可以直接将 PyTorch 模型对象转化为 OpenVINO™ 的模型对象，开发者可以不需要将 ONNX 模型作为中间过渡。

import torchvisionimport torchfrom 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 torchvisionimport torchfrom 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

· list or tuple with tensors (openvino.runtime.Tensor / torch.Tensor / np.ndarray)

· dictionary where 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"The same meaning: {'yes' if result == 1 else '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 .The same meaning: 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 模型的转换。