模型压缩实践系列之——layer dropout

会员服务 ·

模型压缩实践系列之——layer dropout

2020 年 2 月 21 日 AINLP

作者：邱震宇（华泰证券股份有限公司算法工程师）

知乎专栏：我的ai之路

本文算是一篇论文实践小总结，篇幅不大，主要是最近开始在研究模型的压缩工作，看了一些论文，感觉这块内容实践工程的部分很多，因此记录一下，防止忘记。

本次研究重点在于结构性剪枝，并选取了比较简单的layer dropout来实践验证其有效性。参考论文：Reducing Transformer Depth on Demand with Structured Dropout

layer dropout思想概述

layer dropout 属于结构化剪枝方法的范畴。非结构化剪枝包含目前比较经典的weight pruning，即通过对部分权重进行mask计算，间接得对权重进行剪枝。非结构化剪枝会改变模型原有的结构，在某些情况下反而会降低模型的计算效率。因此与此相对的，结构化剪枝正在逐渐被人们关注，结构化剪枝专注于对于相关的结构进行整体的剪枝，以确保最大限度保留模型原有的架构。

layer dropout这篇论文的基本思想很简单。

首先，它介绍了一个概念叫random structured Dropout。对于一个Transformer网络来说，存在很多重叠的相似结构，如attention head，FFN计算，或者相似的transformer层，虽然这些结构的权重参数都是各不相同，但是如果将这些相似结构的权重划分成组，然后以组为单位进行权重的剪枝，将大大减小模型的推断时间。用数学语言来描述就是，假设M是一个用于权重剪枝的mask矩阵，其值要么是0，要么是1。我们预先对网络结构进行组的划分，得到一系列权重组g，组内的所有权重对应的mask矩阵部分均相同： .

其次，它借鉴了DropConnect的思想，即每个权重的裁剪与否都是根据一个伯努利分布采样得到一个裁剪概率。这样做能够在一定程度上给模型增加正则的作用，使得模型更加鲁棒。

综上，layer dropout的主要步骤如下：以bert模型为例，

1、在训练阶段，设定给一个dropout的概率p。假定原先bert为12层，前向计算的时候，数据一层一层往下计算。此时，对于一层的计算，先使用一个随机的均匀分布采样一个概率，若该概率大于p，则略过这个层的计算，否则照常计算。

2、在inference的阶段，可以通过参数传递的方式，传入需要保留的layer层的index列表，然后在load模型后，根据保留的layer层信息，重新设计新的模型参数数据，去掉裁剪掉的layer层。（注意，这里只针对bert中的encoder层，其他embedding等权重不纳入裁剪范畴）

关于inference阶段的裁剪，论文给出了几个策略：

every other。即间隔裁剪，比如裁剪一半的层数，那么需要保留的层数可以为0,2,4,6...等偶数层。这种方法最简单，同时也能到一个相对平衡的网络。
search on valid。类似于调参步骤，根据裁剪的层数，组成不同的子模型，然后用验证集来调试，看哪种组合效果最好。这种方式虽然能够得到最好的效果，但是计算量很大，很费时，成本较高。
data driven strategy。这种方法还是比较巧妙的。首先我们有一个最终裁剪的比率P。然后我们对于每个layer，设定给一个裁剪概率 ,其中所有层的的均值需要等于P。如何实现这个限定条件呢？论文中提出将设计成一个非线性函数，该函数的建模对象是每个layer最后输出的激活值。最后将所有层的用softmax进行概率的归一化得到概率值。在inference的时候，我们选取概率最高的k个层来进行计算。

目前我先尝试了第一种策略，这种策略也是最容易的实现的，而且论文也说明该策略对不同任务的普适性较高，可以作为baseline先进行尝试。

该方法的主要思想还是接近于传统意义上的dropout，在训练时进行根据概率分布对整个layer层进行dropout，能够起到两个作用：一部分是加正则项，一部分是让模型适应layer被剪掉的情况，这样在inference的时候直接去掉某些层，模型就能更鲁棒。实践中也发现，如果直接对没有做过layerdropout的模型直接裁剪layer，会导致模型的精度呈直线型下滑。

上述过程一般是在bert的预训练时进行layer dropout,然后在finetune下游任务的时候，指定裁剪掉具体的层。这样做效果会比较好。然而，在实践时，受限于数据和硬件条件的限制，我们不具备预训练的条件，因此先尝试在finetune的时候做layerdropout，在inference做指定层裁剪。目的还是为了验证下该方法的有效性。

实践

这里主要说一下实践的几个细节注意点，实现方式主要是用pytorch。针对的任务是对话系统中的意图分类。基础模型是bert。目前pytorch的fairseq已经放出了开源的代码，可以参考。

1、在finetune阶段，可以在bertEncoder类中的forward方法中添加layerdropout的逻辑。

for layer_module in self.layer:
     dropout_probability = random.uniform(0, 1)
     if not self.training or (dropout_probability > layer_dropout_prob):
        hidden_states = layer_module(hidden_states, attention_mask)
        if output_all_encoded_layers:
           all_encoder_layers.append(hidden_states)

2、在inference阶段，裁剪layer要注意几点，首先是裁剪的方式主要是先用load将model的state_dict取出来，然后根据传入的保留层参数将保留层进行处理，比如保留0,2,4,6,8,10层，则将这6层重新赋予新的index 0,1,2,3,4,5，并存放到新的state_dict中，非bert_layer的其他权重原样copy到新的state_dict中，最后调用load_state_dict将新的state_dict赋予模型。另外，要注意需要更改bert_config.json的中的层数，因为在loadbert的权重时，层数与state_dict不一致时，会导致报错。inference阶段对state_dict的裁剪重构部分见fairseq/fairseq/checkpoint_utils.py中的prune_state_dict方法。下面贴一下该方法的核心部分代码：

def create_pruning_pass(layers_to_keep, layer_name):
        keep_layers = sorted(
            [int(layer_string) for layer_string in layers_to_keep.split(",")]
        )
        mapping_dict = {}
        for i in range(len(keep_layers)):
            mapping_dict[str(keep_layers[i])] = str(i)

        regex = re.compile("^{layer}.*\.layers\.(\d+)".format(layer=layer_name))
        return {"substitution_regex": regex, "mapping_dict": mapping_dict}

该方法根据传入的保留的layer的index，建立新的layer的index关系。

for pruning_pass in pruning_passes:
            if original_layer_number in pruning_pass["mapping_dict"] and pruning_pass[
                "substitution_regex"
            ].search(layer_name):
                new_layer_number = pruning_pass["mapping_dict"][original_layer_number]
                substitution_match = pruning_pass["substitution_regex"].search(
                    layer_name
                )
                new_state_key = (
                    layer_name[: substitution_match.start(1)]
                    + new_layer_number
                    + layer_name[substitution_match.end(1) :]
                )
                new_state_dict[new_state_key] = state_dict[layer_name]

这里就是在当前所有权重参数中，找bert的encoder层，然后根据create_pruning_pass方法得到的新的层数index映射关系，建立新的state_dict。

实验结果和分析

在做finetune训练时，随着layer dropout的概率的增加，模型的训练的效率是呈正比式提升，但是模型的精度也会有逐渐得下降。当前最优结果是将layer dropout概率设为0.3时，模型精度下降最小，在验证集上大概下降2个点。最后在测试集上做inference的测试，裁剪策略采用隔层裁剪，分别试验了裁剪一半以及裁剪1/4两种方式。最后发现裁剪一半下降的精度会更多，大概有7个点，而裁剪1/4则下降4个点。

究其原因还是在finetune时做layer dropout终究是达不到在预训练时做的效果。bert的整体效果很大一部分归功于预训练的学习，其对于模型整体的性能的贡献是占大头的，因此若还是使用正常方式做预训练，此时bert学到的信息相对完整，并没有学习到去掉某些层带来的变化，所以此时在finetune做裁剪，会让模型信息损失掉一部分。

当然，在finetune阶段做裁剪并不是一点作用没有。我还尝试了用正常的方式finetune一个bert模型，然后在inference阶段用裁剪策略裁掉一些层，最后发现模型的精度下降得很离谱，大概有接近80%的跌幅。这个实验证明在finetune阶段做layer dropout还是能让模型在一定程度上适应layer的裁剪，使得模型鲁棒性增强。

彩蛋

其实我还在之前的中文NER上做了layer dropout的实验。由于此任务使用tensorflow 1.12开发，还是使用的静态图模式，所以要实现训练时的layer drop还需要添加一些工程代码。与pytorch不同，静态图模式下，图一旦建立了就不会被改变，因此要实现每次前向计算时，都能够对不同的层进行dropout，就需要参考经典dropout那样的写法。我实现的大致思路如下：

对于每个bert的encoder layer层，其输出为layer_output,其后会将其赋予prev_output，作为下一层的输入。我在bert包中的modeling.py这个地方做了一个改动，代码如下：

dropout_prob = tf.random_uniform(shape=[], minval=0, maxval=1)
condition = (hidden_dropout_prob == 0.0 or dropout_prob > 0.3)
condition = tf.cast(condition,tf.bool)
gate = tf.cond(condition, lambda: tf.zeros_like(layer_input),
                               lambda: tf.ones_like(layer_input))
prev_output = layer_input*gate + layer_output*(1.-gate)

简单来说，就是设置一个gate，使得本层的输出要么还是本层的输入（即上一层的输出），要么是本层的encoder 输出。gate的取值由两个条件决定，当在训练状态（hidden_dropout_prob==0）或者由均匀分布采样得到的概率大于裁剪概率时，则让gate为全零的矩阵，此时本层计算纳入最后结果；否则，gate为全1的矩阵，则本层相当于计算结果不计入下一层的输入。

这种做法，与pytorch不同，并不能减小训练时的耗时，反而引入了一些计算会稍许增加一些计算。在inference的时候，目前我暂时只实现了使用estimator从checkpoint中读取模型，然后调用model来进行inference。其中，裁剪的部分仍然是在modeling.py中，此时只需要在for layer循环中，将要裁减掉的layer删去就可以。

最终，基于基础模型bert_mrc,在裁剪掉30%的bert层的情况下，模型的inference效率也提升了30%多。然而，模型的精度下降也有点多，相较于之前公开出来的f1分数，整整下降了10个百分点。其中，precision下降是较多的，反而recall下降的比较少。关于这个现象，还需要深入研究。

小结

通过上述两个任务的实验，可以知道layer dropout方法确实可以在提升模型的inference效率上有可观的效果，然而如果其使用的方法不妥，对于整个模型的精度是有影响的。可知，仅在finetune的时候做layer dropout，会在一定程度上降低模型精度，对于不同的任务，可能会有不同的影响。因此，其最好的实践还是在预训练的时候进行该过程。

后续的计划：

1、若有条件，尝试在预训练条件下使用layer dropout，然后在finetune的时候使用直接剪枝。

2、目前，对于TensorFlow 1.12，原先的正常模型是通过pb文件序列化，然后加载inference的，这种方式比加载checkpoint更加高效。但是，对于如何对计算图裁剪掉部分层，然后将其序列化为新的pb文件存在很多坑。目前我遇到的一个问题就是由于之前模型使用estimator接受dataset来进行高效训练，而若重新手动序列化新的pb之后，使用sess进行加载，会报这样的错误:

TensorFlow FailedPreconditionError: iterator has not been initialized

究其原因是estimator将dataset的iterator的初始化操作集成了，因此其operation也会序列到pb中，但是我找不到该op的名称。因此再次加载时，没办法调用这个op进行初始化。关于这个问题目前还在研究中，如果有同学踩过这个坑，欢迎讨论。在stackoverflow上也有类似的问题：

https://stackoverflow.com/questions/51419237/tensorflow-failedpreconditionerror-iterator-has-not-been-initialized stackoverflow.com

Restoring a Tensorflow model that uses Iterators stackoverflow.com

本文由作者授权AINLP原创发布于公众号平台，欢迎投稿，AI、NLP均可。原文链接，点击"阅读原文"直达：

https://zhuanlan.zhihu.com/p/106198038

推荐阅读

AINLP年度阅读收藏清单

中文NER任务实验小结报告——深入模型实现细节

flyai医疗智能问答比赛小结

BottleSum——文本摘要论文系列解读

抛开模型，探究文本自动摘要的本质——ACL2019 论文佳作研读系列

鼠年春节，用 GPT-2 自动写对联和对对联

用 GPT-2 自动写诗，从五言绝句开始

征稿启示 | 稿费+GPU算力+星球嘉宾一个都不少

关于AINLP

AINLP 是一个有趣有AI的自然语言处理社区，专注于 AI、NLP、机器学习、深度学习、推荐算法等相关技术的分享，主题包括文本摘要、智能问答、聊天机器人、机器翻译、自动生成、知识图谱、预训练模型、推荐系统、计算广告、招聘信息、求职经验分享等，欢迎关注！加技术交流群请添加AINLP君微信(id：AINLP2)，备注工作/研究方向+加群目的。