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系列: tinyBERT 介绍与代码训练。_第1张图片

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

Embedding-layer distillation

     BERT系列: tinyBERT 介绍与代码训练。_第2张图片

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

Hidden states based distillation  and Attention based distillation

BERT系列: tinyBERT 介绍与代码训练。_第3张图片

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

BERT系列: tinyBERT 介绍与代码训练。_第4张图片

 Prediction-layer distillation

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

BERT系列: tinyBERT 介绍与代码训练。_第5张图片

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

训练过程中的多次蒸馏

BERT系列: tinyBERT 介绍与代码训练。_第6张图片

 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

BERT系列: tinyBERT 介绍与代码训练。_第7张图片

我用的是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改几个参数就行。

BERT系列: tinyBERT 介绍与代码训练。_第8张图片

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

BERT系列: tinyBERT 介绍与代码训练。_第9张图片

  这是取数据,可以看到取了四个数据。可以在这里介绍一下,一句文本,会有一个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模型的输出从张量图上取下来。

BERT系列: tinyBERT 介绍与代码训练。_第10张图片

                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不同,结果就不同。

BERT系列: tinyBERT 介绍与代码训练。_第11张图片

 输入数据和上面很像。多了个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内部是怎么构成的,向量在里面是如何传递的。

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

你可能感兴趣的:(日常学习,网络搭建,bert,人工智能,深度学习)