github.com/airaria/TextBrewer @v0.2.1.post1

English | 中文说明

<img src="https://github.com/airaria/TextBrewer/raw/v0.2.1.post1/pics/banner.png" width="500"/>









<a href="https://github.com/airaria/TextBrewer/blob/master/LICENSE">
    <img alt="GitHub" src="https://img.shields.io/github/license/airaria/TextBrewer.svg?color=blue&style=flat-square">
</a>
<a href="https://textbrewer.readthedocs.io/">
    <img alt="Documentation" src="https://img.shields.io/website?down_message=offline&label=Documentation&up_message=online&url=https%3A%2F%2Ftextbrewer.readthedocs.io">
</a>    
<a href="https://pypi.org/project/textbrewer">
    <img alt="PyPI" src="https://img.shields.io/pypi/v/textbrewer">
</a>    
<a href="https://github.com/airaria/TextBrewer/releases">
    <img alt="GitHub release" src="https://img.shields.io/github/v/release/airaria/TextBrewer?include_prereleases">
</a>

TextBrewer is a PyTorch-based model distillation toolkit for natural language processing. It includes various distillation techniques from both NLP and CV field and provides an easy-to-use distillation framework, which allows users to quickly experiment with the state-of-the-art distillation methods to compress the model with a relatively small sacrifice in the performance, increasing the inference speed and reducing the memory usage.

Check our paper through ACL Anthology or arXiv pre-print.

Full Documentation

News

Nov 11, 2020

• Updated to 0.2.1:
• More flexible distillation: Supports feeding different batches to the student and teacher. It means the batches for the student and teacher no longer need to be the same. It can be used for distilling models with different vocabularies (e.g., from RoBERTa to BERT).

• Faster distillation: Users now can pre-compute and cache the teacher outputs, then feed the cache to the distiller to save teacher's forward pass time.

  See Feed Different batches to Student and Teacher, Feed Cached Values for details of the above features.

• MultiTaskDistiller now supports intermediate feature matching loss.

• Tensorboard now records more detailed losses (KD loss, hard label loss, matching losses...).

See details in releases.

August 27, 2020

We are happy to announce that our model is on top of GLUE benchmark, check leaderboard.

Click here to see old news

Aug 24, 2020

• Updated to 0.2.0.1:
• fixed bugs in MultiTaskDistiller and training loops.

Jul 29, 2020

• Updated to 0.2.0:
  • Added the support for distributed data-parallel training with DistributedDataParallel: TrainingConfig now accpects the local_rank argument. See the documentation of TrainingConfig for detail.
• Added an example of distillation on the Chinese NER task to demonstrate distributed data-parallel training. See examples/msra_ner_example.

Jul 14, 2020 * Updated to 0.1.10: * Now supports mixed precision training with Apex! Just set fp16 to True in TrainingConfig. See the documentation of TrainingConfig for detail. * Added data_parallel option in TrainingConfig to enable data parallel training and mixed precision training work together.

Apr 26, 2020

• Added Chinese NER task (MSRA NER) results.
• Added results for distilling to T12-nano model, which has a similar strcuture to Electra-small.
• Updated some results of CoNLL-2003, CMRC 2018 and DRCD.

Apr 22, 2020

• Updated to 0.1.9 (added cache option which speeds up distillation; fixed some bugs). See details in releases.
• Added experimential results for distilling Electra-base to Electra-small on Chinese tasks.
• TextBrewer has been accepted by ACL 2020 as a demo paper, please use our new bib entry.

Mar 17, 2020

• Added CoNLL-2003 English NER distillation example. See examples/conll2003_example.

Mar 11, 2020

• Updated to 0.1.8 (Improvements on TrainingConfig and train method). See details in releases.

Mar 2, 2020

• Initial public version 0.1.7 has been released. See details in releases.

Table of Contents

Section	Contents
Introduction	Introduction to TextBrewer
Installation	How to install
Workflow	Two stages of TextBrewer workflow
Quickstart	Example: distilling BERT-base to a 3-layer BERT
Experiments	Distillation experiments on typical English and Chinese datasets
Core Concepts	Brief explanations of the core concepts in TextBrewer
FAQ	Frequently asked questions
Known Issues	Known issues
Citation	Citation to TextBrewer
Follow Us	-

Introduction

Textbrewer is designed for the knowledge distillation of NLP models. It provides various distillation methods and offers a distillation framework for quickly setting up experiments.

The main features of TextBrewer are:

• Wide-support: it supports various model architectures (especially transformer-based models)
• Flexibility: design your own distillation scheme by combining different techniques; it also supports user-defined loss functions, modules, etc.
• Easy-to-use: users don't need to modify the model architectures
• Built for NLP: it is suitable for a wide variety of NLP tasks: text classification, machine reading comprehension, sequence labeling, ...

TextBrewer currently is shipped with the following distillation techniques:

• Mixed soft-label and hard-label training
• Dynamic loss weight adjustment and temperature adjustment
• Various distillation loss functions: hidden states MSE, attention-matrix-based loss, neuron selectivity transfer, ...
• Freely adding intermediate features matching losses
• Multi-teacher distillation
• ...

TextBrewer includes:

1. Distillers: the cores of distillation. Different distillers perform different distillation modes. There are GeneralDistiller, MultiTeacherDistiller, BasicTrainer, etc.
2. Configurations and presets: Configuration classes for training and distillation, and predefined distillation loss functions and strategies.
3. Utilities: auxiliary tools such as model parameters analysis.

To start distillation, users need to provide

1. the models (the trained teacher model and the un-trained student model)
2. datasets and experiment configurations

TextBrewer has achieved impressive results on several typical NLP tasks. See Experiments.

See Full Documentation for detailed usages.

Architecture

Installation

• Requirements
• Python >= 3.6
• PyTorch >= 1.1.0
• TensorboardX or Tensorboard
• NumPy
• tqdm
• Transformers >= 2.0 (optional, used by some examples)

• Apex == 0.1.0 (optional, mixed precision training)

• Install from PyPI

shell pip install textbrewer

• Install from the Github source

shell git clone https://github.com/airaria/TextBrewer.git pip install ./textbrewer

Workflow

• Stage 1: Preparation:
• Train the teacher model
• Define and initialize the student model

• Construct a dataloader, an optimizer, and a learning rate scheduler

• Stage 2: Distillation with TextBrewer:

• Construct a TraningConfig and a DistillationConfig, initialize a distiller
• Define an adaptor and a callback. The adaptor is used for adaptation of model inputs and outputs. The callback is called by the distiller during training
• Call the train method of the distiller

Quickstart

Here we show the usage of TextBrewer by distilling BERT-base to a 3-layer BERT.

Before distillation, we assume users have provided:

• A trained teacher model teacher_model (BERT-base) and a to-be-trained student model student_model (3-layer BERT).
• a dataloader of the dataset, an optimizer and a learning rate builder or class scheduler_class and its args dict scheduler_dict.

Distill with TextBrewer:

import textbrewer
from textbrewer import GeneralDistiller
from textbrewer import TrainingConfig, DistillationConfig

# Show the statistics of model parameters
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)

# Define an adaptor for interpreting the model inputs and outputs
def simple_adaptor(batch, model_outputs):
      # The second and third elements of model outputs are the logits and hidden states
    return {'logits': model_outputs[1],
            'hidden': model_outputs[2]}

# Training configuration 
train_config = TrainingConfig()
# Distillation configuration
# Matching different layers of the student and the teacher
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}])

# Build 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)

# Start!
with distiller:
    distiller.train(optimizer, dataloader, num_epochs=1, scheduler_class=scheduler_class, scheduler_args = scheduler_args, callback=None)

Examples can be found in the examples directory :

• examples/random_token_example : a simple runable toy example which demonstrates the usage of TextBrewer. This example performs distillation on the text classification task with random tokens as inputs.
• examples/cmrc2018_example (Chinese): distillation on CMRC 2018, a Chinese MRC task, using DRCD as data augmentation.
• examples/mnli_example (English): distillation on MNLI, an English sentence-pair classification task. This example also shows how to perform multi-teacher distillation.
• examples/conll2003_example (English): distillation on CoNLL-2003 English NER task, which is in form of sequence labeling.
• examples/msra_ner_example (Chinese): This example distills a Chinese-ELECTRA-base model on the MSRA NER task with distributed data-parallel training(single node, muliti-GPU).

Experiments

We have performed distillation experiments on several typical English and Chinese NLP datasets. The setups and configurations are listed below.

Models

• For English tasks, the teacher model is BERT-base-cased.
• For Chinese tasks, the teacher models are RoBERTa-wwm-ext and Electra-base released by the Joint Laboratory of HIT and iFLYTEK Research.

We have tested different student models. To compare with public results, the student models are built with standard transformer blocks except for BiGRU which is a single-layer bidirectional GRU. The architectures are listed below. Note that the number of parameters includes the embedding layer but does not include the output layer of each specific task.

English models

Model	#Layers	Hidden size	Feed-forward size	#Params	Relative size
BERT-base-cased (teacher)	12	768	3072	108M	100%
T6 (student)	6	768	3072	65M	60%
T3 (student)	3	768	3072	44M	41%
T3-small (student)	3	384	1536	17M	16%
T4-Tiny (student)	4	312	1200	14M	13%
T12-nano (student)	12	256	1024	17M	16%
BiGRU (student)	-	768	-	31M	29%

Chinese models

Model	#Layers	Hidden size	Feed-forward size	#Params	Relative size
RoBERTa-wwm-ext (teacher)	12	768	3072	102M	100%
Electra-base (teacher)	12	768	3072	102M	100%
T3 (student)	3	768	3072	38M	37%
T3-small (student)	3	384	1536	14M	14%
T4-Tiny (student)	4	312	1200	11M	11%
Electra-small (student)	12	256