本文共享自华为云社区《TextBrewer:融兼并改进了NLP和CV中的多种常识蒸馏技能、供给快捷快速的常识蒸馏结构、提升模型的推理速度,减少内存占用》,作者:汀丶。

TextBrewer是一个根据PyTorch的、为完成NLP中的常识蒸馏使命而规划的工具包,
融兼并改进了NLP和CV中的多种常识蒸馏技能,供给快捷快速的常识蒸馏结构,用于以较低的功能丢失紧缩神经网络模型的巨细,提升模型的推理速度,减少内存占用。

1.简介

TextBrewer 为NLP中的常识蒸馏使命规划,交融了多种常识蒸馏技能,供给便利快捷的常识蒸馏结构。

主要特点:

  • 模型无关:适用于多种模型结构(主要面向Transfomer结构)
  • 便利灵活:可自由组合多种蒸馏办法;可便利添加自界说丢失等模块
  • 非侵入式:无需对教师与学生模型自身结构进行修正
  • 支撑典型的NLP使命:文本分类、阅览理解、序列标示等

TextBrewer目前支撑的常识蒸馏技能有:

  • 软标签与硬标签混合练习
  • 动态丢失权重调整与蒸馏温度调整
  • 多种蒸馏丢失函数: hidden states MSE, attention-based loss, neuron selectivity transfer, …
  • 任意构建中间层特征匹配计划
  • 多教师常识蒸馏

TextBrewer的主要功能与模块分为3块:

  1. Distillers:进行蒸馏的中心部件,不同的distiller供给不同的蒸馏模式。目前包含GeneralDistiller, MultiTeacherDistiller, MultiTaskDistiller等
  2. Configurations and Presets:练习与蒸馏办法的装备,并供给预界说的蒸馏策略以及多种常识蒸馏丢失函数
  3. Utilities:模型参数剖析显示等辅助工具

用户需求准备:

  1. 已练习好的教师模型, 待蒸馏的学生模型
  2. 练习数据与必要的试验装备, 即可开端蒸馏

在多个典型NLP使命上,TextBrewer都能取得较好的紧缩作用。相关试验见蒸馏作用。

2.TextBrewer结构

一文详解TextBrewer

2.1 装置要求

  • 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 github.com/airaria/Tex… pip install ./textbrewer

2.2作业流程

一文详解TextBrewer

一文详解TextBrewer

  • Stage 1 : 蒸馏之前的准备作业:

    1. 练习教师模型
    2. 界说与初始化学生模型(随机初始化,或载入预练习权重)
    3. 结构蒸馏用数据集的dataloader,练习学生模型用的optimizer和learning rate scheduler
  • Stage 2 : 运用TextBrewer蒸馏:

    1. 结构练习装备(TrainingConfig)和蒸馏装备(DistillationConfig),初始化distiller
    2. 界说adaptorcallback ,分别用于适配模型输入输出和练习进程中的回调
    3. 调用distillertrain办法开端蒸馏

2.3 以蒸馏BERT-base到3层BERT为例展现TextBrewer用法

在开端蒸馏之前准备:

  • 练习好的教师模型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)

2.4蒸馏使命示例

  • Transformers 4示例

    • examples/notebook_examples/sst2.ipynb (英文): SST-2文本分类使命上的BERT模型练习与蒸馏。
    • examples/notebook_examples/msra_ner.ipynb (中文): MSRA NER中文命名实体辨认使命上的BERT模型练习与蒸馏。
    • examples/notebook_examples/sqaudv1.1.ipynb (英文): SQuAD 1.1英文阅览理解使命上的BERT模型练习与蒸馏。
  • 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模型蒸馏。

2.4.1蒸馏作用

咱们在多个中英文文本分类、阅览理解、序列标示数据集上进行了蒸馏试验。试验的装备和作用如下。

  • 模型

  • 对于英文使命,教师模型为BERT-base-cased

  • 对于中文使命,教师模型为HFL发布的RoBERTa-wwm-extElectra-base

咱们测验了不同的学生模型,为了与已有公开成果相比较,除了BiGRU都是和BERT一样的多层Transformer结构。模型的参数如下表所示。需求注意的是,参数量的核算包括了embedding层,但不包括终究适配各个使命的输出层。

  • 英文模型

一文详解TextBrewer

  • 中文模型

一文详解TextBrewer

  • T6的结构与DistilBERT[1], BERT6-PKD[2], BERT-of-Theseus[3] 相同。
  • T4-tiny的结构与 TinyBERT[4] 相同。
  • T3的结构与BERT3-PKD[2] 相同。

2.4.2 蒸馏装备

distill_config = DistillationConfig(temperature = 8, intermediate_matches = matches)
#其他参数为默认值

不同的模型用的matches咱们选用了以下装备:

一文详解TextBrewer

各种matches的界说在examples/matches/matches.py中。均运用GeneralDistiller进行蒸馏。

2.4.3练习装备

蒸馏用的学习率 lr=1e-4(除非特殊阐明)。练习30~60轮。

2.4.4英文试验成果

在英文试验中,咱们运用了如下三个典型数据集。

一文详解TextBrewer

咱们在下面两表中列出了DistilBERT, BERT-PKD, BERT-of-Theseus, TinyBERT 等公开的蒸馏成果,并与咱们的成果做比照。

Public results:

一文详解TextBrewer

Our results:

一文详解TextBrewer

阐明:

  1. 公开模型的名称后括号内是其等价的模型结构
  2. 蒸馏到T4-tiny的试验中,SQuAD使命上运用了NewsQA作为增强数据;CoNLL-2003上运用了HotpotQA的华章作为增强数据
  3. 蒸馏到T12-nano的试验中,CoNLL-2003上运用了HotpotQA的华章作为增强数据

2.4.5中文试验成果

在中文试验中,咱们运用了如下典型数据集。

一文详解TextBrewer

试验成果如下表所示。

一文详解TextBrewer

一文详解TextBrewer

阐明:

  1. 以RoBERTa-wwm-ext为教师模型蒸馏CMRC 2018和DRCD时,不选用学习率衰减
  2. CMRC 2018和DRCD两个使命上蒸馏时他们互作为增强数据
  3. Electra-base的教师模型练习设置参考自Chinese-ELECTRA
  4. Electra-small学生模型选用预练习权重初始化

3.中心概念

3.1Configurations

  • TrainingConfigDistillationConfig:练习和蒸馏相关的装备。

3.2Distillers

Distiller担任执行实践的蒸馏进程。目前完成了以下的distillers:

  • BasicDistiller: 供给单模型单使命蒸馏方法。可用作测验或简略试验。
  • GeneralDistiller (常用): 供给单模型单使命蒸馏方法,并且支撑中间层特征匹配,一般情况下引荐运用
  • MultiTeacherDistiller: 多教师蒸馏。将多个(同使命)教师模型蒸馏到一个学生模型上。暂不支撑中间层特征匹配
  • MultiTaskDistiller:多使命蒸馏。将多个(不同使命)单使命教师模型蒸馏到一个多使命学生模型。
  • BasicTrainer:用于单个模型的有监督练习,而非蒸馏。可用于练习教师模型

3.3用户界说函数

蒸馏试验中,有两个组件需求由用户供给,分别是callbackadaptor :

3.3.1Callback

回调函数。在每个checkpoint,保存模型后会被distiller调用,并传入当时模型。能够借由回调函数在每个checkpoint评测模型作用。

3.3.2Adaptor

将模型的输入和输出转换为指定的格式,向distiller解说模型的输入和输出,以便distiller根据不同的策略进行不同的核算。在每个练习步,batch和模型的输出model_outputs会作为参数传递给adaptoradaptor担任重新组织这些数据,返回一个字典。

更多细节可拜见完好文档中的阐明。

4.FAQ

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 章节。

点击关注,第一时间了解华为云新鲜技能~