BERT系列： tinyBERT 介绍与代码训练。

前情提要 ：

CSDN

上一篇文章讲述了BERT的全流程，但我们要做的是复现tinyBERT。BERT是一个大家族，里面有BERT-Tiny,BERT-Base，BERT-large等等。 他们的主要区别仅仅是结构不一样， 但是我们今天复现的tinyBERT是和他们不一样的，他的BERT在后面。 这就决定了它不只是结构不同，训练方式也是不同的。

结构差异 ：

        为了介绍结构的差异，我们先来读一个BERT的设置文档BERT config，一个config便可以决定一个BERT的结构。

{
  "hidden_size": 384,                   #决定token被编码的长度，即特征长度
  "intermediate_size": 1536,            # MLP第一次映射的长度，这里特征长度乘以4
  "max_position_embeddings": 512,       # 最大输入长度。
  "model_type": "tiny_bert",                
  "num_attention_heads": 12,            # 注意力头个数
  "num_hidden_layers": 4,               # 堆叠多少层
  "vocab_size": 30522                   # 训练词典个数，与训练语料有关
}

{
  "hidden_size": 768,           
  "intermediate_size": 3072,
  "max_position_embeddings": 512,
  "model_type": "bert",
  "num_attention_heads": 12,
  "num_hidden_layers": 12,
  "vocab_size": 30522
}

我主要把bert家族改变时主要改变的属性拿出来。我们可以看到他们结构上的不同之处。参数的具体意思，可以看前文得到。 这样tinyBERT自然会参数减少很多。 为了性能保持不变，tinyBERT并不像BERT家族那样在无标注上预训练，而是采取了蒸馏学习的方式进行训练。通过对BERT-base的蒸馏，得到了很好的性能。 （可以百度学习一下知识蒸馏）。

蒸馏学习

        BERT系列都是通过MLM和SEQ任务进行学习的。而tinyBERT则是蒸馏。蒸馏可以看作让参数的分布尽量靠近的意思。 很多网络是在网络的输出端蒸馏，希望整个网络的参数都学习到teacher模型。而tinyBERT则是在模型中间进行多次蒸馏。而且在整个训练过程也进行多次蒸馏。

前向过程中的多次蒸馏

        tinyBERT，对于n层的teacher bert，设计了一个mapping function ：n = g ( m )， 将student bert的第m层映射为原来的teacher的第n层。其实这个映射函数非常简单，就是n = k*g。k就是多少层当作tinyBERT的一层。当k=0时，对应的就是embedding layer。我们可以通过下图理解。图中仅为示例，tinyBERT每层的输出都去蒸馏学习Teacher net三层的输出，就是“一层顶三层”。
 

实际上的BERT-base有12层， 对于4层的tinyBERT，正好是三层对一层。 对于蒸馏学习，我们需要根据两个模型的参数或者输出来计算loss，更新模型。从上图中，我们可以看到一共有四种loss，下面分别介绍。

Embedding-layer distillation

     

 这个是对embedding 矩阵的蒸馏loss，说是矩阵，其实是计算两个模型embedding输出的MSEloss。 而因为Student的embedding层的特征维度和Teacher是不一样的，因此要乘上一个转换的映射矩阵，此矩阵在训练时学习。

Hidden states based distillation  and Attention based distillation

 前文我们提到，BERTlayer每一层包含两部分，一个是自注意力层，一个是MLP（也就是两层全连接）。 hidden states层的蒸馏就是指的全连接层的层间蒸馏，而attention自然指的是注意力层注意力分数矩阵，也就是q和k算出来的那个值 的蒸馏学习。

 Prediction-layer distillation

        这个蒸馏，属于是回归本质了。是对最终输出层输出结果的 softmax进行蒸馏学习。

 T是蒸馏学习中的温度系数。对输出蒸馏时，采用的是带温度系数的交叉熵loss。

训练过程中的多次蒸馏

 tinyBERT 并不是像其他蒸馏那样，直接根据成品的Teacher model在分类时蒸馏，而是采用了模仿teacher训练过程的方法。他的蒸馏分两步：General Distillation 与 Task-specific Distillation.  前者是模仿BERT在大规模语料库进行预训练，后者则是在特定的任务上进行分类训练。值得注意的是，在预训练蒸馏阶段，使用的Teacher 模型是仅仅经过预训练未微调的BERT，而在特定任务分类蒸馏训练阶段，使用的Teacher 模型是在特定任务上经过微调的BERT。

到这里，我们已经完成了tinyBERT的多阶段，多层次蒸馏，得到了性能很好，又很快的tinyBERT模型。oumeideta！

上代码

 下载基于torch的代码，让我们看看具体的训练过程吧！

Pretrained-Language-Model/TinyBERT at master · huawei-noah/Pretrained-Language-Model · GitHub

https://huggingface.co/bert-base-uncased/tree/main

事前准备

如果你想跟我一样跑通代码，你需要下载一下内容

1， BERT-base模型 

这个就是teacher模型啦。bert-base-uncased at main

2，wiki数据集。

关于WIKI下载

维基下载页面说明（指南） - 走看看

我下载的这个：

https://dumps.wikimedia.org/enwiki/20220601/enwiki-20220601-pages-articles-multistream1.xml-p1p41242.bz2

然后使用WikiExtractor提取和整理数据集中的文本，使用步骤如下

• pip install wikiextractor
• python -m wikiextractor.WikiExtractor -o 【目标文件】-b 【大小】 【源文件】

 注意这个 大小 他是够多少就形成一个文件，也就是如果写了1M 就会形成很多个1M的输出文件 我用的就是1M

​

我用的是AA里的00号文件，我们只是跑示例而已。

glue 数据集

https://github.com/nyu-mll/GLUE-baselines

可以只下载一个任务，比如QNLI。 自行参照readme 更改任务名即可。

代码阅读

预训练蒸馏。

--pregenerated_data data  --teacher_model /home/model/bert  --do_lower_case --train_batch_size 4 --output_dir model --student_model config

代码文件是general_distill.py

在pycharm 运行编辑配置中   输入上面的参数 。 大部分都是位置参数。自己训练的时候需要按自己的文件夹调整。 注意tiny_bert的config文件 在开头参数里，把bert的config改几个参数就行。

​

我们主要看的是模型的训练，所以数据，设置什么的都先不看了。 所以直接到这一步。

​

  这是取数据，可以看到取了四个数据。可以在这里介绍一下，一句文本，会有一个tokenizer的东西对他做处理，生成3个东西。

        input_ids :  代码里给每个词都编了号 ，这个就是那个号

        input_mask: 输入中标记哪些词模型需要考虑，哪些不需要考虑。比如，输入长度不够时会被pad，这些pad的mask就为0，不需要考虑。

        segment_ids： 标记属于哪个句子

其余的两个则是预训练任务需要的。

        lm_labels_ids: 被遮盖的词的编号

        is_next:输出的两个句子是不是上下文。

                student_atts, student_reps = student_model(input_ids, segment_ids, input_mask)

        tinyBERT的前向。我们可以看看student模型的样子。

class TinyBertForPreTraining(BertPreTrainedModel):
    def __init__(self, config, fit_size=768):
        super(TinyBertForPreTraining, self).__init__(config)
        self.bert = BertModel(config)
        self.cls = BertPreTrainingHeads(
            config, self.bert.embeddings.word_embeddings.weight)
        self.fit_dense = nn.Linear(config.hidden_size, fit_size)
        self.apply(self.init_bert_weights)

    def forward(self, input_ids, token_type_ids=None,
                attention_mask=None, masked_lm_labels=None,
                next_sentence_label=None, labels=None):
        sequence_output, att_output, pooled_output = self.bert(
            input_ids, token_type_ids, attention_mask)
        tmp = []
        for s_id, sequence_layer in enumerate(sequence_output):
            tmp.append(self.fit_dense(sequence_layer))
        sequence_output = tmp

        return att_output, sequence_output

很简单的写法。BertModel 会根据你的设置返回一个BERT模型。他的输出是可以调控的。这里输出了三个东西。

      sequence output 包含了每一层bertlayer的输出

      att_output  包含了每一层的注意力分数。

      pooled_output 最后一层的输出的pool值。可以取meanpool  也可以取cls。

而fit_dense就是前文所提到的，因为两个模型的特征维度不同，需要使用的转换矩阵。

cls等会再用到。（好吧 也没看到用在哪里）

                teacher_reps, teacher_atts, _ = teacher_model(input_ids, segment_ids, input_mask)

得到teacher的注意力分数和输出。

                teacher_reps = [teacher_rep.detach() for teacher_rep in teacher_reps]  # speedup 1.5x
                teacher_atts = [teacher_att.detach() for teacher_att in teacher_atts]

因为是不更新teacher模型的，所以这里要把teacher模型的输出从张量图上取下来。

                new_teacher_reps = [teacher_reps[i * layers_per_block] for i in range(student_layer_num + 1)]
                new_student_reps = student_reps

可以看到 层数对应的计算。 取对应层的最后一层的输出。也就是第3，6，9，12层的输出。

                for student_att, teacher_att in zip(student_atts, new_teacher_atts):
                    student_att = torch.where(student_att <= -1e2, torch.zeros_like(student_att).to(device),
                                              student_att)
                    teacher_att = torch.where(teacher_att <= -1e2, torch.zeros_like(teacher_att).to(device),
                                              teacher_att)
                    att_loss += loss_mse(student_att, teacher_att)

计算注意力分数的loss。

                new_teacher_reps = [teacher_reps[i * layers_per_block] for i in range(student_layer_num + 1)]
                new_student_reps = student_reps

                for student_rep, teacher_rep in zip(new_student_reps, new_teacher_reps):
                    rep_loss += loss_mse(student_rep, teacher_rep)

你会发现 BERT的输出有13个，其中第一个是embedding的输出。 上面这个代码段，会保留0，3，6，9，12层的输出。这个loss计算 可以得到 L_hide 和L_embd。在预训练阶段 没有pred的loss。 只需要这三个loss进行回传更新模型。

任务蒸馏

   在特定任务上进行蒸馏。代码是 task_distill   参数是下面的，我用的时QNLI数据集。去掉了数据增广，因为数据增广只是提前处理，而不是在训练时增广。 记得将前面预训练阶段得到的tinybert的模型保存文件名字改为: pytorch_model.bin.。torch只认这个。

--teacher_model /home/model/bert  --student_model /home/reapper/tinybert/model --data_dir /home/dataset/glue/QNLI --task_name qnli --output_dir model/qnli  --do_lower_case --learning_rate 3e-5  --num_train_epochs  3  --eval_step 100 --max_seq_length 128 --train_batch_size 32 --pred_distill

数据部分我就不说了,GLUE的数据都是有现成的读的方法。介绍一下QNLI，QNLI的一条数据是有两句话，然后判断两句话是蕴含关系还是非蕴含关系，也就是二分类任务。

我们直接上模型。

    if not args.do_eval:
        teacher_model = TinyBertForSequenceClassification.from_pretrained(args.teacher_model, num_labels=num_labels)
        teacher_model.to(device)

    student_model = TinyBertForSequenceClassification.from_pretrained(args.student_model, num_labels=num_labels)
    student_model.to(device)

我们发现他是用同一个模型读入的，前面提过了，只要config不同，结果就不同。

 输入数据和上面很像。多了个seq_lengths，表示文本本身词的个数（padding前）。

                student_logits, student_atts, student_reps = student_model(input_ids, segment_ids, input_mask,
                                                                           is_student=True)

                with torch.no_grad():
                    teacher_logits, teacher_atts, teacher_reps = teacher_model(input_ids, segment_ids, input_mask)

与前文一样，执行模型前向，teacher不计算梯度。不同的是，输出多了一个logits，熟悉的人都知道，这个就是分类头的输出，我们进模型内部看一下。

class TinyBertForSequenceClassification(BertPreTrainedModel):
    def __init__(self, config, num_labels, fit_size=768):
        super(TinyBertForSequenceClassification, self).__init__(config)
        self.num_labels = num_labels
        self.bert = BertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, num_labels)
        self.fit_dense = nn.Linear(config.hidden_size, fit_size)
        self.apply(self.init_bert_weights)

    def forward(self, input_ids, token_type_ids=None, attention_mask=None,
                labels=None, is_student=False):

        sequence_output, att_output, pooled_output = self.bert(input_ids, token_type_ids, attention_mask,
                                                               output_all_encoded_layers=True, output_att=True)

        logits = self.classifier(torch.relu(pooled_output))

        tmp = []
        if is_student:
            for s_id, sequence_layer in enumerate(sequence_output):
                tmp.append(self.fit_dense(sequence_layer))
            sequence_output = tmp
        return logits, att_output, sequence_output

看到 出了bert外，还有一个分类头，一个用于蒸馏的线性映射fit_dense。没了，就这么简简单单。从bert得到pool的特征后，直接relu分类就完事了。

                    teacher_layer_num = len(teacher_atts)
                    student_layer_num = len(student_atts)
                    assert teacher_layer_num % student_layer_num == 0
                    layers_per_block = int(teacher_layer_num / student_layer_num)
                    new_teacher_atts = [teacher_atts[i * layers_per_block + layers_per_block - 1]
                                        for i in range(student_layer_num)]

                    for student_att, teacher_att in zip(student_atts, new_teacher_atts):
                        student_att = torch.where(student_att <= -1e2, torch.zeros_like(student_att).to(device),
                                                  student_att)
                        teacher_att = torch.where(teacher_att <= -1e2, torch.zeros_like(teacher_att).to(device),
                                                  teacher_att)

                        tmp_loss = loss_mse(student_att, teacher_att)
                        att_loss += tmp_loss

att_loss： 与预训练时一致。

                    new_teacher_reps = [teacher_reps[i * layers_per_block] for i in range(student_layer_num + 1)]
                    new_student_reps = student_reps
                    for student_rep, teacher_rep in zip(new_student_reps, new_teacher_reps):
                        tmp_loss = loss_mse(student_rep, teacher_rep)
                        rep_loss += tmp_loss

hideloss 和 embdloss 与预训练一致。

                if not args.pred_distill:
                else:
                    if output_mode == "classification":
                        cls_loss = soft_cross_entropy(student_logits / args.temperature,
                                                      teacher_logits / args.temperature)
                    elif output_mode == "regression":
                        loss_mse = MSELoss()
                        cls_loss = loss_mse(student_logits.view(-1), label_ids.view(-1))

                    loss = cls_loss

pred loss 我们可以看到，他是有开关的，也就是可以选择不训练。 他这里是 如果是蒸馏前面的层，这里就不用蒸馏分类层了。 相当于bert输出已经蒸馏过了。 如果不蒸馏前面的层，就蒸馏分类层，。

我们现在是分类模式，所以loss是用softCE计算的，如果你看了模型蒸馏，你肯定知道这个东西。其实就是本来对于分类任务，他的target都是one-hot形式的：【0，1】 这样的，现在变成了【0.1，0.9】这样的。 这也是蒸馏的原理之一。   温度系数这里是1.

至此，我们得到了所有的四个loss 回传即可完成模型训练。

事后

   至此，我们完成了tinyBERT的训练，得到了一个快速版的BERT。

    代码我读起来很快，主要因为我对bert和transformer已经有了很深的了解。 如果对BERT不了解，这个博客也是很好的机会，你可以慢慢调试，一步一步的看BERT内部是怎么构成的，向量在里面是如何传递的。

        这篇博文大部分内容来自个人经验和论文，如果有什么错误，请联系我，谢谢。