PyTorch模型转ONNX的3种方法

每天都有数百个 Transformers 实验和模型上传到 Hugging Face Hub。进行这些实验的机器学习工程师和学生使用各种框架，例如 PyTorch、TensorFlow/Keras 等。这些模型已被数千家公司使用，并构成了 AI 驱动产品的基础。

如果你在生产环境中部署 Transformers 模型，我们建议先将它们导出为序列化格式，以便在专用运行时和硬件上加载、优化和执行。

在本指南中，我们将了解：

• 什么是 ONNX？
• 什么是 Hugging Face Optimum？
• 支持哪些 Transformers 架构？
• 如何将 Transformers 模型 (BERT) 转换为 ONNX？
• 下一步是什么？

让我们开始吧！

NSDT工具推荐： Three.js AI纹理开发包 - YOLO合成数据生成器 - GLTF/GLB在线编辑 - 3D模型格式在线转换 - 可编程3D场景编辑器 - REVIT导出3D模型插件 - 3D模型语义搜索引擎 - Three.js虚拟轴心开发包 - 3D模型在线减面 - STL模型在线切割 

1、什么是 ONNX？

ONNX（开放神经网络交换）是一种用于表示机器学习模型的开放标准和格式。 ONNX 定义了一组通用运算符和一种通用文件格式，用于表示各种框架中的深度学习模型，包括 PyTorch 和 TensorFlow。

当将模型导出为 ONNX 格式时，这些运算符用于构建计算图（通常称为中间表示），该图表示数据通过神经网络的流动。

重要提示： ONNX 不是运行时 ONNX 只是可以与 ONNX Runtime 等运行时一起使用的表示。你可以在 此 处找到受支持的加速器列表。

2、什么是 Hugging Face Optimum？

Hugging Face Optimum 是一个开源库，也是 Hugging Face Transformers 的扩展，它提供了统一的性能优化工具 API，以实现在加速硬件上训练和运行模型的最高效率，包括用于在 Graphcore IPU 和 Habana Gaudi 上优化性能的工具包。

Optimum 可用于转换、量化、图形优化、加速训练和推理，并且支持transformers管线。

下面你可以看到如何利用 ONNX 的 Optimum 的典型客户旅程：

可以在 Transformers 文档的 ONNX 部分找到所有受支持的 Transformers 架构的列表。可以转换为 ONNX 并使用 Hugging Face Optimum 进行优化。

3、如何将 Transformers 模型转换为 ONNX？

目前有三种方法可以将 Hugging Face Transformers 模型转换为 ONNX。在本节中，我们将学习如何使用所有三种方法导出 distilbert-base-uncased-finetuned-sst-2-english 进行文本分类，从低级 torch API 到最用户友好的高级 API optimal。每种方法都完全相同

3.1 使用 torch.onnx 导出（低级）

torch.onnx 使你能够通过导出方法将模型检查点转换为 ONNX 图。但你必须提供很多值，例如 input_names、 dynamic_axes 等。

首先需要安装一些依赖项：

pip install transformers torch

使用 export 导出我们的检查点：

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
# load model and tokenizer
model_id = "distilbert-base-uncased-finetuned-sst-2-english"
model = AutoModelForSequenceClassification.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
dummy_model_input = tokenizer("This is a sample", return_tensors="pt")
 
# export
torch.onnx.export(
    model,
    tuple(dummy_model_input.values()),
    f="torch-model.onnx",
    input_names=['input_ids', 'attention_mask'],
    output_names=['logits'],
    dynamic_axes={'input_ids': {0: 'batch_size', 1: 'sequence'},
                  'attention_mask': {0: 'batch_size', 1: 'sequence'},
                  'logits': {0: 'batch_size', 1: 'sequence'}},
    do_constant_folding=True,
    opset_version=13,
)

3.2 使用 transformers.onnx 导出（中级）

transformers.onnx 可让你利用配置对象将模型检查点转换为 ONNX 图。这样，你就不必为 dynamic_axes 等提供复杂的配置。

首先需要安装一些依赖项：

pip install transformers[onnx] torch

使用 transformers.onnx 导出我们的检查点：

from pathlib import Path
import transformers
from transformers.onnx import FeaturesManager
from transformers import AutoConfig, AutoTokenizer, AutoModelForSequenceClassification
 
# load model and tokenizer
model_id = "distilbert-base-uncased-finetuned-sst-2-english"
feature = "sequence-classification"
base_model = AutoModelForSequenceClassification.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
 
# load config
model_kind, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model, feature=feature)
onnx_config = model_onnx_config(model.config)
 
# export
onnx_inputs, onnx_outputs = transformers.onnx.export(
        preprocessor=tokenizer,
        model=model,
        config=onnx_config,
        opset=13,
        output=Path("trfs-model.onnx")
)

3.3 使用 Optimum 导出（高级）

Optimum Inference 包括使用 ORTModelForXxx 类将原始 Transformers 模型转换为 ONNX 的方法。要将 Transformers 模型转换为 ONNX，只需将 from_transformers=True 传递给 from_pretrained() 方法，你的模型就会加载并转换为 ONNX，并利用底层的 transformers.onnx 包。

首先需要安装一些依赖项：

pip install optimum[onnxruntime]

使用 ORTModelForSequenceClassification 导出检查点：

from optimum.onnxruntime import ORTModelForSequenceClassification
 
model = ORTModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english",from_transformers=True)

使用 Optimum 进行转换的最好之处在于，你可以立即使用该模型来运行预测或将其加载到管线中。

---

原文链接：PyTorch转ONNX 3种方法 - BimAnt