Hugging Face Optimum
项目
Optimum
是 Hugging Face 提供的 优化和加速工具,用于 高效推理、量化、图优化、多硬件支持,适用于 CPU、GPU、TPU、NPU 和专用 AI 硬件(如 Habana Gaudi, Intel Neural Compressor, OpenVINO, ONNX Runtime, TensorRT)。
1. Optimum
的核心功能
功能 | 作用 |
---|---|
ONNX Runtime | 将 Transformer 模型转换为 ONNX 格式,优化推理速度 |
TensorRT | 在 NVIDIA GPU 上进行高效推理 |
OpenVINO | 在 Intel CPU 和 NPU 上优化推理 |
Intel Neural Compressor | 进行 混合精度推理 和 量化(INT8, FP16) |
Habana Gaudi | 在 Gaudi AI 处理器上进行训练和推理 |
Graphcore IPU | 在 Graphcore 设备上优化 Transformer 运行 |
2. 安装 Optimum
Optimum
需要根据不同的硬件后端安装相应的依赖:
pip install optimum
如果需要 ONNX Runtime
:
pip install optimum[onnxruntime]
如果需要 TensorRT
:
pip install optimum[tensorrt]
如果需要 OpenVINO
:
pip install optimum[openvino]
3. Optimum
的核心模块
模块 | 作用 |
---|---|
optimum.onnxruntime | ONNX Runtime 加速 |
optimum.tensorrt | TensorRT 加速 |
optimum.intel | Intel Neural Compressor 量化 |
optimum.openvino | OpenVINO 优化 |
optimum.habana | Habana Gaudi 支持 |
optimum.graphcore | Graphcore IPU 支持 |
4. Optimum
详细用法
4.1. 使用 ONNX Runtime
进行加速
ONNX 可以大幅提高 Transformer 推理速度,适用于 CPU 和 NVIDIA GPU。
from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer
model_id = "distilbert-base-uncased-finetuned-sst-2-english"
# 加载 ONNX 版本的模型
model = ORTModelForSequenceClassification.from_pretrained(model_id, export=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)
# 推理
inputs = tokenizer("Hugging Face is amazing!", return_tensors="pt")
outputs = model(**inputs)
print(outputs.logits)
4.2. 使用 TensorRT
进行加速
TensorRT 适用于 NVIDIA GPU,优化 Transformer 计算效率。
from optimum.tensorrt import TRTModelForSequenceClassification
# 加载 TensorRT 版本的模型
model = TRTModelForSequenceClassification.from_pretrained(model_id)
4.3. 在 Intel 设备上使用 OpenVINO
from optimum.openvino import OVModelForSequenceClassification
model = OVModelForSequenceClassification.from_pretrained(model_id, export=True)
4.4. 量化 Transformer 模型
量化(INT8, FP16)可以 减少显存占用,提高推理速度。
from optimum.intel.neural_compressor import INCQuantizer
quantizer = INCQuantizer.from_pretrained("bert-base-uncased")
quantized_model = quantizer.quantize()
quantized_model.save_pretrained("./quantized_model")
4.5. 在 Habana Gaudi 设备上运行
Habana Gaudi 是 AI 训练专用加速器,支持 Transformer 训练和推理:
from optimum.habana import GaudiTrainer, GaudiConfig
gaudi_config = GaudiConfig()
trainer = GaudiTrainer(model=model, args=training_args, train_dataset=train_dataset)
trainer.train()
5. Optimum
适用于哪些任务?
任务 | 适用 |
---|---|
文本分类 (AutoModelForSequenceClassification ) | 是 |
问答任务 (AutoModelForQuestionAnswering ) | 是 |
命名实体识别 (AutoModelForTokenClassification ) | 是 |
机器翻译 (AutoModelForSeq2SeqLM ) | 是 |
文本生成 (AutoModelForCausalLM ) | 部分支持 |
语音识别 (AutoModelForSpeechSeq2Seq ) | 是 |
6. Optimum
vs transformers
方案 | 适用场景 | 主要作用 |
---|---|---|
transformers | 普通训练和推理 | 加载和微调模型 |
optimum.onnxruntime | CPU/GPU 推理加速 | 转换 ONNX 格式,提高推理速度 |
optimum.tensorrt | NVIDIA GPU 加速 | 使用 TensorRT 进行高效推理 |
optimum.intel | Intel CPU/NPU 加速 | 量化 Transformer 模型(INT8) |
optimum.openvino | Intel 设备优化 | 兼容 OpenVINO 计算框架 |
optimum.habana | Habana Gaudi 设备 | AI 训练加速 |
optimum.graphcore | Graphcore IPU | 计算优化 |
7. 总结
Optimum
是 Hugging Face 提供的 专门优化 Transformer 训练和推理的库,支持 ONNX, TensorRT, OpenVINO, Habana Gaudi 等高效计算框架。
常见方法:
ORTModelForSequenceClassification.from_pretrained(model_id, export=True)
使用 ONNX RuntimeTRTModelForSequenceClassification.from_pretrained(model_id)
使用 TensorRTOVModelForSequenceClassification.from_pretrained(model_id, export=True)
使用 OpenVINOINCQuantizer.from_pretrained("bert-base-uncased").quantize()
进行模型量化GaudiTrainer(model, args, train_dataset)
在 Habana Gaudi 设备上训练
如果你的 推理速度较慢、显存不足,或者希望优化计算,推荐使用 Optimum
。