从零实现BERT中的两个预训任务（下）

1 引言

各位朋友大家好，欢迎来到月来客栈，我是掌柜空字符。

经过上一篇文章[1]的介绍，我们已经详细地了解了如何从零来构建训练BERT模型所需要的NSP和MLM任务数据集。总的来说，整个数据集的构建流程如图1所示。

在接下来的这一篇文章中，掌柜将带着大家如何来一步步地编码实现NSP和MLM这两个任务以及整个模型的训练和复用过程。

2 预训练任务实现

根据文章BERT原理与NSP和MLM可知，BERT的预训练过程包括两个任务：NSP和MLM。同时，为了使得大家能够对这两部分的代码实现有着更加清晰的认识与理解，掌柜将先分别来实现这两个任务，最后再将两者结合到一起来实现BERT的预训练任务。

2.1 NSP任务实现

由于NSP任务实现起来较为简单，所以掌柜这里就直接贴出代码：

xxxxxxxxxx
30
1
class BertForNextSentencePrediction(nn.Module):
2

3
    def __init__(self, config, bert_pretrained_model_dir=None):
4
        super(BertForNextSentencePrediction, self).__init__()
5
        if bert_pretrained_model_dir is not None:
6
            self.bert = BertModel.from_pretrained(config, bert_pretrained_model_dir)
7
        else:
8
            self.bert = BertModel(config)
9
        self.classifier = nn.Linear(config.hidden_size, 2)
10

11
    def forward(self,
12
                input_ids,  # [src_len, batch_size]
13
                attention_mask=None,  # [batch_size, src_len] mask掉padding部分的内容
14
                token_type_ids=None,  # [src_len, batch_size] 如果输入模型的只有一个序列，那么这个参数也不用传值
15
                position_ids=None,
16
                next_sentence_labels=None):  # [batch_size,]
17
        pooled_output, _ = self.bert(
18
            input_ids=input_ids,
19
            attention_mask=attention_mask,
20
            token_type_ids=token_type_ids,
21
            position_ids=position_ids)
22
        # pooled_output: [batch_size, hidden_size]
23
        seq_relationship_score = self.classifier(pooled_output)
24
        # seq_relationship_score: [batch_size, 2]
25
        if next_sentence_labels is not None:
26
            loss_fct = nn.CrossEntropyLoss()
27
            loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_labels.view(-1))
28
            return loss
29
        else:
30
            return seq_relationship_score

上述代码便是整个NSP任务的实现，可以看到其本质上就是一个文本分类任务，仅仅只用取BERT模型最后一层输出的[CLS]做一个分类任务即可。详细介绍可以参见文章基于BERT预训练模型的文本分类任务，这里掌柜就不再赘述了。

2.2 MLM任务实现

相比较于NSP，对于实现MLM任务来说则稍微复杂了一点点。它需要将BERT模型整个最后一层的输出进行一次变换和标准化，然后再做Token级的分类任务来预测被掩盖部分对应的Token值，这个网络结构如图2所示。

如图2所示便是构造MLM任务的流程示意。首先取BERT模型最后一层的输出，形状为[src_len,batch_size,hidden_size]；然后再经过一次（非）线性变换和标准化，形状同样为[src_len,batch_size,hidden_size]；最后再经过一个分类层对每个Token进行分类处理便得到了最后的预测结果，形状为[src_len,batch_size,vocab_size]。

此时我们便可以定义类BertForLMTransformHead来完成上述3个步骤：

xxxxxxxxxx
20
1
class BertForLMTransformHead(nn.Module):
2
    def __init__(self, config, bert_model_embedding_weights=None):
3
        super(BertForLMTransformHead, self).__init__()
4
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
5
        if isinstance(config.hidden_act, str):
6
            self.transform_act_fn = get_activation(config.hidden_act)
7
        else:
8
            self.transform_act_fn = config.hidden_act
9
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=1e-12)
10
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
11
        if bert_model_embedding_weights is not None:
12
            self.decoder.weight = nn.Parameter(bert_model_embedding_weights)
13
        # [hidden_size, vocab_size]
14
        self.decoder.bias = nn.Parameter(torch.zeros(config.vocab_size))
15
    def forward(self, hidden_states):
16
        hidden_states = self.dense(hidden_states)  # [src_len, batch_size, hidden_size]
17
        hidden_states = self.transform_act_fn(hidden_states)  # [src_len, batch_size, hidden_size]
18
        hidden_states = self.LayerNorm(hidden_states)  # [src_len, batch_size, hidden_size]
19
        hidden_states = self.decoder(hidden_states)
20
        return hidden_states # [src_len, batch_size, vocab_size]

在上述代码中，第4-8行用来定义相应的（非）线性变换。第8-9行则是用来定义对应的标准化和最后的分类层。第11-12则是用来判断最后分类层中的权重参数是否复用BERT模型Token Embedding中的权重参数，因为MLM任务最后的预测类别就等于Token Embedding中的各个词，所以最后分类层中的权重参数可以复用Token Embedding中的权重参数[2]。

同时，还有一个细节的地方在于，在Token Embedding中词表的形状为[vocab_size,hidden_size]，而最后一个分类层权重参数的形状为[hidden_size,vocab_size]，因此按道理说这里的权重bert_model_embedding_weights应该要转置后才能复用。不过在第12行代码中我们却直接将其赋值给了最后一层的权重参数，这是为什么呢？

原来在PyTorch的nn.Linear()中，权重参数默认的初始化方式就是

xxxxxxxxxx
1
1
self.weight = Parameter(torch.Tensor(out_features, in_features))

也就是将两个维度进行了交换，所以刚好我们在复用时就不用做其它任何操作了。

紧接着，我们便可以通过如下代码来实现MLM任务：

xxxxxxxxxx
33
1
class BertForMaskedLM(nn.Module):
2
    def __init__(self, config, bert_pretrained_model_dir=None):
3
        super(BertForMaskedLM, self).__init__()
4
        if bert_pretrained_model_dir is not None:
5
            self.bert = BertModel.from_pretrained(config, bert_pretrained_model_dir)
6
        else:
7
            self.bert = BertModel(config)
8
        weights = None
9
        if config.use_embedding_weight:
10
            weights = self.bert.bert_embeddings.word_embeddings.embedding.weight
11
            logging.info(f"## 使用token embedding中的权重矩阵作为输出层的权重！{weights.shape}")
12
        self.classifier = BertForLMTransformHead(config, weights)
13
        self.config = config
14
    def forward(self,
15
                input_ids,  # [src_len, batch_size]
16
                attention_mask=None,  # [batch_size, src_len] mask掉padding部分的内容
17
                token_type_ids=None,  # [src_len, batch_size]
18
                position_ids=None,
19
                masked_lm_labels=None):  # [src_len,batch_size]
20
        _, all_encoder_outputs = self.bert(
21
            input_ids=input_ids,attention_mask=attention_mask,
22
            token_type_ids=token_type_ids,position_ids=position_ids)
23
        sequence_output = all_encoder_outputs[-1]  # 取Bert最后一层的输出
24
        # sequence_output: [src_len, batch_size, hidden_size]
25
        prediction_scores = self.classifier(sequence_output)
26
        # prediction_scores: [src_len, batch_size, vocab_size]
27
        if masked_lm_labels is not None:
28
            loss_fct = nn.CrossEntropyLoss(ignore_index=-1)
29
            masked_lm_loss = loss_fct(prediction_scores.reshape(-1, self.config.vocab_size),
30
                                      masked_lm_labels.reshape(-1))
31
            return masked_lm_loss
32
        else:
33
            return prediction_scores  # [src_len, batch_size, vocab_size]

在上述代码中，第4-7行用于返回得到原始的BERT模型；第9-10行则是取Token Embedding中的权重参数；第12行则是返回得到MLM任务实例化后的类对象；第20-24行为返回得到BERT模型的所有层输出，并只取最后一层；第25行则是完成最后MLM中的分类任务；第27-33行则是根据标签是否为空来返回不同的输出结果。

到此，对于MLM任务的实现掌柜就介绍完了。

2.3 整体实现

经过上面两节内容的介绍，此时再来整体实现NSP和MLM任务那就变得十分容易了。整体实现代码如下所示：

xxxxxxxxxx
43
1
class BertForPretrainingModel(nn.Module):
2
    def __init__(self, config, bert_pretrained_model_dir=None):
3
        super(BertForPretrainingModel, self).__init__()
4
        if bert_pretrained_model_dir is not None:
5
            self.bert = BertModel.from_pretrained(config, bert_pretrained_model_dir)
6
        else:  # 如果没有指定预训练模型路径，则随机初始化整个网络权重
7
            self.bert = BertModel(config)
8
        weights = None
9
        if 'use_embedding_weight' in config.__dict__ and config.use_embedding_weight:
10
            weights = self.bert.bert_embeddings.word_embeddings.embedding.weight
11
            logging.info(f"## 使用token embedding中的权重矩阵作为输出层的权重！{weights.shape}")
12
        self.mlm_prediction = BertForLMTransformHead(config, weights)
13
        self.nsp_prediction = nn.Linear(config.hidden_size, 2)
14
        self.config = config
15

16
    def forward(self, input_ids,  # [src_len, batch_size]
17
                attention_mask=None,  # [batch_size, src_len] mask掉padding部分的内容
18
                token_type_ids=None,  # [src_len, batch_size]
19
                position_ids=None,
20
                masked_lm_labels=None,  # [src_len,batch_size]
21
                next_sentence_labels=None):  # [batch_size]
22
        pooled_output, all_encoder_outputs = self.bert(
23
            input_ids=input_ids,
24
            attention_mask=attention_mask,
25
            token_type_ids=token_type_ids,
26
            position_ids=position_ids)
27
        sequence_output = all_encoder_outputs[-1]  # 取Bert最后一层的输出
28
        # sequence_output: [src_len, batch_size, hidden_size]
29
        mlm_prediction_logits = self.mlm_prediction(sequence_output)
30
        # mlm_prediction_logits: [src_len, batch_size, vocab_size]
31
        nsp_pred_logits = self.nsp_prediction(pooled_output)
32
        # nsp_pred_logits： [batch_size, 2]
33
        if masked_lm_labels is not None and next_sentence_labels is not None:
34
            loss_fct = nn.CrossEntropyLoss(ignore_index=-1)
35
            mlm_loss = loss_fct(mlm_prediction_logits.reshape(-1, self.config.vocab_size),
36
                                masked_lm_labels.reshape(-1))
37
            nsp_loss = loss_fct(nsp_pred_logits.reshape(-1, 2),
38
                                next_sentence_labels.reshape(-1))
39
            total_loss = mlm_loss + nsp_loss
40
            return total_loss, mlm_prediction_logits, nsp_pred_logits
41
        else:
42
            return mlm_prediction_logits, nsp_pred_logits
43
        # [src_len, batch_size, vocab_size], [batch_size, 2]

在上述代码中，第12-13行分别用来返回得到实例化后的MLM和NSP任务模型；第22-26行是返回BERT模型的所有输出；第27-32行是分别取BERT模型输出的不同部分来分别进行后续的MLM和NSP任务；第33-43行是根据是否有标签输入来返回不同的输出结果，同时需要注意的是第39行返回的是NSP+MLM两个任务的损失和作为整体模型的损失值；第40-42行是根据条件返回模型不同的结果。

到此对于NSP和MLM任务模型的实现部分就介绍完了。不过掌柜在这里同样要提醒大家的是，在逐行阅读代码的时候最好是将各个变量的维度一起带进去，弄清楚每一步计算后各个变量维度的变化，这样才能更好的理解整个模型。

3 模型训练与微调

在实现完整个模型的前向传播过程后，接下来便可以实现模型训练部分的代码了。

3.1 训练实现

对于整个模型训练部分的代码其实和掌柜在前面几个微调任务中介绍的差不多，只是为了能更加清楚地知道训练模的训练过程型掌柜在这里加入了一些通过Tensorboard可视化的代码。由于这部分代码较长掌柜就分两部分来进行介绍。

xxxxxxxxxx
34
1
def train(config):
2
    model = BertForPretrainingModel(config,config.pretrained_model_dir)
3
    last_epoch = -1
4
    if os.path.exists(config.model_save_path):
5
        checkpoint = torch.load(config.model_save_path)
6
        last_epoch = checkpoint['last_epoch']
7
        loaded_paras = checkpoint['model_state_dict']
8
        model.load_state_dict(loaded_paras)
9
        logging.info("## 成功载入已有模型，进行追加训练......")
10
    model.train()
11
    bert_tokenize = BertTokenizer.from_pretrained(config.pretrained_model_dir).tokenize
12
    data_loader = LoadBertPretrainingDataset(vocab_path=config.vocab_path, ...)
13
    train_iter, test_iter, val_iter = \
14
        data_loader.load_train_val_test_data(test_file_path=config.test_file_path, ...)
15
      
16
    no_decay = ["bias", "LayerNorm.weight"]
17
    optimizer_grouped_parameters = [
18
        {
19
            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
20
            "weight_decay": config.weight_decay,
21
            "initial_lr": config.learning_rate
22
        },
23
        {
24
            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
25
            "weight_decay": 0.0,
26
            "initial_lr": config.learning_rate
27
        },
28
    ]      
29
    optimizer = AdamW([{"params": model.parameters(),
30
                        "initial_lr": config.learning_rate}])
31
    scheduler = get_polynomial_decay_schedule_with_warmup(optimizer,
32
                                                          config.num_warmup_steps,
33
                                                          config.num_train_steps,
34
                                                          last_epoch=last_epoch)

在上述代码中，第2行是实例化模型对象；第4-9行是查看本地是否存在相关模型（指对训练过程中保存的模型进行追加训练），这里值得一说的是之所以也要保存last_epoch这个参数是为了同时能够恢复学习率、以及Tensoboard中相关可视化变量在上一次模型保存是时的状态；第11-14行是载入训练模型时所需要用到的数据集；第16-28行是筛选模型中哪些参数需要进行权重衰减（也就是$L_2$$L_2$正则化），哪些参数不需要进行权重衰减，同时根据筛选条件可知所有bias和LayerNorm相关的参数都不需要进行正则化处理；第29行是定义优化器，并且是通过initial_lr来指定的学习率的，因为后续需要用到动态学习率调整的策略；第31-34行是指定动态学习率的调整策略，关于这部分内容的介绍可以参考文章Transformers之自定义学习率动态调整。

进一步，实现模型的迭代正反向传播过程，代码如下：

xxxxxxxxxx
34
1
    for epoch in range(config.epochs):
2
        losses = 0
3
        for idx, (b_token_ids, b_segs, b_mask, b_mlm_label, b_nsp_label) in enumerate(train_iter):
4
          #......
5
            loss, mlm_logits, nsp_logits = model(input_ids=b_token_ids,
6
                                                 attention_mask=b_mask,
7
                                                 token_type_ids=b_segs,
8
                                                 masked_lm_labels=b_mlm_label,
9
                                                 next_sentence_labels=b_nsp_label)
10
            optimizer.zero_grad()
11
            loss.backward()
12
            optimizer.step()
13
            scheduler.step()
14
            losses += loss.item()
15
            mlm_acc, _, _, nsp_acc, _, _ = accuracy(mlm_logits, nsp_logits, b_mlm_label,
16
                                                    b_nsp_label, data_loader.PAD_IDX)
17
            if idx % 20 == 0:
18
                config.writer.add_scalar('Training/Loss', loss.item(), scheduler.last_epoch)
19
                config.writer.add_scalar('Training/Learning Rate', 
20
                                         scheduler.get_last_lr()[0], scheduler.last_epoch)
21
                config.writer.add_scalars(main_tag='Training/Accuracy',
22
                                          tag_scalar_dict={'NSP': nsp_acc,'MLM': mlm_acc},
23
                                          global_step=scheduler.last_epoch)
24
        train_loss = losses / len(train_iter)
25
        if (epoch + 1) % config.model_val_per_epoch == 0:
26
            mlm_acc, nsp_acc = evaluate(config, val_iter, model, data_loader.PAD_IDX)
27
            config.writer.add_scalars(main_tag='Testing/Accuracy',
28
                                      tag_scalar_dict={'NSP': nsp_acc,'MLM': mlm_acc},
29
                                      global_step=scheduler.last_epoch)
30
            if mlm_acc > max_acc:
31
                max_acc = mlm_acc
32
                state_dict = model.state_dict()
33
            torch.save({'last_epoch': scheduler.last_epoch,'model_state_dict': state_dict},
34
                         config.model_save_path)

在上述代码中，第5-9行是模型的前向传播过程；第10-13行是反向传播及相关参数的更新过程；第14-16行分别是损失的累计和两个预训练任务准确率的计算；第18-23行是对训练过程中模型的损失、学习率和准确率进行可视化；第26-29行是计算模型在测试集上的准确率，并同时进行可视化处理；第30-34行则是保存在测试集上取得最大准确率时所对应的模型。

到此，对于模型训练部分的内容就介绍完了。同时，模型在训练过程中将会有类似如下所示的输出：

xxxxxxxxxx
4
1
- INFO: Epoch: [1/120], Batch[ 0/7836], Train loss : 1.897, Train mlm acc: 0.0, nsp acc: 0.537
2
- INFO: Epoch: [1/120], Batch[20/7836], Train loss : 1.786, Train mlm acc: 0.0, nsp acc: 0.562
3
- INFO: Epoch: [1/120], Batch[40/7836], Train loss : 1.812, Train mlm acc: 0.0, nsp acc: 0.489
4
          ......

最终，模型的在宋词数据集上损失值和准确率的变化如图3所示。

如图3所示，左边为NSP和MLM两个预训练任务整体损失的变化情况；右边蓝色曲线和橙色曲线分别为NSP和MLM这两个任务在训练集上的准确率的变化情况。不过虽然模型看似在训练集上有着不错的训练效果，但是在测试集上却显得不那么尽如人意，如图4所示。

如图4所示便是模型在测试集上的表现情况，其中蓝色曲线和粉色曲线分别表示NSP和MLM这两个任务在测试集上的准确率变化情况。从图中可以发现，模型大约在20万次迭代后准确率就没有发生明显地变换了。在调整过多次参数组合后依旧没有得到一个好的效果，掌柜猜测这可能是由于使用的宋词数据集比较特别，不像普通的白话文那样测试集和训练集的句式比较类似。不过后续掌柜也会继续尝试调整模型，并将相应的结果推送到代码仓库中。

3.2 推理实现

在模型训练部分的内容介绍完毕后，下面我们再来看模型推理部分的实现。对于推理部分的实现总体思路为：①将测试样本构造为模型所接受的输入格式；②通过模型前向传播得到预测结果输出；③对模型输出结果进行格式化处理得到最终的预测结果。对于模型预测部分的实现代码如下：

xxxxxxxxxx
17
1
def inference(config, sentences=None, masked=False, language='en'):
2
    bert_tokenize = BertTokenizer.from_pretrained(config.pretrained_model_dir).tokenize
3
    data_loader = LoadBertPretrainingDataset(vocab_path=config.vocab_path,
4
                                             ....
5
                                             masked_rate=0.15)  # 15% Mask掉
6
    token_ids, pred_idx, mask = data_loader.make_inference_samples(sentences,masked=masked,
7
                                                                   language=language)
8
    model = BertForPretrainingModel(config,config.pretrained_model_dir)
9
    if os.path.exists(config.model_save_path):
10
        checkpoint = torch.load(config.model_save_path)
11
        loaded_paras = checkpoint['model_state_dict']
12
        model.load_state_dict(loaded_paras)
13
        logging.info("## 成功载入已有模型进行推理......")
14
    with torch.no_grad():
15
        mlm_logits, _ = model(input_ids=token_ids,
16
                              attention_mask=mask)
17
    pretty_print(token_ids, mlm_logits, pred_idx,data_loader.vocab.itos, sentences, language)

在上述代码中，第3-5行为初始化类LoadBertPretrainingDataset，同时需要说明的是由于是预测场景，所以构造样本时masked_rate可以是任意值；第6-7行则是将传入的测试样本转换为模型所接受的形式，其中第6行masked是用来指定输入的测试样本有没进行mask操作，如果没有则自动按masked_rate的比例进行mask操作，第7行是指定测试样本的语种类型；第9-12行则是载入本地保存好的模型；第15-16行是得到模型前向传播的输出结果；第17行是根据模型的前向传播输出结果来格式化得到最终的输出形式。

最终，可以通过如下方式来完成模型的推理过程，代码如下：

xxxxxxxxxx
11
1
if __name__ == '__main__':
2
    config = ModelConfig()
3
    train(config)
4
    sentences_1 = ["I no longer love her, true, but perhaps I love her.",
5
                   "Love is so short and oblivion so long."]
6

7
    sentences_2 = ["我住长江头，君住长江尾。",
8
                   "日日思君不见君，共饮长江水。",
9
                   "此水几时休，此恨何时已。",
10
                   "只愿君心似我心，定不负相思意。"]
11
    inference(config, sentences_2, masked=False, language='zh')

上述代码运行结束后将会看到类似如下所示的结果：

xxxxxxxxxx
9
1
- INFO: ## 成功载入已有模型进行推理……
2
- INFO:  ### 原始：我住长江头，君住长江尾。
3
- INFO:   ## 掩盖：我住长江头，[MASK]住长[MASK]尾。
4
- INFO:   ## 预测：我住长江头，君住长河尾。  
5
- INFO: ====================
6
- INFO:  ### 原始：日日思君不见君，共饮长江水。
7
- INFO:   ## 掩盖：日日思君不[MASK]君，共[MASK]长江水。
8
- INFO:   ## 预测：日日思君不见君，共饮长江水。
9
#   ......

需要注意的是，由于目前模型在测试集上的结果并不理想，所以上述展示的是模型在训练集上的推理结果。

3.3 模型微调

在介绍完整个预训练模型实现过程后，最后一步就是如何将训练得到的模型继续运用在下游任务中。当然，实现这一目的也非常简单，只需要将保存好的模型重新命名为pytorch_model.bin，然后替换掉之前的文件即可。这样就可以想前面介绍的几个下游任务一样对模型进行微调了。

4 总结

在这篇文章中，掌柜首先回顾了BERT预训练任务中MLM和NSP数据集的构建流程；然后分别介绍了NSP和MLM任务模型的实现过程；最后，掌柜逐一详细介绍了BERT预训练模型训练部分以及推理部分的代码实现。到此，对于整个BERT部分内容的介绍就结束了。

本次内容就到此结束，感谢您的阅读！如果你觉得上述内容对你有所帮助，欢迎点赞分享！若有任何疑问与建议，请添加掌柜微信nulls8（备注来源）或留言进行交流。青山不改，绿水长流，我们月来客栈见！

引用

[1] 基于BERT预训练模型的NSL和MLM任务（上）

[2] https://github.com/google-research/bert/blob/master/run_pretraining.py

推荐阅读

[1]BERT原理与NSL和MLM

[2]BERT模型的分步实现

[3]基于BERT预训练模型的文本分类任务

[4]基于BERT预训练模型的文本蕴含任务

[5]基于BERT预训练模型的SWAG问答任务

[6]基于BERT预训练模型的SQuAD回答任务

[7] This Post Is All You Need（上卷）——层层剥开Transformer