【Hugging Face】Optimum 库的常用方法和属性

Hugging Face Optimum 的常用方法和属性

Optimum 是 Hugging Face 优化 Transformer 训练和推理的库，适用于 ONNX Runtime、TensorRT、OpenVINO、Habana Gaudi、Intel Neural Compressor 等高效计算框架，可 提高推理速度、降低显存占用、支持量化和加速训练。

1. Optimum 的常见属性

在加载 Optimum 支持的模型后，可以访问以下属性：

from optimum.onnxruntime import ORTModelForSequenceClassification

model = ORTModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english", export=True)

2. Optimum 的常用方法

3. Optimum 详细用法

3.1. 使用 ONNX Runtime 进行推理

from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer

# 加载 ONNX 版本的 BERT
model = ORTModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english", export=True)
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")

# 处理输入
inputs = tokenizer("Hugging Face is great!", return_tensors="pt")

# 进行推理
outputs = model(**inputs)
print(outputs.logits)

3.2. 使用 TensorRT 进行加速

适用于 NVIDIA GPU。

from optimum.tensorrt import TRTModelForSequenceClassification

# 加载 TensorRT 版本的 BERT
model = TRTModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")

3.3. 在 Intel CPU/NPU 上使用 OpenVINO

from optimum.openvino import OVModelForSequenceClassification

# 加载 OpenVINO 版本的 BERT
model = OVModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english", export=True)

3.4. 量化 Transformer 模型

量化（INT8、FP16）可以 减少显存占用，提高推理速度。

from optimum.intel.neural_compressor import INCQuantizer

# 加载量化器
quantizer = INCQuantizer.from_pretrained("bert-base-uncased")

# 量化模型
quantized_model = quantizer.quantize(dtype="int8")

# 保存量化后的模型
quantized_model.save_pretrained("./quantized_model")

3.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()

3.6. 释放 GPU 内存

model.to("cpu")
torch.cuda.empty_cache()

3.7. 保存和加载优化后的模型

保存

model.save_pretrained("./optimized_model")

重新加载

from optimum.onnxruntime import ORTModelForSequenceClassification

model = ORTModelForSequenceClassification.from_pretrained("./optimized_model")

4. Optimum 适用于哪些任务？

5. Optimum vs transformers

6. 总结

Optimum 是 Hugging Face 提供的 优化 Transformer 训练和推理的库，支持 ONNX, TensorRT, OpenVINO, Habana Gaudi 等高效计算框架。

常见方法：

• ORTModelForSequenceClassification.from_pretrained(model_id, export=True) 使用 ONNX Runtime
• TRTModelForSequenceClassification.from_pretrained(model_id) 使用 TensorRT
• OVModelForSequenceClassification.from_pretrained(model_id, export=True) 使用 OpenVINO
• INCQuantizer.from_pretrained("bert-base-uncased").quantize() 进行模型量化
• GaudiTrainer(model, args, train_dataset) 在 Habana Gaudi 设备上训练
• model.save_pretrained(path) 保存优化模型
• model.to(device) 移动到 GPU/CPU 设备
• optimize_model(model, optimization_level="O2") 自动优化 Transformer 模型
• quantize_model(model, dtype="int8") 进行模型量化

如果你的 推理速度较慢、显存不足，或者希望优化计算，推荐使用 Optimum。