HuggingFace Transformers 模型导出到 ONNX 格式

<think>嗯，用户想了解如何将机器学习模型导出为ONNX格式。让我先理清楚已有的信息和引用内容。根据提供的引用，有几个方法可以导出ONNX模型，比如使用scikit-learn的skl2onnx库，或者HuggingFace的Optimum库，还有PyTorch自带的导出功能。 首先，用户的问题比较广泛，可能需要分步骤回答。我应该先介绍通用的步骤，再具体到不同的框架。比如，对于scikit-learn模型，可以用skl2onnx，而PyTorch模型可以用torch.onnx.export。同时，HuggingFace的transformers模型可以通过Optimum或transformers.onnx来处理。 然后，我需要确保每个步骤都符合用户的要求，比如行内公式用$...$，独立公式用$$...$$。例如，提到线性回归模型时，可以插入公式$y = WX + b$。此外，要注意引用标识的正确添加，比如在提到安装依赖的地方引用[^1]，或者在优化kv_cache时引用[^4]。 可能用户还会关心导出时的常见问题，比如数据类型转换，如引用[5]提到的INT64错误，需要转换为INT32或FP32。这也是需要注意的点，应该在注意事项里提到。 另外，用户提供的示例中包括了Python代码，所以在回答中也应该包含具体的代码示例，比如scikit-learn导出ONNX的代码，以及PyTorch模型的导出代码。同时，要确保代码格式正确，使用反引号包裹。 还要考虑用户可能的后续问题，比如如何验证导出的模型是否正确，或者如何进行性能优化。这些可以作为相关问题列出，比如导出后的验证步骤，或者如何利用ONNX Runtime进行推理。 最后，需要检查所有引用是否正确对应，例如引用[1]是关于安装依赖，引用[2]是关于scikit-learn导出，引用[3]和[4]涉及HuggingFace模型和LLM导出，引用[5]处理转换中的错误。确保每个引用的地方都是自然衔接的，比如在提到模型部署优化时引用，在数据类型问题时引用[^5]。 总结一下，回答的结构应该是：简介ONNX，分框架介绍导出方法（scikit-learn、PyTorch、HuggingFace），注意事项，常见问题处理，最后相关问题。每个部分都插入对应的代码示例、公式和引用标识。</think>### 如何将机器学习模型导出为ONNX格式 #### 一、ONNX格式简介 ONNX（Open Neural Network Exchange）是一种开放的模型表示格式，支持跨框架和平台部署。通过ONNX，用户可以将PyTorch、TensorFlow、Scikit-learn等框架训练的模型转换为统一格式，方便在ONNX Runtime等推理引擎中运行[^2]。 --- #### 二、通用导出步骤 1. **安装依赖库** 根据框架选择安装对应工具： ```bash pip install onnx onnxruntime skl2onnx # 通用库 pip install torch # PyTorch模型 pip install optimum # HuggingFace模型[^3] ``` 2. **模型训练与准备** 以线性回归为例，模型公式为： $$y = WX + b$$ 训练完成后需确保模型已序列化（如`.pkl`文件）。 3. **导出为ONNX格式** **方法1：Scikit-learn模型** 使用`skl2onnx`库： ```python from skl2onnx import convert_sklearn from skl2onnx.common.data_types import FloatTensorType # 训练模型 from sklearn.linear_model import LinearRegression model = LinearRegression().fit(X_train, y_train) # 定义输入类型（假设输入维度为3） initial_type = [('float_input', FloatTensorType([None, 3]))] onnx_model = convert_sklearn(model, initial_types=initial_type) # 保存文件 with open("model.onnx", "wb") as f: f.write(onnx_model.SerializeToString()) ``` [^2] **方法2：PyTorch模型** 使用`torch.onnx.export`： ```python import torch # 定义示例输入 dummy_input = torch.randn(1, 3) # 输入维度为3 model = torch.nn.Linear(3, 1) # 定义模型 # 导出ONNX torch.onnx.export( model, dummy_input, "model.onnx", input_names=["input"], output_names=["output"], dynamic_axes={"input": {0: "batch_size"}} # 支持动态批次 ) ``` **方法3：HuggingFace Transformers模型** 使用`optimum`库： ```python from optimum.onnxruntime import ORTModelForSequenceClassification from transformers import AutoTokenizer, AutoModelForSequenceClassification model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased") tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") # 导出为ONNX model.save_pretrained("onnx_model/", use_onnx=True) ``` [^3] 4. **验证ONNX模型** 使用ONNX Runtime推理测试： ```python import onnxruntime as ort import numpy as np sess = ort.InferenceSession("model.onnx") input_name = sess.get_inputs()[0].name output = sess.run(None, {input_name: np.random.randn(1,3).astype(np.float32)}) ``` --- #### 三、注意事项 1. **数据类型兼容性** ONNX不支持INT64类型权重，需强制转换为INT32或FP32（参考代码）： ```python from onnx import helper, numpy_helper # 遍历模型节点替换INT64类型 ``` 2. **优化导出结构** - 对复杂模型（如LLM），需简化冗余操作（如Qwen模型的`rearrange`操作）。 - 使用`opset_version`参数指定兼容的操作集版本。 --- #### 四、常见问题 1. **动态维度支持** 在导出时通过`dynamic_axes`参数定义可变维度（如动态批次大小）。 2. **输入/输出对齐** 确保导出时的输入名称和维度与原始模型一致，避免推理错误。 ---