本文分享自华为云社区《TextBrewer:融合并改进了NLP和CV中的多种知识蒸馏技术、提供便捷快速的知识蒸馏框架、提升模型的推理速度,减少内存占用》,作者:汀丶。
TextBrewer是一个基于PyTorch的、为实现NLP中的知识蒸馏任务而设计的工具包,
融合并改进了NLP和CV中的多种知识蒸馏技术,提供便捷快速的知识蒸馏框架,用于以较低的性能损失压缩神经网络模型的大小,提升模型的推理速度,减少内存占用。
TextBrewer 为NLP中的知识蒸馏任务设计,融合了多种知识蒸馏技术,提供方便快捷的知识蒸馏框架。
主要特点:
TextBrewer目前支持的知识蒸馏技术有:
TextBrewer的主要功能与模块分为3块:
用户需要准备:
在多个典型NLP任务上,TextBrewer都能取得较好的压缩效果。相关实验见蒸馏效果。
Python >= 3.6
PyTorch >= 1.1.0
TensorboardX or Tensorboard
NumPy
tqdm
Transformers >= 2.0 (可选, Transformer相关示例需要用到)
Apex == 0.1.0 (可选,用于混合精度训练)
从PyPI自动下载安装包安装:
pip install textbrewer
git clone https://github.com/airaria/TextBrewer.git pip install ./textbrewer
Stage 1 : 蒸馏之前的准备工作:
Stage 2 : 使用TextBrewer蒸馏:
TrainingConfig
)和蒸馏配置(DistillationConfig
),初始化distiller在开始蒸馏之前准备:
teacher_model
(BERT-base),待训练学生模型student_model
(3-layer BERT)dataloader
,优化器optimizer
,学习率调节器类或者构造函数scheduler_class
和构造用的参数字典 scheduler_args
使用TextBrewer蒸馏:
import textbrewer from textbrewer import GeneralDistiller from textbrewer import TrainingConfig, DistillationConfig #展示模型参数量的统计 print("\nteacher_model's parametrers:") result, _ = textbrewer.utils.display_parameters(teacher_model,max_level=3) print (result) print("student_model's parametrers:") result, _ = textbrewer.utils.display_parameters(student_model,max_level=3) print (result) #定义adaptor用于解释模型的输出 def simple_adaptor(batch, model_outputs): # model输出的第二、三个元素分别是logits和hidden states return {'logits': model_outputs[1], 'hidden': model_outputs[2]} #蒸馏与训练配置 # 匹配教师和学生的embedding层;同时匹配教师的第8层和学生的第2层 distill_config = DistillationConfig( intermediate_matches=[ {'layer_T':0, 'layer_S':0, 'feature':'hidden', 'loss': 'hidden_mse','weight' : 1}, {'layer_T':8, 'layer_S':2, 'feature':'hidden', 'loss': 'hidden_mse','weight' : 1}]) train_config = TrainingConfig() #初始化distiller distiller = GeneralDistiller( train_config=train_config, distill_config = distill_config, model_T = teacher_model, model_S = student_model, adaptor_T = simple_adaptor, adaptor_S = simple_adaptor) #开始蒸馏 with distiller: distiller.train(optimizer, dataloader, num_epochs=1, scheduler_class=scheduler_class, scheduler_args = scheduler_args, callback=None)
Transformers 4示例
examples/random_token_example: 一个可运行的简单示例,在文本分类任务上以随机文本为输入,演示TextBrewer用法。
examples/cmrc2018_example (中文): CMRC 2018上的中文阅读理解任务蒸馏,并使用DRCD数据集做数据增强。
examples/mnli_example (英文): MNLI任务上的英文句对分类任务蒸馏,并展示如何使用多教师蒸馏。
examples/conll2003_example (英文): CoNLL-2003英文实体识别任务上的序列标注任务蒸馏。
examples/msra_ner_example (中文): MSRA NER(中文命名实体识别)任务上,使用分布式数据并行训练的Chinese-ELECTRA-base模型蒸馏。
我们在多个中英文文本分类、阅读理解、序列标注数据集上进行了蒸馏实验。实验的配置和效果如下。
模型
对于英文任务,教师模型为BERT-base-cased
对于中文任务,教师模型为HFL发布的RoBERTa-wwm-ext 与 Electra-base
我们测试了不同的学生模型,为了与已有公开结果相比较,除了BiGRU都是和BERT一样的多层Transformer结构。模型的参数如下表所示。需要注意的是,参数量的统计包括了embedding层,但不包括最终适配各个任务的输出层。
distill_config = DistillationConfig(temperature = 8, intermediate_matches = matches) #其他参数为默认值
不同的模型用的matches
我们采用了以下配置:
各种matches的定义在examples/matches/matches.py中。均使用GeneralDistiller进行蒸馏。
蒸馏用的学习率 lr=1e-4(除非特殊说明)。训练30~60轮。
在英文实验中,我们使用了如下三个典型数据集。
我们在下面两表中列出了DistilBERT, BERT-PKD, BERT-of-Theseus, TinyBERT 等公开的蒸馏结果,并与我们的结果做对比。
Public results:
Our results:
说明:
在中文实验中,我们使用了如下典型数据集。
实验结果如下表所示。
说明:
TrainingConfig
和 DistillationConfig
:训练和蒸馏相关的配置。Distiller负责执行实际的蒸馏过程。目前实现了以下的distillers:
BasicDistiller
: 提供单模型单任务蒸馏方式。可用作测试或简单实验。GeneralDistiller
(常用): 提供单模型单任务蒸馏方式,并且支持中间层特征匹配,一般情况下推荐使用。MultiTeacherDistiller
: 多教师蒸馏。将多个(同任务)教师模型蒸馏到一个学生模型上。暂不支持中间层特征匹配。MultiTaskDistiller
:多任务蒸馏。将多个(不同任务)单任务教师模型蒸馏到一个多任务学生模型。BasicTrainer
:用于单个模型的有监督训练,而非蒸馏。可用于训练教师模型。蒸馏实验中,有两个组件需要由用户提供,分别是callback 和 adaptor :
回调函数。在每个checkpoint,保存模型后会被distiller
调用,并传入当前模型。可以借由回调函数在每个checkpoint评测模型效果。
将模型的输入和输出转换为指定的格式,向distiller
解释模型的输入和输出,以便distiller
根据不同的策略进行不同的计算。在每个训练步,batch
和模型的输出model_outputs
会作为参数传递给adaptor
,adaptor
负责重新组织这些数据,返回一个字典。
更多细节可参见完整文档中的说明。
Q: 学生模型该如何初始化?
A: 知识蒸馏本质上是“老师教学生”的过程。在初始化学生模型时,可以采用随机初始化的形式(即完全不包含任何先验知识),也可以载入已训练好的模型权重。例如,从BERT-base模型蒸馏到3层BERT时,可以预先载入RBT3模型权重(中文任务)或BERT的前三层权重(英文任务),然后进一步进行蒸馏,避免了蒸馏过程的“冷启动”问题。我们建议用户在使用时尽量采用已预训练过的学生模型,以充分利用大规模数据预训练所带来的优势。
Q: 如何设置蒸馏的训练参数以达到一个较好的效果?
A: 知识蒸馏的比有标签数据上的训练需要更多的训练轮数与更大的学习率。比如,BERT-base上训练SQuAD一般以lr=3e-5训练3轮左右即可达到较好的效果;而蒸馏时需要以lr=1e-4训练30~50轮。当然具体到各个任务上肯定还有区别,我们的建议仅是基于我们的经验得出的,仅供参考。
Q: 我的教师模型和学生模型的输入不同(比如词表不同导致input_ids不兼容),该如何进行蒸馏?
A: 需要分别为教师模型和学生模型提供不同的batch,参见完整文档中的 Feed Different batches to Student and Teacher, Feed Cached Values 章节。
Q: 我缓存了教师模型的输出,它们可以用于加速蒸馏吗?
A: 可以, 参见完整文档中的 Feed Different batches to Student and Teacher, Feed Cached Values 章节。
|