本文将会介绍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：

1. 集合：比如： qconfig_spec={nn.LSTM, nn.Linear} 。列出要量化的神经网络模块。
2. 字典： 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]]]])

模型量化策略

  当前，由于量化算子的覆盖有限，因此，对于不同的深度学习模型，其量化策略不同，见下表：

   下面对BERT模型进行训练后动态量化，分析模型在量化前后，推理效果和推理性能的变化。

实验

   我们使用的训练后的模型为中文文本分类模型，其训练过程可以参考文章：NLP（六十六）使用HuggingFace中的Trainer进行BERT模型微调 。
   训练后的BERT模型动态量化实验的设置如下：

1. base model: bert-base-chinese
2. CPU info: x86-64, Intel® Core™ i5-10210U CPU @ 1.60GHz
3. batch size: 1
4. 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()

• 量化前模型的推理时间及评估指标

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次实验，结果如下：

  综上所述，对于训练后的BERT模型（文本分类模型）进行动态量化，其结论如下：

• 模型推理效果：量化前后基本相同，量化后略有下降
• 模型推理时间：量化后平均提速约1.52倍

总结

  本文介绍了量化基本概念，PyTorch模型量化方式，以及对BERT模型训练后进行动态量化后在推理效果和推理性能上的实验。
  本文项目已开源至Github项目：https://github.com/percent4/dynamic_quantization_on_bert 。
  本人已开通个人博客网站，网址为：https://percent4.github.io/ ，欢迎大家访问~