代码拉取完成，页面将自动刷新
TextPruner是一个为预训练语言模型设计的模型裁剪工具包，通过轻量、快速的裁剪方法对模型进行结构化剪枝，从而实现压缩模型体积、提升模型速度。
其他相关资源：
[Mar 21, 2022] （1.1版本功能更新) 添加了针对XLM、BART、T5和mT5模型的词表裁剪功能。
[Mar 4, 2022] TextPruner论文 TextPruner: A Model Pruning Toolkit for Pre-Trained Language Models 被 ACL 2022 demo录用。
[Jan 26, 2022] （1.0.1版本功能更新） 添加了对自监督裁剪的支持。通过TransformerPruningConfig中的use_logits设置。
| 章节 | 内容 |
|---|---|
| 简介 | TextPruner介绍 |
| 安装 | 安装要求与方法 |
| 裁剪模式 | 三种裁剪模式说明 |
| 使用方法 | TextPruner快速上手 |
| 实验结果 | 典型任务上的裁剪效果 |
| 常见问题 | 常见问题 |
| 关注我们 | - |
TextPruner是一个为预训练语言模型设计，基于PyTorch实现的模型裁剪工具包。它提供了针对预训练模型的结构化裁剪功能，通过识别并移除模型结构中不重要的结构与神经元，达到压缩模型大小、提升模型推理速度的目的。
TextPruner的主要特点包括：
TextPruner目前支持词表裁剪和transformer裁剪，参见裁剪模式。
要使用TextPruner，用户可以在python脚本中导入TextPruner或直接在命令行运行TextPruner命令行工具，参见使用方法。
TextPruner在典型任务上的实验效果，参见实验结果。
论文链接：TextPruner: A Model Pruning Toolkit for Pre-Trained Language Models
TextPruner目前支持Transformers库中的如下预训练模型:
| 模型 | 词表裁剪 | Transformer裁剪 |
|---|---|---|
| BERT | ||
| ALBERT | ||
| RoBERTa | ||
| ELECTRA | ||
| XLM-RoBERTa | ||
| XLM | ||
| BART | ||
| T5 | ||
| mT5 |
API文档参见在线文档
安装要求
使用pip安装
pip install textpruner
从源代码安装
git clone https://github.com/airaria/TextPruner.git
pip install ./textpruner
TextPruner提供了3种裁剪模式，分别为词表裁剪（Vocabulary Pruning），Transformer裁剪（Transformer Pruning）和流水线裁剪（Pipeline Pruning）。
预训练模型通常包含对具体任务来说冗余的词表。通过移除词表中未在具体任务未出现的token，可以实现减小模型体积，提升MLM等任务训练速度的效果。
另一种裁剪方式是裁剪每个transformer模块的大小。一些研究表明transformer中的注意力头（attention heads）并不是同等重要，移除不重要的注意力头并不会显著降低模型性能。TextPruner找到并移除每个transformer中“不重要”的注意力头和全连接层神经元，从而在减小模型体积的同时把对模型性能的影响尽可能降到最低。
在该模式中，TextPruner对给定模型依次分别进行Transformer裁剪和词表裁剪，对模型体积做全面的压缩。
Pruners执行具体的裁剪过程，configurations设置裁剪参数。它们名称的含义是不言自明的：
textpruner.VocabularyPrunertextpruner.TransformerPrunertextpruner.PipelinePrunertextpruner.GeneralConfigtextpruner.VocabularyPruningConfigtextpruner.TransformerPruningConfigTextPruner的API文档请参见在线文档。 Configurations的说明参见Configurations。 下面展示基本用法。
要进行词表裁剪，用户应提供一个文本文件或字符串列表（list of strings）。TextPruner将从model和tokenizer中移除未在文本文件或列表中出现过的token。
具体的例子参见examples/vocabulary_pruning和examples/vocabulary_pruning_xnli.
词表裁剪仅需3行代码：
from textpruner import VocabularyPruner
pruner = VocabularyPruner(model, tokenizer)
pruner.prune(dataiter=texts)
model和tokenizer是要裁剪的模型和对应的分词器texts是字符串列表（list of strings），一般为任务相关数据的文本，用以确定裁剪后的词表大小。TextPruner将从model和tokenizer中移除未在其中出现过的token。textpruner-cli \
--pruning_mode vocabulary \
--configurations gc.json vc.json \
--model_class XLMRobertaForSequenceClassification \
--tokenizer_class XLMRobertaTokenizer \
--model_path /path/to/model/and/config/directory \
--vocabulary /path/to/a/text/file
configurations：JSON格式的配置文件。model_class : 模型的完整类名，要求该类在当前目录下可访问。例如model_class是modeling.ModelClassName，那么当前目录下应存在modeling.py。如果model_class 中无模块名，那么TextPruner会试图从transformers库中导入model_class，如上面的例子。tokenizer_class : tokenizer的完整类名。要求该类在当前目录下可访问。如果tokenizer_class 中无模块名，那么TextPruner会试图从transformers库中导入tokenizer_class。model_path：模型、tokenizer和相关配置文件存放目录。vocabulary : 用于定义新词表的文本文件。TextPruner将从model和tokenizer中移除未在其中出现过的token。要在一个数据集上进行transformer裁剪，需要一个dataloader对象。每次迭代dataloader应返回一个batch，batch的格式应与训练模型时相同：包括inputs和labels（batch内容本身不必用和训练时相同）。
TextPruner需要模型返回的loss用以计算神经元的重要性指标。TextPruner会尝试猜测模型输出中的哪个元素是loss。如以下皆不成立：
output['loss']或output.loss得到，其中output是模型的输出那么用户应提供一个adaptor函数（以模型的输出为输入，返回loss）给TransformerPruner。
当运行于自监督裁剪模式，TextPruner需要模型返回的logits。此时需要adaptor函数返回logits。详细参见TransformerPruningConfig中的use_logits选项。
具体的例子参见examples/transformer_pruning
自监督裁剪的例子参见examples/transformer_pruning_xnli
裁剪一个12层预训练模型，每层的注意力头目标数为8，全连接层的目标维数为2048，通过4次迭代裁剪到目标大小：
from textpruner import TransformerPruner, TransformerPruningConfig
transformer_pruning_config = TransformerPruningConfig(
target_ffn_size=2048,
target_num_of_heads=8,
pruning_method='iterative',
n_iters=4)
pruner = TransformerPruner(model,transformer_pruning_config=transformer_pruning_config)
pruner.prune(dataloader=dataloader, save_model=True)
dataloader用于向pruner提供数据用于计算各个注意力头的神经元的重要性，从而决定裁剪顺序。textpruner-cli \
--pruning_mode transformer \
--configurations gc.json tc.json \
--model_class XLMRobertaForSequenceClassification \
--tokenizer_class XLMRobertaTokenizer \
--model_path ../models/xlmr_pawsx \
--dataloader_and_adaptor dataloader_script
configurations：JSON格式的配置文件。model_class : 模型的完整类名，要求该类在当前目录下可访问。例如model_class是modeling.ModelClassName，那么当前目录下应存在modeling.py。如果model_class 中无模块名，那么TextPruner会试图从transformers库中导入model_class，如上面的例子。tokenizer_class : tokenizer的完整类名。要求该类在当前目录下可访问。如果tokenizer_class 中无模块名，那么TextPruner会试图从transformers库中导入tokenizer_class。model_path：模型、tokenizer和相关配置文件存放目录。dataloader_and_adaptor : Python脚本文件，其中定义并初始化了dataloader和adaptor（adaptor可选）。流水线裁剪结合了transformer裁剪和词表裁剪。
具体的例子参见examples/pipeline_pruning
from textpruner import PipelinePruner, TransformerPruningConfig
transformer_pruning_config = TransformerPruningConfig(
target_ffn_size=2048, target_num_of_heads=8,
pruning_method='iterative',n_iters=4)
pruner = PipelinePruner(model, tokenizer, transformer_pruning_config=transformer_pruning_config)
pruner.prune(dataloader=dataloader, dataiter=texts, save_model=True)
textpruner-cli \
--pruning_mode pipeline \
--configurations gc.json tc.json vc.json \
--model_class XLMRobertaForSequenceClassification \
--tokenizer_class XLMRobertaTokenizer \
--model_path ../models/xlmr_pawsx \
--vocabulary /path/to/a/text/file \
--dataloader_and_adaptor dataloader_script
裁剪过程受配置对象（configuration objects）控制：
GeneralConfig：设置使用的device和输出目录。VocabularyPruningConfig：设置裁剪的阈值（token的词频低于此阈值将被裁减），以及是否裁剪lm_head。TransformerPruningConfig：Transformer裁剪过程参数的各种配置。它们用于不同的裁剪模式：
词表裁剪可接受GeneralConfig and VocabularyPruningConfig
VocabularyPruner(vocabulary_pruning_config= ..., general_config = ...)
Transformer裁剪可接受GeneralConfig and TransformerPruningConfig
TransformerPruner(transformer_pruning_config= ..., general_config = ...)
流水线裁剪可接受全部3种Config：
TransformerPruner(transformer_pruning_config= ..., vocabulary_pruning_config= ..., general_config = ...)
在Python脚本中，配置对象是dataclass对象；在命令行中，配置对象是JSON文件。
如果未向pruner提供相应的配置对象，TextPruner将使用默认配置。
配置对象的各个参数详细意义请参见GeneralConfig，VocabularyPruningConfig和TransformerPruningConfig API文档。
在Python脚本定义：
from textpruner import GeneralConfig, VocabularyPruningConfig, TransformerPruningConfig
from textpruner import VocabularyPruner, TransformerPruner, PipelinePruner
#GeneralConfig
general_config = GeneralConfig(device='auto',output_dir='./pruned_models')
#VocabularyPruningConfig
vocabulary_pruning_config = VocabularyPruningConfig(min_count=1,prune_lm_head='auto')
#TransformerPruningConfig
#Pruning with the given masks
transformer_pruning_config = TransformerPruningConfig(pruning_method = 'masks')
#TransformerPruningConfig
#Pruning on labeled dataset iteratively
transformer_pruning_config = TransformerPruningConfig(
target_ffn_size = 2048,
target_num_of_heads = 8,
pruning_method = 'iterative',
ffn_even_masking = True,
head_even_masking = True,
n_iters = 1,
multiple_of = 1
)
作为JSON文件：
GeneralConfig：gc.jsonVocabularyPruningConfig：vc.jsonTransformerPruningConfig：
textpruner.summary：显示模型参数摘要。textpruner.inference_time：测量与显示模型的推理耗时。例子：
from transformers import BertForMaskedLM
import textpruner
import torch
model = BertForMaskedLM.from_pretrained('bert-base-uncased')
print("Model summary:")
print(textpruner.summary(model,max_level=3))
dummy_inputs = [torch.randint(low=0,high=10000,size=(32,512))]
print("Inference time:")
textpruner.inference_time(model.to('cuda'),dummy_inputs)
Outputs:
Model summary:
LAYER NAME #PARAMS RATIO MEM(MB)
--model: 109,514,810 100.00% 417.77
--bert: 108,892,160 99.43% 415.39
--embeddings: 23,837,696 21.77% 90.94
--position_ids: 512 0.00% 0.00
--word_embeddings: 23,440,896 21.40% 89.42
--position_embeddings: 393,216 0.36% 1.50
--token_type_embeddings: 1,536 0.00% 0.01
--LayerNorm: 1,536 0.00% 0.01
--encoder
--layer: 85,054,464 77.66% 324.46
--cls
--predictions(partially shared): 622,650 0.57% 2.38
--bias: 30,522 0.03% 0.12
--transform: 592,128 0.54% 2.26
--decoder(shared): 0 0.00% 0.00
Inference time:
Device: cuda:0
Mean inference time: 1214.41ms
Standard deviation: 2.39ms
使用基于XLM-RoBERTa-base的分类模型，我们在多语言NLI任务PAWS-X的英文数据集上训练与测试，并使用TextPruner对训练好的模型进行裁剪。
我们从XNLI英文训练集中采样10万条样本作为词表，将XLM-RoBERTa模型的词表大小裁剪至这10万条样本的范围内，裁剪前后模型对比如下所示。
| Model | Total size (MB) | Vocab size | Acc on en (%) |
|---|---|---|---|
| XLM-RoBERTa-base | 1060 (100%) | 250002 | 94.65 |
| + Vocabulary Pruning | 398 (37.5%) | 23936 | 94.20 |
XLM-RoBERTa作为多语言模型，词表占了模型的很大一部分。通过只保留相关任务和语言的词表，可以显著减小模型体积,并且对模型准确率只有微弱影响。
使用（H,F）指示模型结构，其中H是平均每层注意力头数量，F是全连接层的维数（intermediate hidden size）。原始的XLM-RoBERTa-base模型可记为（12, 3072）。我们考虑裁剪到另外两种结构（8,2048）和（6，1536）。
使用长度512，batch size 32的数据作为输入测量推理时间：
| Model | Total size (MB) | Encoder size (MB) | Inference time (ms) | Speed up |
|---|---|---|---|---|
| (12, 3072) | 1060 | 324 | 1012 | 1.0x |
| (8, 2048) | 952 | 216 | 666 | 1.5x |
| (6, 1536) | 899 | 162 | 504 | 2.0x |
我们尝试使用不同的迭代次数进行transformer裁剪，各个模型的准确率变化如下表所示：
| Model | n_iters=1 | n_iters=2 | n_iters=4 | n_iters=8 | n_iters=16 |
|---|---|---|---|---|---|
| (12, 3072) | 94.65 | - | - | - | - |
| (8, 2048) | 93.30 | 93.60 | 93.60 | 93.85 | 93.95 |
| (8, 2048) with uneven heads | 92.95 | 93.50 | 93.95 | 94.05 | 94.25 |
| (6, 1536) | 85.15 | 89.10 | 90.90 | 90.60 | 90.85 |
| (6, 1536) with uneven heads | 45.35 | 86.45 | 90.55 | 90.90 | 91.95 |
表中的uneven heads指允许模型在不同层有不同的注意力头数。 可以看到，随着迭代次数的增加，裁剪后的模型的性能也随之提升。
最后，我们用PipelinePruner同时裁剪词表和transformer：
| Model | Total size (MB) | Speed up | Acc on en (%) |
|---|---|---|---|
| XLM-RoBERTa-base | 1060 (100%) | 1.0x | 94.65 |
| + Pipeline pruning to (8, 2048) with uneven heads | 227 (22%) | 1.5x | 93.75 |
Transformer裁剪过程使用了16次迭代，词表裁剪过程使用XNLI英文训练集中采样的10万条样本作为词表。整个裁剪过程在单张T4 GPU上耗时10分钟。
Q: TextPruner 是否支持 Tensorflow 2 ？
A: 不支持。
Q: 对于知识蒸馏和模型裁剪，能否给一些使用建议 ？
A: 知识蒸馏与模型裁剪都是减小模型体积的主流手段：
知识蒸馏通常可以获得更好的模型效果和更高的压缩率，但是蒸馏过程较消耗算力与时间；为了获得好的蒸馏效果，对大量数据的访问也是必不可少的。
在相同目标模型体积下，结构化无训练裁剪方法的性能通常低于知识蒸馏，但其优点是快速与轻量。裁剪过程最短可以在数分钟内完成，并且只需要少量标注数据进行指导。
（还有一类包含训练过程的裁剪方法，它们在保证模型性能的同时也可以取得很好的压缩效果。）
如果你对知识蒸馏感兴趣，可以参见我们的知识蒸馏工具包TextBrewer。
如果你想取得最好的模型压缩效果，或许可以尝试同时采用蒸馏与裁剪这两种手段。
如果TextPruner有助于您的研究，请引用以下论文:
@inproceedings{yang-etal-2022-textpruner,
title = "{T}ext{P}runer: A Model Pruning Toolkit for Pre-Trained Language Models",
author = "Yang, Ziqing and
Cui, Yiming and
Chen, Zhigang",
booktitle = "Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics: System Demonstrations",
month = may,
year = "2022",
address = "Dublin, Ireland",
publisher = "Association for Computational Linguistics",
url = "https://aclanthology.org/2022.acl-demo.4",
pages = "35--43"
}
欢迎关注哈工大讯飞联合实验室官方微信公众号，了解最新的技术动态
此处可能存在不合适展示的内容，页面不予展示。您可通过相关编辑功能自查并修改。
如您确认内容无涉及 不当用语 / 纯广告导流 / 暴力 / 低俗色情 / 侵权 / 盗版 / 虚假 / 无价值内容或违法国家有关法律法规的内容，可点击提交进行申诉，我们将尽快为您处理。
1. 开源生态
2. 协作、人、软件
3. 评估模型
马建仓 AI 助手