TinyBERT（官网介绍）

https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/TinyBERT
TinyBERT比BERT-base小7.5倍，推理速度快9.4倍，在自然语言理解任务中表现出色。它在训练前和任务特定的学习阶段执行一种新的transformer蒸馏。TinyBERT学习概述如下:

image.png
细节看论文：TinyBERT: Distilling BERT for Natural Language Understanding

安装依赖

python版本：python3

  1. pip install -r requirements.txt

一般蒸馏方法：

一般来说，我们使用未经微调的原始bert模型作为教师，使用大型文本语料库作为学习数据。通过对一般域的文本进行变压器蒸馏，我们得到了一个一般的TinyBERT，它为特定任务的蒸馏提供了一个良好的初始化。
一般蒸馏方法分为两个步骤:(1)生成json格式的语料库;(2)运行变压器蒸馏;
步骤1:使用pregenerate_training_data.py。生成json格式的语料

  1. # ${BERT_BASE_DIR}$ includes the BERT-base teacher model.
  2. python pregenerate_training_data.py --train_corpus ${CORPUS_RAW} \ 
  3.                   --bert_model ${BERT_BASE_DIR}$ \
  4.                   --reduce_memory --do_lower_case \
  5.                   --epochs_to_generate 3 \
  6.                   --output_dir ${CORPUS_JSON_DIR}$

第二步:使用general-distill.py来进行一般的蒸馏

  1. # ${STUDENT_CONFIG_DIR}$ includes the config file of student_model.
  3. python general_distill.py --pregenerated_data ${CORPUS_JSON}$ \ 
  4.                           --teacher_model ${BERT_BASE}$ \
  5.                           --student_model ${STUDENT_CONFIG_DIR}$ \
  6.                           --reduce_memory --do_lower_case \
  7.                           --train_batch_size 256 \
  8.                           --output_dir ${GENERAL_TINYBERT_DIR}$

我们还提供了一般的TinyBERT模型，用户可以跳过一般的蒸馏。
=================第一个版本, 我们在论文中所使用的结果 ===========================
General_TinyBERT(4layer-312dim)
General_TinyBERT(6layer-768dim)
=================第二版(2019/11/18)使用更多(book+wiki)和没有“[MASK]”语料库进行训练 =======
General_TinyBERT_v2(4layer-312dim)
General_TinyBERT_v2(6layer-768dim)

数据扩张

数据扩充的目的是扩展任务特定的训练集。学习更多的任务相关的例子，可以进一步提高学生模型的泛化能力。我们结合一个预先训练好的语言模型BERT和手套嵌入来做单词级的替换来增加数据。
使用data_augmentation.py 运行数据扩充和扩充数据集,结果会自动保存到${GLUE_DIR/TASK_NAME}$/train_aug.tsv ：

  1. python data_augmentation.py --pretrained_bert_model ${BERT_BASE_DIR}$ \
  2.                             --glove_embs ${GLOVE_EMB}$ \
  3.                             --glue_dir ${GLUE_DIR}$ \  
  4.                             --task_name ${TASK_NAME}$

在运行GLUE任务的数据扩充之前，您应该通过运行这个脚本下载GLUE数据，并将其解压缩到GLUE_DIR目录。TASK_NAME可以是CoLA，SST-2, MRPC, STS-B, QQP, MNLI, QNLI, RTE的任意一个。

特定于任务的蒸馏

在任务特定精馏中，我们重新执行提出的变压器精馏，通过重点学习任务特定的知识来进一步改进TinyBERT。
具体任务精馏包括两个步骤:(1)中间层精馏;(2)预测层蒸馏。

  • 第一步:使用’task_distill .py。运行中间层蒸馏。
  1. # ${FT_BERT_BASE_DIR}$ contains the fine-tuned BERT-base model.
  2. python task_distill.py --teacher_model ${FT_BERT_BASE_DIR}$ \
  3.                        --student_model ${GENERAL_TINYBERT_DIR}$ \
  4.                        --data_dir ${TASK_DIR}$ \
  5.                        --task_name ${TASK_NAME}$ \ 
  6.                        --output_dir ${TMP_TINYBERT_DIR}$ \
  7.                        --max_seq_length 128 \
  8.                        --train_batch_size 32 \
  9.                        --num_train_epochs 10 \
  10.                        --aug_train \
  11.                        --do_lower_case
  • 第二步:使用task_distill.py运行预测层蒸馏。
  1. python task_distill.py --pred_distill  \
  2.                        --teacher_model ${FT_BERT_BASE_DIR}$ \
  3.                        --student_model ${TMP_TINYBERT_DIR}$ \
  4.                        --data_dir ${TASK_DIR}$ \
  5.                        --task_name ${TASK_NAME}$ \
  6.                        --output_dir ${TINYBERT_DIR}$ \
  7.                        --aug_train  \  
  8.                        --do_lower_case \
  9.                        --learning_rate 3e-5  \
  10.                        --num_train_epochs  3  \
  11.                        --eval_step 100 \
  12.                        --max_seq_length 128 \
  13.                        --train_batch_size 32

我们还提供了所有GLUE任务的蒸馏TinyBERT(4layer-312dim和6layer-768dim)供评估。每个任务都有自己的文件夹，其中保存了相应的模型。
TinyBERT(4layer-312dim)
TinyBERT(6layer-768dim)

评估

task_distill.py 也提供了评估运行以下命令:

  1. ${TINYBERT_DIR}$ includes the config file, student model and vocab file.
  2. python task_distill.py --do_eval \
  3.                        --student_model ${TINYBERT_DIR}$ \
  4.                        --data_dir ${TASK_DIR}$ \
  5.                        --task_name ${TASK_NAME}$ \
  6.                        --output_dir ${OUTPUT_DIR}$ \
  7.                        --do_lower_case \
  8.                        --eval_batch_size 32 \
  9.                        --max_seq_length 128

改进

根据transformers修改了数据预处理流程，增加了保存特征的代码，使得程序运行更加顺畅