使用ONNX/tensorflow-onnx将HuggingFace BERT模型转换为ONNX格式

【免费下载链接】tensorflow-onnx 项目地址: https://gitcode.com/gh_mirrors/ten/tensorflow-onnx

前言

在深度学习领域，模型格式的转换是一个常见需求，特别是当我们需要将训练好的模型部署到不同平台时。ONNX（Open Neural Network Exchange）作为一种开放的模型格式，能够帮助我们在不同框架之间转换和部署模型。本文将详细介绍如何使用tensorflow-onnx工具将HuggingFace的BERT问答模型转换为ONNX格式。

准备工作

环境配置

首先需要安装必要的Python包：

pip install tensorflow transformers tf2onnx onnxruntime

这些包的作用分别是：

• tensorflow：深度学习框架
• transformers：HuggingFace的Transformer模型库
• tf2onnx：TensorFlow到ONNX的转换工具
• onnxruntime：ONNX模型的推理运行时

禁用GPU

为了确保示例能在所有环境中运行，我们显式禁用GPU：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = ""

加载HuggingFace BERT模型

我们使用HuggingFace提供的TFBertForQuestionAnswering模型，这是一个基于TensorFlow实现的BERT问答模型：

from transformers import BertTokenizer, TFBertForQuestionAnswering

# 加载预训练的tokenizer和模型
tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
model = TFBertForQuestionAnswering.from_pretrained('bert-base-cased')

# 准备问答输入
question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
input_dict = tokenizer(question, text, return_tensors='tf')

# 运行模型获取结果
tf_results = model(input_dict)

这段代码会输出模型的预测结果，包括开始位置和结束位置的logits值。

模型转换到ONNX

定义输入规格

ONNX转换需要明确指定模型的输入规格。对于BERT模型，我们需要定义三个输入张量：

input_spec = (
    tf.TensorSpec((None, None), tf.int32, name="input_ids"),
    tf.TensorSpec((None, None), tf.int32, name="token_type_ids"),
    tf.TensorSpec((None, None), tf.int32, name="attention_mask")
)

• input_ids：输入token的ID
• token_type_ids：区分不同句子的标记
• attention_mask：注意力掩码，标识哪些位置是有效输入

执行转换

使用tf2onnx进行转换：

_, _ = tf2onnx.convert.from_keras(
    model, 
    input_signature=input_spec, 
    opset=13, 
    output_path="bert.onnx"
)

参数说明：

• opset=13：指定ONNX操作集版本
• output_path：输出的ONNX模型路径

验证ONNX模型

使用ONNX Runtime加载模型

import onnxruntime as rt

# 获取输出名称
output_names = list(tf_results.keys())

# 准备输入数据（转换为numpy数组）
input_dict_np = {k: v.numpy() for k, v in input_dict.items()}

# 创建ONNX Runtime会话
opt = rt.SessionOptions()
sess = rt.InferenceSession("bert.onnx")
onnx_results = sess.run(output_names, input_dict_np)

结果对比

为了确保转换后的模型与原始TensorFlow模型行为一致，我们需要对比两者的输出：

for i, name in enumerate(output_names):
    np.testing.assert_allclose(
        tf_results[name], 
        onnx_results[i], 
        rtol=1e-5, 
        atol=1e-5
    )

如果对比通过，说明ONNX模型转换成功，且与原始模型行为一致。

技术细节解析

1. 输入规格定义：BERT模型的输入包含三个部分，这在转换时需要明确指定。不同BERT变体可能有不同的输入要求。

2. OPSET版本：ONNX的操作集版本会影响模型的兼容性。选择13是因为它支持大多数现代神经网络操作。

3. 精度对比：由于不同框架实现细节可能略有差异，我们使用相对宽松的容差（1e-5）来验证结果一致性。

4. 模型输出：BERT问答模型的输出包含开始和结束位置的logits，这些值用于确定答案在文本中的位置。

常见问题

1. 转换失败：可能是由于某些操作不被ONNX支持，可以尝试更新tf2onnx或调整opset版本。

2. 精度差异：如果结果差异较大，可以尝试降低容差或检查模型是否在转换过程中被优化。

3. 性能问题：ONNX Runtime提供了多种优化选项，可以通过配置SessionOptions来优化推理性能。

扩展应用

这种方法不仅适用于BERT问答模型，也可以应用于HuggingFace提供的其他Transformer模型，如GPT、RoBERTa等。只需替换模型加载部分，并相应调整输入规格即可。

总结

本文详细介绍了如何使用tensorflow-onnx将HuggingFace的BERT模型转换为ONNX格式。通过这种转换，我们可以将训练好的模型部署到支持ONNX的各种推理引擎上，实现跨平台部署。关键步骤包括模型加载、输入规格定义、转换执行和结果验证。掌握这一流程对于模型部署和跨框架应用开发具有重要意义。

【免费下载链接】tensorflow-onnx 项目地址: https://gitcode.com/gh_mirrors/ten/tensorflow-onnx