使用Optimum将BART-large-mnli模型转换为ONNX格式实现零样本分类

背景介绍

Optimum是Hugging Face推出的一个优化库，专门用于在各种硬件和运行时环境下高效运行Transformer模型。其中一项重要功能是将PyTorch或TensorFlow模型转换为ONNX格式，以获得更好的推理性能。

本文将详细介绍如何将facebook/bart-large-mnli模型转换为ONNX格式，并用于零样本文本分类任务。

模型转换步骤

首先需要安装必要的依赖包：

pip install torch==2.1.2
pip install --upgrade-strategy eager install optimum[onnxruntime]

使用Optimum命令行工具进行模型转换：

optimum-cli export onnx --task zero-shot-classification --model facebook/bart-large-mnli bart-large-mnli_onnx/

这个命令会：

1. 下载facebook/bart-large-mnli预训练模型
2. 根据指定的zero-shot-classification任务
3. 将模型转换为ONNX格式
4. 保存到bart-large-mnli_onnx目录

正确使用转换后的模型

转换完成后，需要使用正确的模型类加载ONNX模型。对于零样本分类任务，应该使用ORTModelForSequenceClassification而不是ORTModelForQuestionAnswering。

from transformers import AutoTokenizer, pipeline
from optimum.onnxruntime import ORTModelForSequenceClassification

# 加载tokenizer和模型
tokenizer = AutoTokenizer.from_pretrained("bart-large-mnli_onnx")
model = ORTModelForSequenceClassification.from_pretrained("bart-large-mnli_onnx")

# 创建零样本分类pipeline
classifier = pipeline("zero-shot-classification", 
                     model=model, 
                     tokenizer=tokenizer)

# 示例使用
sequence = "What are your thoughts on current events?"
candidate_labels = ["technology", "healthcare", "social issues", "sports"]
results = classifier(sequence, candidate_labels)

常见问题解决

在尝试使用ONNX模型时，可能会遇到以下错误：

KeyError: 'start_logits'

这个错误通常是因为使用了错误的模型类。对于零样本分类任务：

• 错误做法：使用ORTModelForQuestionAnswering
• 正确做法：使用ORTModelForSequenceClassification

因为这两种任务有不同的输出结构：

• 问答模型输出start_logits和end_logits
• 序列分类模型输出分类logits

性能优化建议

使用ONNX格式模型相比原始PyTorch模型有以下优势：

1. 推理速度更快：ONNX Runtime针对推理进行了优化
2. 内存占用更低：模型权重以更高效的格式存储
3. 跨平台兼容性：可在多种硬件和操作系统上运行

对于生产环境部署，还可以进一步：

• 量化模型减小体积
• 使用ONNX Runtime的图优化
• 针对特定硬件进行调优

总结

通过Optimum库可以方便地将Hugging Face模型转换为ONNX格式。对于零样本分类任务，关键是要选择正确的模型类ORTModelForSequenceClassification。这种转换方式既保持了模型的强大能力，又获得了ONNX运行时的性能优势，是生产部署的理想选择。