本文将会介绍BERT模型训练后动态量化（Post Training Dynamic Quantization，PTDQ）。
量化
在深度学习中，量化（Quantization）指的是使用更少的bit来存储原本以浮点数存储的tensor，以及使用更少的bit来完成原本以浮点数完成的计算。这么做的好处主要有如下几点：
- 更少的模型体积，接近4倍的减少
- 可以更快地计算，由于更少的内存访问和更快的int8计算，可以快2~4倍
PyTorch中的模型参数默认以FP32精度储存。对于量化后的模型，其部分或者全部的tensor操作会使用int类型来计算，而不是使用量化之前的float类型。当然，量化还需要底层硬件支持，x86 CPU(支持AVX2)、ARM CPU、Google TPU、Nvidia Volta/Turing/Ampere、Qualcomm DSP这些主流硬件都对量化提供了支持。
PTDQ
PyTorch对量化的支持目前有如下三种方式：
- Post Training Dynamic Quantization：模型训练完毕后的动态量化
- Post Training Static Quantization：模型训练完毕后的静态量化
- QAT (Quantization Aware Training)：模型训练中开启量化
本文仅介绍Post Training Dynamic Quantization（PTDQ）。
对训练后的模型权重执行动态量化，将浮点模型转换为动态量化模型，仅对模型权重进行量化，偏置不会量化。默认情况下，仅对Linear和RNN变体量化 (因为这些layer的参数量很大，收益更高)。
torch.quantization.quantize_dynamic(model, qconfig_spec=None, dtype=torch.qint8, mapping=None, inplace=False)
参数解释：
- model：模型（默认为FP32）
- qconfig_spec：
- 集合：比如： qconfig_spec={nn.LSTM, nn.Linear} 。列出要量化的神经网络模块。
- 字典： qconfig_spec = {nn.Linear: default_dynamic_qconfig, nn.LSTM: default_dynamic_qconfig}
- dtype： float16 或 qint8
- mapping：就地执行模型转换，原始模块发生变异
- inplace：将子模块的类型映射到需要替换子模块的相应动态量化版本的类型
例子：
# -*- coding: utf-8 -*-
# 动态量化模型，只量化权重
import torch
from torch import nn
class DemoModel(torch.nn.Module):
def __init__(self):
super(DemoModel, self).__init__()
self.conv = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=1)
self.relu = nn.ReLU()
self.fc = torch.nn.Linear(2, 2)
def forward(self, x):
x = self.conv(x)
x = self.relu(x)
x = self.fc(x)
return x
if __name__ == "__main__":
model_fp32 = DemoModel()
# 创建一个量化的模型实例
model_int8 = torch.quantization.quantize_dynamic(model=model_fp32, # 原始模型
qconfig_spec={torch.nn.Linear}, # 要动态量化的算子
dtype=torch.qint8) # 将权重量化为：qint8
print(model_fp32)
print(model_int8)
# 运行模型
input_fp32 = torch.randn(1, 1, 2, 2)
output_fp32 = model_fp32(input_fp32)
print(output_fp32)
output_int8 = model_int8(input_fp32)
print(output_int8)
输出结果如下：
DemoModel(
(conv): Conv2d(1, 1, kernel_size=(1, 1), stride=(1, 1))
(relu): ReLU()
(fc): Linear(in_features=2, out_features=2, bias=True)
)
DemoModel(
(conv): Conv2d(1, 1, kernel_size=(1, 1), stride=(1, 1))
(relu): ReLU()
(fc): DynamicQuantizedLinear(in_features=2, out_features=2, dtype=torch.qint8, qscheme=torch.per_tensor_affine)
)
tensor([[[[0.3120, 0.3042],
[0.3120, 0.3042]]]], grad_fn=<AddBackward0>)
tensor([[[[0.3120, 0.3042],
[0.3120, 0.3042]]]])
模型量化策略
当前，由于量化算子的覆盖有限，因此，对于不同的深度学习模型，其量化策略不同，见下表：
| 模型 | 量化策略 | 原因 |
|---|---|---|
| LSTM/RNN | Dynamic Quantization | 模型吞吐量由权重的计算/内存带宽决定 |
| BERT/Transformer | Dynamic Quantization | 模型吞吐量由权重的计算/内存带宽决定 |
| CNN | Static Quantization | 模型吞吐量由激活函数的内存带宽决定 |
| CNN | Quantization Aware Training | 模型准确率不能由Static Quantization获取的情况 |
下面对BERT模型进行训练后动态量化，分析模型在量化前后，推理效果和推理性能的变化。
实验
我们使用的训练后的模型为中文文本分类模型，其训练过程可以参考文章：NLP（六十六）使用HuggingFace中的Trainer进行BERT模型微调 。
训练后的BERT模型动态量化实验的设置如下：
- base model: bert-base-chinese
- CPU info: x86-64, Intel® Core™ i5-10210U CPU @ 1.60GHz
- batch size: 1
- thread: 1
具体的实验过程如下：
- 加载模型及tokenizer
import torch
from transformers import AutoModelForSequenceClassification
MAX_LENGTH = 128
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
checkpoint = f"./sougou_test_trainer_{MAX_LENGTH}/checkpoint-96"
model = AutoModelForSequenceClassification.from_pretrained(checkpoint).to(device)
from transformers import AutoTokenizer, DataCollatorWithPadding
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
- 测试数据集
import pandas as pd
test_df = pd.read_csv("./data/sougou/test.csv")
test_df.head()
| text | label | |
|---|---|---|
| 0 | 届数比赛时间比赛地点参加国家和地区冠军亚军决赛成绩第一届1956-1957英国11美国丹麦6... | 0 |
| 1 | 商品属性材质软橡胶带加浮雕工艺+合金彩色队徽吊牌规格162mm数量这一系列产品不限量发行图案... | 0 |
| 2 | 今天下午，沈阳金德和长春亚泰队将在五里河相遇。在这两支球队中沈阳籍球员居多，因此这场比赛实际... | 0 |
| 3 | 本报讯中国足协准备好了与特鲁西埃谈判的合同文本，也在北京给他预订好了房间，但特鲁西埃爽约了！... | 0 |
| 4 | 网友点击发表评论祝贺中国队夺得五连冠搜狐体育讯北京时间5月6日，2006年尤伯杯羽毛球赛在日... | 0 |
- 量化前模型的推理时间及评估指标
import numpy as np
import time
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=MAX_LENGTH, truncation=True, padding=True, return_tensors='pt').to(device)
# print(encoded_text)
logits = model(**encoded_text)
label_id = np.argmax(logits[0].detach().cpu().numpy(), axis=1)[0]
pred_labels.append(label_id)
print(i, (time.time() - row_s_time)*1000, label_id)
print("avg time: ", (time.time() - s_time) * 1000 / test_df.shape[0])
0 229.3872833251953 0
100 362.0314598083496 1
200 311.16747856140137 2
300 324.13792610168457 3
400 406.9099426269531 4
avg time: 352.44047810332944
from sklearn.metrics import classification_report
print(classification_report(true_labels, pred_labels, digits=4))
precision recall f1-score support
0 0.9900 1.0000 0.9950 99
1 0.9691 0.9495 0.9592 99
2 0.9900 1.0000 0.9950 99
3 0.9320 0.9697 0.9505 99
4 0.9895 0.9495 0.9691 99
accuracy 0.9737 495
macro avg 0.9741 0.9737 0.9737 495
weighted avg 0.9741 0.9737 0.9737 495
- 设置量化后端
# 模型量化
cpu_device = torch.device("cpu")
torch.backends.quantized.supported_engines
['none', 'onednn', 'x86', 'fbgemm']
torch.backends.quantized.engine = 'x86'
- 量化后模型的推理时间及评估指标
# 8-bit 量化
quantized_model = torch.quantization.quantize_dynamic(
model, {torch.nn.Linear}, dtype=torch.qint8
).to(cpu_device)
q_s_time = time.time()
q_true_labels, q_pred_labels = [], []
for i, row in test_df.iterrows():
row_s_time = time.time()
q_true_labels.append(row["label"])
encoded_text = tokenizer(row['text'], max_length=MAX_LENGTH, truncation=True, padding=True, return_tensors='pt').to(cpu_device)
logits = quantized_model(**encoded_text)
label_id = np.argmax(logits[0].detach().numpy(), axis=1)[0]
q_pred_labels.append(label_id)
print(i, (time.time() - row_s_time) * 1000, label_id)
print("avg time: ", (time.time() - q_s_time) * 1000 / test_df.shape[0])
0 195.47462463378906 0
100 247.33805656433105 1
200 219.41304206848145 2
300 206.44831657409668 3
400 187.4992847442627 4
avg time: 217.63229466447928
from sklearn.metrics import classification_report
print(classification_report(q_true_labels, q_pred_labels, digits=4))
precision recall f1-score support
0 0.9900 1.0000 0.9950 99
1 0.9688 0.9394 0.9538 99
2 0.9900 1.0000 0.9950 99
3 0.9320 0.9697 0.9505 99
4 0.9896 0.9596 0.9744 99
accuracy 0.9737 495
macro avg 0.9741 0.9737 0.9737 495
weighted avg 0.9741 0.9737 0.9737 495
- 量化前后模型大小对比
import os
def print_size_of_model(model):
torch.save(model.state_dict(), "temp.p")
print("Size (MB): ", os.path.getsize("temp.p")/1e6)
os.remove("temp.p")
print_size_of_model(model)
print_size_of_model(quantized_model)
Size (MB): 409.155273
Size (MB): 152.627621
量化后端（Quantization backend）取决于CPU架构，不同计算机的CPU架构不同，因此，默认的动态量化不一定在所有的CPU上都能生效，需根据自己计算机的CPU架构设置好对应的量化后端。另外，不同的量化后端也有些许差异。Linux服务器使用uname -a可查看CPU信息。
重复上述实验过程，以模型的最大输入长度为变量，取值为128,256,384，每种情况各做3次实验，结果如下：
| 实验 | 最大长度 | 量化前平均推理时间(ms) | 量化前weighted F1值 | 量化前平均推理时间(ms) | 量化前weighted F1值 |
|---|---|---|---|---|---|
| 实验1 | 384 | 1066 | 0.9797 | 686 | 0.9838 |
| 实验2 | 384 | 1047.6 | 0.9899 | 738.1 | 0.9879 |
| 实验3 | 384 | 1020.9 | 0.9817 | 714.0 | 0.9838 |
| 实验1 | 256 | 668.7 | 0.9717 | 431.4 | 0.9718 |
| 实验2 | 256 | 675.1 | 0.9717 | 449.9 | 0.9718 |
| 实验3 | 256 | 656.0 | 0.9717 | 446.5 | 0.9718 |
| 实验1 | 128 | 335.8 | 0.9737 | 200.5 | 0.9737 |
| 实验2 | 128 | 336.5 | 0.9737 | 227.2 | 0.9737 |
| 实验3 | 128 | 352.4 | 0.9737 | 217.6 | 0.9737 |
综上所述，对于训练后的BERT模型（文本分类模型）进行动态量化，其结论如下：
- 模型推理效果：量化前后基本相同，量化后略有下降
- 模型推理时间：量化后平均提速约1.52倍
总结
本文介绍了量化基本概念，PyTorch模型量化方式，以及对BERT模型训练后进行动态量化后在推理效果和推理性能上的实验。
本文项目已开源至Github项目：https://github.com/percent4/dynamic_quantization_on_bert 。
本人已开通个人博客网站，网址为：https://percent4.github.io/ ，欢迎大家访问~
