本文将会介绍如何使用HuggingFace的Optimum，来对微调后的BERT模型进行量化（Quantization）。
在文章NLP（六十七）BERT模型训练后动态量化（PTDQ）中，我们使用PyTorch自带的PTDQ（Post Training Dynamic Quantization）量化策略对微调后的BERT模型进行量化，取得了模型推理性能的提升（大约1.5倍）。本文将尝试使用Optimum量化工具。
Optimum介绍
Optimum 是 Transformers 的扩展，它提供了一组性能优化工具，可以在目标硬件上以最高效率训练和运行模型。
Optimum针对不同的硬件，提供了不同的优化方案，如下表：
| 硬件 | 安装命令 |
|---|---|
| ONNX runtime | python -m pip install optimum[onnxruntime] |
| Intel Neural Compressor (INC) | python -m pip install optimum[neural-compressor] |
| Intel OpenVINO | python -m pip install optimum[openvino,nncf] |
| Graphcore IPU | python -m pip install optimum[graphcore] |
| Habana Gaudi Processor (HPU) | python -m pip install optimum[habana] |
| GPU | python -m pip install optimum[onnxruntime-gpu] |
本文将会介绍基于ONNX的模型量化技术。ONNX（英语：Open Neural Network Exchange）是一种针对机器学习所设计的开放式的文件格式，用于存储训练好的模型。它使得不同的人工智能框架（如Pytorch、MXNet）可以采用相同格式存储模型数据并交互。
模型量化
我们使用的微调后的BERT模型采用文章NLP（六十六）使用HuggingFace中的Trainer进行BERT模型微调中给出的文本分类模型。
首先，我们先加载PyTorch中的设备（CPU）。
# load device
import torch
device = torch.device("cpu")
接着，我们使用optimum.onnxruntime模块加载模型和tokenizer，并将模型保存为onnx格式。
from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer
import torch
model_id = "./sougou_test_trainer_256/checkpoint-96"
onnx_path = "./sougou_test_trainer_256/onnx_256"
# load vanilla transformers and convert to onnx
model = ORTModelForSequenceClassification.from_pretrained(model_id, from_transformers=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)
# save onnx checkpoint and tokenizer
model.save_pretrained(onnx_path)
tokenizer.save_pretrained(onnx_path)
此时，会多出onnx_256文件夹，保存模型为model.onnx。
输出结果为：
('./sougou_test_trainer_256/onnx_256\\tokenizer_config.json',
'./sougou_test_trainer_256/onnx_256\\special_tokens_map.json',
'./sougou_test_trainer_256/onnx_256\\vocab.txt',
'./sougou_test_trainer_256/onnx_256\\added_tokens.json',
'./sougou_test_trainer_256/onnx_256\\tokenizer.json')
使用transfomers中的pipeline对模型进行快速推理。
from transformers import pipeline
vanilla_clf = pipeline("text-classification", model=model, tokenizer=tokenizer)
vanilla_clf("这期节目继续关注中国篮球的话题。众所周知，我们已经结束了男篮世界杯的所有赛程，一胜四负的一个成绩，甚至比上一届的世界杯成绩还要差。因为这一次我们连奥运会落选赛也都没有资格参加，所以，连续两次错过了巴黎奥运会的话，对于中国篮协，还有对于姚明来说，确实成为了他任职的一个最大的败笔。对于球迷非常关注的一个话题，乔尔杰维奇是否下课，可能对于这个悬念来说也都是暂时有答案了。")
输出结果如下：
[{'label': 'LABEL_0', 'score': 0.9963239431381226}]
对ONNX模型进行优化。
from optimum.onnxruntime import ORTOptimizer
from optimum.onnxruntime.configuration import OptimizationConfig
# create ORTOptimizer and define optimization configuration
optimizer = ORTOptimizer.from_pretrained(model)
optimization_config = OptimizationConfig(optimization_level=99) # enable all optimizations
# apply the optimization configuration to the model
optimizer.optimize(
save_dir=onnx_path,
optimization_config=optimization_config,
)
此时，优化后的模型为model_optimized.onnx。
对优化后的模型进行推理。
from transformers import pipeline
# load optimized model
optimized_model = ORTModelForSequenceClassification.from_pretrained(onnx_path, file_name="model_optimized.onnx")
# create optimized pipeline
optimized_clf = pipeline("text-classification", model=optimized_model, tokenizer=tokenizer)
optimized_clf("今年7月，教育部等四部门联合印发了《关于在深化非学科类校外培训治理中加强艺考培训规范管理的通知》（以下简称《通知》）。《通知》针对近年来校外艺术培训的状况而发布，并从源头就校外艺术培训机构的“培训主体、从业人员、招生行为、安全底线”等方面进行严格规范。校外艺术培训之所以火热，主要在于高中阶段艺术教育发展迟滞于学生需求。分析教育部数据，2021年艺术学科在校生占比为9.84%，高于2020年的9.73%；2020至2021年艺术学科在校生的年增长率为5.04%，远高于4.28%的总在校生年增长率。增长的数据，是近年来艺考招生连年火热的缩影，在未来一段时间内，艺考或将在全国范围内继续保持高热度。")
输出结果为：
[{'label': 'LABEL_3', 'score': 0.9926980137825012}]
对优化后的ONNX模型再进行量化，代码为：
from optimum.onnxruntime import ORTQuantizer
from optimum.onnxruntime.configuration import AutoQuantizationConfig
# create ORTQuantizer and define quantization configuration
dynamic_quantizer = ORTQuantizer.from_pretrained(optimized_model)
dqconfig = AutoQuantizationConfig.avx2(is_static=False, per_channel=False)
# apply the quantization configuration to the model
model_quantized_path = dynamic_quantizer.quantize(
save_dir=onnx_path,
quantization_config=dqconfig,
)
此时量化后的模型为model_optimized_quantized.onnx。比较量化前后的模型大小，代码为：
import os
# get model file size
size = os.path.getsize(os.path.join(onnx_path, "model_optimized.onnx"))/(1024*1024)
quantized_model = os.path.getsize(os.path.join(onnx_path, "model_optimized_quantized.onnx"))/(1024*1024)
print(f"Model file size: {size:.2f} MB")
print(f"Quantized Model file size: {quantized_model:.2f} MB")
输出结果为：
Model file size: 390.17 MB
Quantized Model file size: 97.98 MB
最后，加载量化后的模型，代码为：
# load quantization model
from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import pipeline, AutoTokenizer
quantized_model = ORTModelForSequenceClassification.from_pretrained(onnx_path, file_name="model_optimized_quantized.onnx").to(device)
tokenizer = AutoTokenizer.from_pretrained(onnx_path)
推理实验
在进行模型推理实验前，先加载测试数据集。
import pandas as pd
test_df = pd.read_csv("./data/sougou/test.csv")
使用量化前的模型进行推理，记录推理时间，代码如下：
# original model evaluate
import numpy as np
import time
cost_time_list = []
s_time = time.time()
true_labels, pred_labels = [], []
for i, row in test_df.iterrows():
row_s_time = time.time()
true_labels.append(row["label"])
encoded_text = tokenizer(row['text'], max_length=256, truncation=True, padding=True, return_tensors='pt')
# print(encoded_text)
logits = model(**encoded_text)
label_id = np.argmax(logits[0].detach().numpy(), axis=1)[0]
pred_labels.append(label_id)
cost_time_list.append((time.time() - row_s_time) * 1000)
if i % 100:
print(i, (time.time() - row_s_time) * 1000, label_id)
print("avg time:", (time.time() - s_time) * 1000 / test_df.shape[0])
print("P50 time:", np.percentile(np.array(cost_time_list), 50))
print("P95 time:", np.percentile(np.array(cost_time_list), 95))
输出结果为：
0 710.2577686309814 0
100 477.72765159606934 1
200 616.3530349731445 2
300 509.63783264160156 3
400 531.57639503479 4
avg time: 501.0757282526806
P50 time: 504.6522617340088
P95 time: 623.9353895187337
对输出结果进行指标评级，代码为：
from sklearn.metrics import classification_report
print(classification_report(true_labels, pred_labels, digits=4))
重复上述代码，将模型替换为量化前ONNX模型（model.onnx），优化后ONNX模型（model_oprimized.onnx），量化后ONNX模型（model_optimized_quantized.onnx），进行推理时间（单位：ms）统计和推理指标评估，结果见下表：
| 模型 | 平均推理时间 | P95推理时间 | weighted F1 |
|---|---|---|---|
| 量化前ONNX模型 | 501.1 | 623.9 | 0.9717 |
| 优化后ONNX模型 | 484.6 | 629.6 | 0.9717 |
| 量化后ONNX模型 | 361.5 | 426.9 | 0.9738 |
对比文章NLP（六十七）BERT模型训练后动态量化（PTDQ）中的推理结果，原始模型的平均推理时间为666.6ms，weighted F1值为0.9717，我们有如下结论：
- ONNX模型不影响推理效果，但在平均推理时间上提速约1.33倍
- 优化ONNX模型不影响推理效果，但在平均推理时间上提速约1.38倍
- 量化后的ONNX模型影响推理效果，一般会略有下降，本次实验结果为提升，但在平均推理时间上提速约1.84倍，由于PyTorch的PTDQ（模型训练后动态量化）
总结
本文介绍了如何使用HuggingFace的Optimum，来对微调后的BERT模型进行量化（Quantization），在optimum.onnxruntime模块中，平均推理时间提速约1.8倍。
本文已开源至Github，网址为：https://github.com/percent4/dynamic_quantization_on_bert 。
本文已开通个人博客，欢迎大家访问：https://percent4.github.io/ 。
欢迎关注我的公众号NLP奇幻之旅，原创技术文章第一时间推送。文章来源：https://www.toymoban.com/news/detail-697347.html
欢迎关注我的知识星球“自然语言处理奇幻之旅”，笔者正在努力构建自己的技术社区。文章来源地址https://www.toymoban.com/news/detail-697347.html
- NLP（六十六）使用HuggingFace中的Trainer进行BERT模型微调：https://blog.csdn.net/jclian91/article/details/132644042
- NLP（六十七）BERT模型训练后动态量化（PTDQ）：https://blog.csdn.net/jclian91/article/details/132644042
- Optimum: https://huggingface.co/docs/optimum/index
- Optimizing Transformers with Hugging Face Optimum: https://www.philschmid.de/optimizing-transformers-with-optimum
到了这里，关于NLP（六十八）使用Optimum进行模型量化的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！
