Optimum 文档
快速入门
并获得增强的文档体验
开始使用
快速入门
本快速入门指南旨在帮助那些准备好深入代码并查看如何将 🤗 Optimum 集成到其模型训练和推理工作流程中的开发者。
加速推理
OpenVINO
要加载模型并使用 OpenVINO Runtime 运行推理,您只需将 AutoModelForXxx
类替换为相应的 OVModelForXxx
类。如果您想加载 PyTorch 检查点,请设置 export=True
将您的模型转换为 OpenVINO IR (中间表示)。
- from transformers import AutoModelForSequenceClassification
+ from optimum.intel.openvino import OVModelForSequenceClassification
from transformers import AutoTokenizer, pipeline
# Download a tokenizer and model from the Hub and convert to OpenVINO format
tokenizer = AutoTokenizer.from_pretrained(model_id)
model_id = "distilbert-base-uncased-finetuned-sst-2-english"
- model = AutoModelForSequenceClassification.from_pretrained(model_id)
+ model = OVModelForSequenceClassification.from_pretrained(model_id, export=True)
# Run inference!
classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
results = classifier("He's a dreadful magician.")
ONNX Runtime
为了使用 ONNX Runtime 加速推理,🤗 Optimum 使用配置对象来定义图优化和量化的参数。然后,这些对象被用于实例化专用的优化器和量化器。
在应用量化或优化之前,我们首先需要加载我们的模型。要加载模型并使用 ONNX Runtime 运行推理,您只需将标准的 Transformers AutoModelForXxx
类替换为相应的 ORTModelForXxx
类。如果您想从 PyTorch 检查点加载,请设置 export=True
将您的模型导出为 ONNX 格式。
>>> from optimum.onnxruntime import ORTModelForSequenceClassification
>>> from transformers import AutoTokenizer
>>> model_checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
>>> save_directory = "tmp/onnx/"
>>> # Load a model from transformers and export it to ONNX
>>> tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
>>> ort_model = ORTModelForSequenceClassification.from_pretrained(model_checkpoint, export=True)
>>> # Save the ONNX model and tokenizer
>>> ort_model.save_pretrained(save_directory)
>>> tokenizer.save_pretrained(save_directory)
现在让我们看看如何使用 ONNX Runtime 应用动态量化
>>> from optimum.onnxruntime.configuration import AutoQuantizationConfig
>>> from optimum.onnxruntime import ORTQuantizer
>>> # Define the quantization methodology
>>> qconfig = AutoQuantizationConfig.arm64(is_static=False, per_channel=False)
>>> quantizer = ORTQuantizer.from_pretrained(ort_model)
>>> # Apply dynamic quantization on the model
>>> quantizer.quantize(save_dir=save_directory, quantization_config=qconfig)
在此示例中,我们量化了来自 Hugging Face Hub 的一个模型,同样地,我们也可以通过提供包含模型权重的目录路径来量化本地托管的模型。应用 quantize()
方法的结果是一个 model_quantized.onnx
文件,可用于运行推理。以下是如何加载 ONNX Runtime 模型并用它生成预测的示例
>>> from optimum.onnxruntime import ORTModelForSequenceClassification
>>> from transformers import pipeline, AutoTokenizer
>>> model = ORTModelForSequenceClassification.from_pretrained(save_directory, file_name="model_quantized.onnx")
>>> tokenizer = AutoTokenizer.from_pretrained(save_directory)
>>> classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
>>> results = classifier("I love burritos!")
加速训练
Habana
为了在 Habana 的 Gaudi 处理器上训练 transformers,🤗 Optimum 提供了一个 GaudiTrainer
,它与 🤗 Transformers 的 Trainer 非常相似。这里是一个简单的例子
- from transformers import Trainer, TrainingArguments
+ from optimum.habana import GaudiTrainer, GaudiTrainingArguments
# Download a pretrained model from the Hub
model = AutoModelForXxx.from_pretrained("bert-base-uncased")
# Define the training arguments
- training_args = TrainingArguments(
+ training_args = GaudiTrainingArguments(
output_dir="path/to/save/folder/",
+ use_habana=True,
+ use_lazy_mode=True,
+ gaudi_config_name="Habana/bert-base-uncased",
...
)
# Initialize the trainer
- trainer = Trainer(
+ trainer = GaudiTrainer(
model=model,
args=training_args,
train_dataset=train_dataset,
...
)
# Use Habana Gaudi processor for training!
trainer.train()
ONNX Runtime
要使用 ONNX Runtime 的加速功能训练 transformers,🤗 Optimum 提供了一个 ORTTrainer
,它与 🤗 Transformers 的 Trainer 非常相似。这里是一个简单的例子
- from transformers import Trainer, TrainingArguments
+ from optimum.onnxruntime import ORTTrainer, ORTTrainingArguments
# Download a pretrained model from the Hub
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")
# Define the training arguments
- training_args = TrainingArguments(
+ training_args = ORTTrainingArguments(
output_dir="path/to/save/folder/",
optim="adamw_ort_fused",
...
)
# Create a ONNX Runtime Trainer
- trainer = Trainer(
+ trainer = ORTTrainer(
model=model,
args=training_args,
train_dataset=train_dataset,
+ feature="text-classification", # The model type to export to ONNX
...
)
# Use ONNX Runtime for training!
trainer.train()
开箱即用的 ONNX 导出
Optimum 库为 Transformers 和 Diffusers 模型提供了开箱即用的 ONNX 导出功能!
将模型导出到 ONNX 非常简单
optimum-cli export onnx --model gpt2 gpt2_onnx/
查看帮助以获取更多选项
optimum-cli export onnx --help
请查看文档了解更多信息。
torch.fx 集成
Optimum 与 torch.fx
集成,提供了一行代码即可实现的多种图变换。我们的目标是通过 torch.fx
更好地支持量化,包括量化感知训练 (QAT) 和训练后量化 (PTQ)。