原文：PyTorch (實驗）BERT 上的動態(tài)量化

作者：黃建宇

被審核： Raghuraman Krishnamoorthi

由編輯：林 ess 琳

介紹

在本教程中，我們將動態(tài)量化應(yīng)用在 BERT 模型上，緊跟 HuggingFace Transformers 示例中的 BERT 模型。 通過這一循序漸進(jìn)的過程，我們將演示如何將 BERT 等眾所周知的最新模型轉(zhuǎn)換為動態(tài)量化模型。

• BERT，或者說 Transformers 的雙向嵌入表示法，是一種預(yù)訓(xùn)練語言表示法的新方法，可以在許多流行的自然語言處理(NLP）任務(wù)(例如問題解答，文本分類， 和別的。 可以在此處找到。
• PyTorch 中的動態(tài)量化支持將權(quán)重模型的浮點模型轉(zhuǎn)換為具有靜態(tài) int8 或 float16 數(shù)據(jù)類型的量化模型，并為激活提供動態(tài)量化。 當(dāng)權(quán)重量化為 int8 時，激活(每批）動態(tài)量化為 int8。 在 PyTorch 中，我們有 torch.quantization.quantize_dynamic API ，它用僅動態(tài)權(quán)重的量化版本替換了指定的模塊，并輸出了量化模型。
• 我們在通用語言理解評估基準(zhǔn)(GLUE）中演示了 Microsoft Research Paraphrase 語料庫(MRPC）任務(wù)的準(zhǔn)確性和推理性能結(jié)果。 MRPC(Dolan 和 Brockett，2005 年）是從在線新聞源中自動提取的句子對的語料庫，帶有人工注釋，說明句子中的句子在語義上是否等效。 由于班級不平衡(正向為 68％，負(fù)向為 32％），我們遵循常規(guī)做法并報告 F1 得分。 MRPC 是用于語言對分類的常見 NLP 任務(wù)，如下所示。

1.設(shè)定

1.1 安裝 PyTorch 和 HuggingFace 變壓器

要開始本教程，首先請遵循 PyTorch (此處）和 HuggingFace Github Repo (此處）中的安裝說明。 此外，我們還將安裝 scikit-learn 軟件包，因為我們將重復(fù)使用其內(nèi)置的 F1 分?jǐn)?shù)計算幫助器功能。

pip install sklearn
pip install transformers

由于我們將使用 PyTorch 的實驗部分，因此建議安裝最新版本的 Torch 和 Torchvision。 您可以在此處找到有關(guān)本地安裝的最新說明。 例如，要在 Mac 上安裝：

yes y | pip uninstall torch tochvision
yes y | pip install --pre torch -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html

1.2 導(dǎo)入必要的模塊

在這一步中，我們將導(dǎo)入本教程所需的 Python 模塊。

from __future__ import absolute_import, division, print_function


import logging
import numpy as np
import os
import random
import sys
import time
import torch


from argparse import Namespace
from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,
                              TensorDataset)
from tqdm import tqdm
from transformers import (BertConfig, BertForSequenceClassification, BertTokenizer,)
from transformers import glue_compute_metrics as compute_metrics
from transformers import glue_output_modes as output_modes
from transformers import glue_processors as processors
from transformers import glue_convert_examples_to_features as convert_examples_to_features


## Setup logging
logger = logging.getLogger(__name__)
logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt = '%m/%d/%Y %H:%M:%S',
                    level = logging.WARN)


logging.getLogger("transformers.modeling_utils").setLevel(
   logging.WARN)  # Reduce logging


print(torch.__version__)

我們設(shè)置線程數(shù)以比較 FP32 和 INT8 性能之間的單線程性能。 在本教程的最后，用戶可以通過使用右側(cè)并行后端構(gòu)建 PyTorch 來設(shè)置其他線程數(shù)量。

torch.set_num_threads(1)
print(torch.__config__.parallel_info())

1.3 了解助手功能

幫助器功能內(nèi)置在轉(zhuǎn)換器庫中。 我們主要使用以下輔助函數(shù)：一個用于將文本示例轉(zhuǎn)換為特征向量的函數(shù)； 另一個用于測量預(yù)測結(jié)果的 F1 分?jǐn)?shù)。

gum_convert_examples_to_features 函數(shù)將文本轉(zhuǎn)換為輸入特征：

• 標(biāo)記輸入序列；
• 在開頭插入[CLS]；
• 在第一句和第二句之間并在最后插入[SEP]；
• 生成令牌類型 ID，以指示令牌是屬于第一序列還是第二序列。

gum_compute_metrics 函數(shù)的計算指標(biāo)為 F1 得分，可以將其解釋為精度和召回率的加權(quán)平均值，其中 F1 得分在 1 和最差處達(dá)到最佳值 得分為 0。精度和召回率對 F1 得分的相對貢獻(xiàn)相等。

• F1 分?jǐn)?shù)的公式為：

1.4 下載數(shù)據(jù)集

在運行 MRPC 任務(wù)之前，我們通過運行腳本并下載 GLUE 數(shù)據(jù)并將其解壓縮到目錄glue_data中。

python download_glue_data.py --data_dir='glue_data' --tasks='MRPC'

2.微調(diào) BERT 模型

BERT 的精神是預(yù)訓(xùn)練語言表示形式，然后以最小的任務(wù)相關(guān)參數(shù)對各種任務(wù)上的深層雙向表示形式進(jìn)行微調(diào)，并獲得最新的結(jié)果。 在本教程中，我們將專注于對預(yù)訓(xùn)練的 BERT 模型進(jìn)行微調(diào)，以對 MRPC 任務(wù)上的語義等效句子對進(jìn)行分類。

要為 MRPC 任務(wù)微調(diào)預(yù)訓(xùn)練的 BERT 模型(HuggingFace 變壓器中的bert-base-uncased模型），可以按照示例中的命令進(jìn)行操作：

export GLUE_DIR=./glue_data
export TASK_NAME=MRPC
export OUT_DIR=./$TASK_NAME/
python ./run_glue.py \
    --model_type bert \
    --model_name_or_path bert-base-uncased \
    --task_name $TASK_NAME \
    --do_train \
    --do_eval \
    --do_lower_case \
    --data_dir $GLUE_DIR/$TASK_NAME \
    --max_seq_length 128 \
    --per_gpu_eval_batch_size=8   \
    --per_gpu_train_batch_size=8   \
    --learning_rate 2e-5 \
    --num_train_epochs 3.0 \
    --save_steps 100000 \
    --output_dir $OUT_DIR

我們在此處為 MRPC 任務(wù)提供了經(jīng)過微調(diào)的 BERT 模型。 為了節(jié)省時間，您可以將模型文件(?400 MB）直接下載到本地文件夾$OUT_DIR中。

2.1 設(shè)置全局配置

在這里，我們設(shè)置了全局配置，用于評估動態(tài)量化前后的微調(diào) BERT 模型。

configs = Namespace()


## The output directory for the fine-tuned model, $OUT_DIR.
configs.output_dir = "./MRPC/"


## The data directory for the MRPC task in the GLUE benchmark, $GLUE_DIR/$TASK_NAME.
configs.data_dir = "./glue_data/MRPC"


## The model name or path for the pre-trained model.
configs.model_name_or_path = "bert-base-uncased"
## The maximum length of an input sequence
configs.max_seq_length = 128


## Prepare GLUE task.
configs.task_name = "MRPC".lower()
configs.processor = processors[configs.task_name]()
configs.output_mode = output_modes[configs.task_name]
configs.label_list = configs.processor.get_labels()
configs.model_type = "bert".lower()
configs.do_lower_case = True


## Set the device, batch size, topology, and caching flags.
configs.device = "cpu"
configs.per_gpu_eval_batch_size = 8
configs.n_gpu = 0
configs.local_rank = -1
configs.overwrite_cache = False


## Set random seed for reproducibility.
def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
set_seed(42)

2.2 加載微調(diào)的 BERT 模型

我們從configs.output_dir加載標(biāo)記器和經(jīng)過微調(diào)的 BERT 序列分類器模型(FP32）。

tokenizer = BertTokenizer.from_pretrained(
    configs.output_dir, do_lower_case=configs.do_lower_case)


model = BertForSequenceClassification.from_pretrained(configs.output_dir)
model.to(configs.device)

2.3 定義標(biāo)記化和評估功能

我們重用了 Huggingface 中的標(biāo)記化和評估函數(shù)。

## coding=utf-8
## Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
## Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
## ## Licensed under the Apache License, Version 2.0 (the "License");
## you may not use this file except in compliance with the License.
## You may obtain a copy of the License at
## ##     http://www.apache.org/licenses/LICENSE-2.0
## ## Unless required by applicable law or agreed to in writing, software
## distributed under the License is distributed on an "AS IS" BASIS,
## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
## See the License for the specific language governing permissions and
## limitations under the License.


def evaluate(args, model, tokenizer, prefix=""):
    # Loop to handle MNLI double evaluation (matched, mis-matched)
    eval_task_names = ("mnli", "mnli-mm") if args.task_name == "mnli" else (args.task_name,)
    eval_outputs_dirs = (args.output_dir, args.output_dir + '-MM') if args.task_name == "mnli" else (args.output_dir,)


    results = {}
    for eval_task, eval_output_dir in zip(eval_task_names, eval_outputs_dirs):
        eval_dataset = load_and_cache_examples(args, eval_task, tokenizer, evaluate=True)


        if not os.path.exists(eval_output_dir) and args.local_rank in [-1, 0]:
            os.makedirs(eval_output_dir)


        args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
        # Note that DistributedSampler samples randomly
        eval_sampler = SequentialSampler(eval_dataset) if args.local_rank == -1 else DistributedSampler(eval_dataset)
        eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)


        # multi-gpu eval
        if args.n_gpu > 1:
            model = torch.nn.DataParallel(model)


        # Eval!
        logger.info("***** Running evaluation {} *****".format(prefix))
        logger.info("  Num examples = %d", len(eval_dataset))
        logger.info("  Batch size = %d", args.eval_batch_size)
        eval_loss = 0.0
        nb_eval_steps = 0
        preds = None
        out_label_ids = None
        for batch in tqdm(eval_dataloader, desc="Evaluating"):
            model.eval()
            batch = tuple(t.to(args.device) for t in batch)


            with torch.no_grad():
                inputs = {'input_ids':      batch[0],
                          'attention_mask': batch[1],
                          'labels':         batch[3]}
                if args.model_type != 'distilbert':
                    inputs['token_type_ids'] = batch[2] if args.model_type in ['bert', 'xlnet'] else None  # XLM, DistilBERT and RoBERTa don't use segment_ids
                outputs = model(**inputs)
                tmp_eval_loss, logits = outputs[:2]


                eval_loss += tmp_eval_loss.mean().item()
            nb_eval_steps += 1
            if preds is None:
                preds = logits.detach().cpu().numpy()
                out_label_ids = inputs['labels'].detach().cpu().numpy()
            else:
                preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
                out_label_ids = np.append(out_label_ids, inputs['labels'].detach().cpu().numpy(), axis=0)


        eval_loss = eval_loss / nb_eval_steps
        if args.output_mode == "classification":
            preds = np.argmax(preds, axis=1)
        elif args.output_mode == "regression":
            preds = np.squeeze(preds)
        result = compute_metrics(eval_task, preds, out_label_ids)
        results.update(result)


        output_eval_file = os.path.join(eval_output_dir, prefix, "eval_results.txt")
        with open(output_eval_file, "w") as writer:
            logger.info("***** Eval results {} *****".format(prefix))
            for key in sorted(result.keys()):
                logger.info("  %s = %s", key, str(result[key]))
                writer.write("%s = %s\n" % (key, str(result[key])))


    return results


def load_and_cache_examples(args, task, tokenizer, evaluate=False):
    if args.local_rank not in [-1, 0] and not evaluate:
        torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache


    processor = processors[task]()
    output_mode = output_modes[task]
    # Load data features from cache or dataset file
    cached_features_file = os.path.join(args.data_dir, 'cached_{}_{}_{}_{}'.format(
        'dev' if evaluate else 'train',
        list(filter(None, args.model_name_or_path.split('/'))).pop(),
        str(args.max_seq_length),
        str(task)))
    if os.path.exists(cached_features_file) and not args.overwrite_cache:
        logger.info("Loading features from cached file %s", cached_features_file)
        features = torch.load(cached_features_file)
    else:
        logger.info("Creating features from dataset file at %s", args.data_dir)
        label_list = processor.get_labels()
        if task in ['mnli', 'mnli-mm'] and args.model_type in ['roberta']:
            # HACK(label indices are swapped in RoBERTa pretrained model)
            label_list[1], label_list[2] = label_list[2], label_list[1]
        examples = processor.get_dev_examples(args.data_dir) if evaluate else processor.get_train_examples(args.data_dir)
        features = convert_examples_to_features(examples,
                                                tokenizer,
                                                label_list=label_list,
                                                max_length=args.max_seq_length,
                                                output_mode=output_mode,
                                                pad_on_left=bool(args.model_type in ['xlnet']),                 # pad on the left for xlnet
                                                pad_token=tokenizer.convert_tokens_to_ids([tokenizer.pad_token])[0],
                                                pad_token_segment_id=4 if args.model_type in ['xlnet'] else 0,
        )
        if args.local_rank in [-1, 0]:
            logger.info("Saving features into cached file %s", cached_features_file)
            torch.save(features, cached_features_file)


    if args.local_rank == 0 and not evaluate:
        torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache


    # Convert to Tensors and build dataset
    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
    all_attention_mask = torch.tensor([f.attention_mask for f in features], dtype=torch.long)
    all_token_type_ids = torch.tensor([f.token_type_ids for f in features], dtype=torch.long)
    if output_mode == "classification":
        all_labels = torch.tensor([f.label for f in features], dtype=torch.long)
    elif output_mode == "regression":
        all_labels = torch.tensor([f.label for f in features], dtype=torch.float)


    dataset = TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_labels)
    return dataset

3.應(yīng)用動態(tài)量化

我們在模型上調(diào)用torch.quantization.quantize_dynamic，將動態(tài)量化應(yīng)用于 HuggingFace BERT 模型。 特別，

• 我們指定要對模型中的 torch.nn.Linear 模塊進(jìn)行量化；
• 我們指定希望將權(quán)重轉(zhuǎn)換為量化的 int8 值。

quantized_model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)
print(quantized_model)

3.1 檢查型號

我們首先檢查一下模型尺寸。 我們可以看到模型大小顯著減少(FP32 總大?。?38 MB； INT8 總大?。?81 MB）：

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)

本教程中使用的 BERT 模型(bert-base-uncased）的詞匯量 V 為 30522。在嵌入量為 768 的情況下，單詞嵌入表的總大小為?4(字節(jié)/ FP32） 30522 768 = 90 MB 。 因此，借助量化，非嵌入表部分的模型大小從 350 MB(FP32 模型）減少到 90 MB(INT8 模型）。

3.2 評估推理的準(zhǔn)確性和時間

接下來，我們比較一下動態(tài)量化后原始 FP32 模型和 INT8 模型之間的推斷時間以及評估精度。

def time_model_evaluation(model, configs, tokenizer):
    eval_start_time = time.time()
    result = evaluate(configs, model, tokenizer, prefix="")
    eval_end_time = time.time()
    eval_duration_time = eval_end_time - eval_start_time
    print(result)
    print("Evaluate total time (seconds): {0:.1f}".format(eval_duration_time))


## Evaluate the original FP32 BERT model
time_model_evaluation(model, configs, tokenizer)


## Evaluate the INT8 BERT model after the dynamic quantization
time_model_evaluation(quantized_model, configs, tokenizer)

在 MacBook Pro 上本地運行此程序，無需進(jìn)行量化，推理(對于 MRPC 數(shù)據(jù)集中的所有 408 個示例）大約需要 160 秒，而進(jìn)行量化則只需大約 90 秒。 我們總結(jié)了在 Macbook Pro 上運行量化 BERT 模型推斷的結(jié)果，如下所示：

| Prec | F1 score | Model Size | 1 thread | 4 threads |
| FP32 |  0.9019  |   438 MB   | 160 sec  | 85 sec    |
| INT8 |  0.8953  |   181 MB   |  90 sec  | 46 sec    |

在 MRPC 任務(wù)的微調(diào) BERT 模型上應(yīng)用訓(xùn)練后動態(tài)量化后，我們的 F1 分?jǐn)?shù)準(zhǔn)確性為 0.6％。 作為比較，在的最新論文(表 1）中，通過應(yīng)用訓(xùn)練后動態(tài)量化，可以達(dá)到 0.8788；通過應(yīng)用量化感知訓(xùn)練，可以達(dá)到 0.8956。 主要區(qū)別在于我們在 PyTorch 中支持非對稱量化，而該論文僅支持對稱量化。

請注意，在本教程中，為了進(jìn)行單線程比較，我們將線程數(shù)設(shè)置為 1。 對于這些量化的 INT8 運算符，我們還支持運算內(nèi)并行化。 用戶現(xiàn)在可以通過torch.set_num_threads(N)設(shè)置多線程(N是內(nèi)部運算并行線程的數(shù)量）。 啟用幀內(nèi)并行支持的一項初步要求是使用正確的后端(例如 OpenMP，Native 或 TBB）構(gòu)建 PyTorch。 您可以使用torch.__config__.parallel_info()檢查并行化設(shè)置。 在使用 PyTorch 和本機后端進(jìn)行并行化的同一臺 MacBook Pro 上，我們可以獲得大約 46 秒的時間來處理 MRPC 數(shù)據(jù)集的評估。

3.3 序列化量化模型

我們可以序列化并保存量化模型，以備將來使用。

quantized_output_dir = configs.output_dir + "quantized/"
if not os.path.exists(quantized_output_dir):
    os.makedirs(quantized_output_dir)
    quantized_model.save_pretrained(quantized_output_dir)

結(jié)論

在本教程中，我們演示了如何演示如何將 BERT 等著名的最新 NLP 模型轉(zhuǎn)換為動態(tài)量化模型。 動態(tài)量化可以減小模型的大小，而對準(zhǔn)確性的影響有限。

謝謝閱讀！ 與往常一樣，我們歡迎您提供任何反饋，因此，如果有任何問題，請在此處創(chuàng)建一個問題。

參考文獻(xiàn)

[1] J.Devlin，M。Chang，K。Lee 和 K. Toutanova， BERT：用于語言理解的深度雙向變壓器的預(yù)訓(xùn)練(2018）。

[2] HuggingFace 變壓器。

[3] O. Zafrir，G。Boudoukh，P。Izsak 和 M. Wasserblat(2019 年）。 Q8BERT：量化的 8 位 BERT 。