Bert (Bi-directional Encoder Representations from Transformers) Pytorch 源码解读（二）

前言

这里是 Bert(Bi-directional Encoder Representations from Transformers) 源码解读的第二部分，第一部分主要介绍了 bert_model.py 文件中， bert 模型的定义。而第二部分为 BERT_Training.py 文件，该部分源码主要实现了 Bert 模型的预训练工作。

---

Bert 源码解读：

1. 模型结构源码： bert_model.py

2. 模型预训练源码：bert_training.py

3. 数据预处理源码：wiki_dataset.py

---

在开始前，先大致的介绍一下 bert 模型的预训练。bert 的预训练过程是个 Multi Task learning 的过程。其中同时进行的两个任务分别为：

1. Masked Language Modeling。

2. Next sentece Predict(Classification)。

Masked Language Mode(MLM)选择输入序列中的随机token样本，并用特殊的token[MASK]替换。MLM的目标是预测遮挡token时的交叉熵损失。BERT一致选择15%的输入token作为可能的替换。在所选的token中，80%替换为[MASK]， 10%保持不变，10%替换为随机选择的词汇表token。MLM 任务在 bert 中的主要任务是建立语言模型，他的原理和 Word2Vec 中的 CBOW 模型以及 Negative Sampling 算法思想类似。

Next Sentece Predict(NSP)是一种二分类损失，用于预测两个片段在原文中是否相互跟随。通过从文本语料库中提取连续的句子来创建积极的例子。反例是通过对来自不同文档的段进行配对来创建的。正、负样本的抽样概率相等。NSP的目标是为了提高下游任务的性能，比如自然语言推理，这需要对句子对之间的关系进行推理。

---

开始

1. Import & Config

from torch.utils.data import DataLoader
from dataset.wiki_dataset import BERTDataset
from models.bert_model import *
import tqdm
import pandas as pd
import numpy as np
import os


config = {}
config["train_corpus_path"] = "./pretraining_data/wiki_dataset/test_wiki.txt"
config["test_corpus_path"] = "./pretraining_data/wiki_dataset/test_wiki.txt"
config["word2idx_path"] = "./pretraining_data/wiki_dataset/bert_word2idx_extend.json"
config["output_path"] = "./output_wiki_bert"

config["batch_size"] = 1
config["max_seq_len"] = 200
config["vocab_size"] = 32162
config["lr"] = 2e-6
config["num_workers"] = 0

首先导入各种所需要的库，之前定义好的 Bert 模型，以及数据处理模块，其次对一些参数及路径进行设置。参数分别有：batch_size，max_seq_len 最大序列长度，vocab_size 字典大小，lr 学习率，num_workers 加载数据时的线程数。

2. Pretrain

class Pretrainer:
    def __init__(self, bert_model,
                 vocab_size,
                 max_seq_len,
                 batch_size,
                 lr,
                 with_cuda=True,
                 ):
        # 词量, 注意在这里实际字(词)汇量 = vocab_size - 20,
        # 因为前20个token用来做一些特殊功能, 如padding等等
        self.vocab_size = vocab_size
        self.batch_size = batch_size
        # 学习率
        self.lr = lr
        # 是否使用GPU
        cuda_condition = torch.cuda.is_available() and with_cuda
        self.device = torch.device("cuda:0" if cuda_condition else "cpu")
        # 限定的单句最大长度
        self.max_seq_len = max_seq_len
        # 初始化超参数的配置
        bertconfig = BertConfig(vocab_size=config["vocab_size"])
        # 初始化bert模型
        self.bert_model = bert_model(config=bertconfig)
        self.bert_model.to(self.device)
        # 初始化训练数据集
        train_dataset = BERTDataset(corpus_path=config["train_corpus_path"],
                                    word2idx_path=config["word2idx_path"],
                                    seq_len=self.max_seq_len,
                                    hidden_dim=bertconfig.hidden_size,
                                    on_memory=False,
                                    )
        # 初始化训练dataloader
        self.train_dataloader = DataLoader(train_dataset,
                                           batch_size=self.batch_size,
                                           num_workers=config["num_workers"],
                                           collate_fn=lambda x: x)
        # 初始化测试数据集
        test_dataset = BERTDataset(corpus_path=config["test_corpus_path"],
                                   word2idx_path=config["word2idx_path"],
                                   seq_len=self.max_seq_len,
                                   hidden_dim=bertconfig.hidden_size,
                                   on_memory=True,
                                   )
        # 初始化测试dataloader
        self.test_dataloader = DataLoader(test_dataset, batch_size=self.batch_size,
                                          num_workers=config["num_workers"],
                                          collate_fn=lambda x: x)
        # 初始化positional encoding
        self.positional_enc = self.init_positional_encoding(hidden_dim=bertconfig.hidden_size,
                                                            max_seq_len=self.max_seq_len)
        # 拓展positional encoding的维度为[1, max_seq_len, hidden_size]
        self.positional_enc = torch.unsqueeze(self.positional_enc, dim=0)

        # 列举需要优化的参数并传入优化器
        optim_parameters = list(self.bert_model.parameters())
        self.optimizer = torch.optim.Adam(optim_parameters, lr=self.lr)

        print("Total Parameters:", sum([p.nelement() for p in self.bert_model.parameters()]))

__init__ 方法中，主要进行 bert 预训练前的一些准备工作，包括：定义字典大小、batch size、learning rate、GPU的指定、最大句子长度、以及 Bert 的超参设置。同时对训练数据、测试数据、positional encoding进行初始化，定义optimizer。这部分代码很简单，每一部分都在注释中进行了说明，这里就不再进行一一说明。

3.Positional encoding

    def init_positional_encoding(self, hidden_dim, max_seq_len):
        position_enc = np.array([
            [pos / np.power(10000, 2 * i / hidden_dim) for i in range(hidden_dim)]
            if pos != 0 else np.zeros(hidden_dim) for pos in range(max_seq_len)])

        position_enc[1:, 0::2] = np.sin(position_enc[1:, 0::2])  # dim 2i
        position_enc[1:, 1::2] = np.cos(position_enc[1:, 1::2])  # dim 2i+1
        denominator = np.sqrt(np.sum(position_enc**2, axis=1, keepdims=True))
        position_enc = position_enc / (denominator + 1e-8)
        position_enc = torch.from_numpy(position_enc).type(torch.FloatTensor)
        return position_enc

Positional encoding 的初始化代码，这里使用的 Positional encoding 的初始化方式与传统 Transformer 相同，通过 sin 与 cos 函数生成固定的值，具体计算方法为：

 这里需要注意的是，在 Google 官方开源的 TensorFlow 版本的 bert 源码中，Positional encoding 是使用的跟字向量相同的方法初始化，经过模型训练得到的，而这里使用了传统 Transformer 的直接计算方法。

4. load_model & train

    def load_model(self, model, dir_path="./output"):
        # 加载模型
        checkpoint_dir = self.find_most_recent_state_dict(dir_path)
        checkpoint = torch.load(checkpoint_dir)
        model.load_state_dict(checkpoint["model_state_dict"], strict=False)
        torch.cuda.empty_cache()
        model.to(self.device)
        print("{} loaded for training!".format(checkpoint_dir))

    def train(self, epoch, df_path="./output_wiki_bert/df_log.pickle"):
        self.bert_model.train()
        self.iteration(epoch, self.train_dataloader, train=True, df_path=df_path)

load_model 用于模型的加载，train 方法对模型进行迭代训练。

5. NSP loss & MLM loss

    def compute_loss(self, predictions, labels, num_class=2, ignore_index=None):
        if ignore_index is None:
            loss_func = CrossEntropyLoss()
        else:
            loss_func = CrossEntropyLoss(ignore_index=ignore_index)
        return loss_func(predictions.view(-1, num_class), labels.view(-1))

    def get_mlm_accuracy(self, predictions, labels):
        predictions = torch.argmax(predictions, dim=-1, keepdim=False)
        mask = (labels > 0).to(self.device)
        mlm_accuracy = torch.sum((predictions == labels) * mask).float()
        mlm_accuracy /= (torch.sum(mask).float() + 1e-8)
        return mlm_accuracy.item()

compute_loss 方法计算 NSP 任务的 loss 值，get_mlm_accuracy 计算 MLM 任务的 loss 值。

6. Padding

    def padding(self, output_dic_lis):
        bert_input = [i["bert_input"] for i in output_dic_lis]
        bert_label = [i["bert_label"] for i in output_dic_lis]
        segment_label = [i["segment_label"] for i in output_dic_lis]
        bert_input = torch.nn.utils.rnn.pad_sequence(bert_input, batch_first=True)
        bert_label = torch.nn.utils.rnn.pad_sequence(bert_label, batch_first=True)
        segment_label = torch.nn.utils.rnn.pad_sequence(segment_label, batch_first=True)
        is_next = torch.cat([i["is_next"] for i in output_dic_lis])
        return {"bert_input": bert_input,
                "bert_label": bert_label,
                "segment_label": segment_label,
                "is_next": is_next}

对输入的数据进行 Padding 补齐，在训练数据初始化时，通过处理数据模块已近将输入序列的最大长度截断至设置好的max_seq_len 长度。而这里进行的是 Padding 补齐的操作。

7. iteration

    def iteration(self, epoch, data_loader, train=True, df_path="./output_wiki_bert/df_log.pickle"):
        if not os.path.isfile(df_path) and epoch != 0:
            raise RuntimeError("log DataFrame path not found and can't create a new one because we're not training from scratch!")
        if not os.path.isfile(df_path) and epoch == 0:
            df = pd.DataFrame(columns=["epoch", "train_next_sen_loss", "train_mlm_loss",
                                       "train_next_sen_acc", "train_mlm_acc",
                                       "test_next_sen_loss", "test_mlm_loss",
                                       "test_next_sen_acc", "test_mlm_acc"
                                       ])
            df.to_pickle(df_path)
            print("log DataFrame created!")

        str_code = "train" if train else "test"

        # Setting the tqdm progress bar
        data_iter = tqdm.tqdm(enumerate(data_loader),
                              desc="EP_%s:%d" % (str_code, epoch),
                              total=len(data_loader),
                              bar_format="{l_bar}{r_bar}")

        total_next_sen_loss = 0
        total_mlm_loss = 0
        total_next_sen_acc = 0
        total_mlm_acc = 0
        total_element = 0

        for i, data in data_iter:
            # print('IDX of data_iter:', i)
            data = self.padding(data)
            # 0. batch_data will be sent into the device(GPU or cpu)
            data = {key: value.to(self.device) for key, value in data.items()}
            positional_enc = self.positional_enc[:, :data["bert_input"].size()[-1], :].to(self.device)

            # 1. forward the next_sentence_prediction and masked_lm model
            mlm_preds, next_sen_preds = self.bert_model.forward(input_ids=data["bert_input"],
                                                                positional_enc=positional_enc,
                                                                token_type_ids=data["segment_label"])

            mlm_acc = self.get_mlm_accuracy(mlm_preds, data["bert_label"])
            next_sen_acc = next_sen_preds.argmax(dim=-1, keepdim=False).eq(data["is_next"]).sum().item()
            mlm_loss = self.compute_loss(mlm_preds, data["bert_label"], self.vocab_size, ignore_index=0)
            next_sen_loss = self.compute_loss(next_sen_preds, data["is_next"])
            loss = mlm_loss + next_sen_loss

            # 3. backward and optimization only in train
            if train:
                self.optimizer.zero_grad()
                loss.backward()
                # for param in self.model.parameters():
                #     print(param.grad.data.sum())
                self.optimizer.step()

            total_next_sen_loss += next_sen_loss.item()
            total_mlm_loss += mlm_loss.item()
            total_next_sen_acc += next_sen_acc
            total_element += data["is_next"].nelement()
            total_mlm_acc += mlm_acc

            if train:
                log_dic = {
                    "epoch": epoch,
                   "train_next_sen_loss": total_next_sen_loss / (i + 1),
                   "train_mlm_loss": total_mlm_loss / (i + 1),
                   "train_next_sen_acc": total_next_sen_acc / total_element,
                   "train_mlm_acc": total_mlm_acc / (i + 1),
                   "test_next_sen_loss": 0, "test_mlm_loss": 0,
                   "test_next_sen_acc": 0, "test_mlm_acc": 0
                }

            else:
                log_dic = {
                    "epoch": epoch,
                   "test_next_sen_loss": total_next_sen_loss / (i + 1),
                   "test_mlm_loss": total_mlm_loss / (i + 1),
                   "test_next_sen_acc": total_next_sen_acc / total_element,
                   "test_mlm_acc": total_mlm_acc / (i + 1),
                   "train_next_sen_loss": 0, "train_mlm_loss": 0,
                   "train_next_sen_acc": 0, "train_mlm_acc": 0
                }

            if i % 10 == 0:
                data_iter.write(str({k: v for k, v in log_dic.items() if v != 0 and k != "epoch"}))

        if train:
            df = pd.read_pickle(df_path)
            df = df.append([log_dic])
            df.reset_index(inplace=True, drop=True)
            df.to_pickle(df_path)
        else:
            log_dic = {k: v for k, v in log_dic.items() if v != 0 and k != "epoch"}
            df = pd.read_pickle(df_path)
            df.reset_index(inplace=True, drop=True)
            for k, v in log_dic.items():
                df.at[epoch, k] = v
            df.to_pickle(df_path)
            return float(log_dic["test_next_sen_loss"])+float(log_dic["test_mlm_loss"])

这部分代码主要为模型预训的每个 epoch 的迭代过程，以及训练中 log 的记录。

源码中，首先对要储存的 log 进行定义，然后将 total loss 和 accuracy 置零，其次就是计算每个 batch 的 mlm loss 与 nsp loss ，在进行反向传播，更新参数的过程。

最后记录每个 epoch 的 训练信息。

8. 模型保存

    def save_state_dict(self, model, epoch, dir_path="./output", file_path="bert.model"):
        if not os.path.exists(dir_path):
            os.mkdir(dir_path)
        save_path = dir_path+ "/" + file_path + ".epoch.{}".format(str(epoch))
        model.to("cpu")
        torch.save({"model_state_dict": model.state_dict()}, save_path)
        print("{} saved!".format(save_path))
        model.to(self.device)

这部分源码也比较简单，主要是模型的保存。

9. 训练代码

    def init_trainer(dynamic_lr, load_model=False):
        trainer = Pretrainer(BertForPreTraining,
                             vocab_size=config["vocab_size"],
                             max_seq_len=config["max_seq_len"],
                             batch_size=config["batch_size"],
                             lr=dynamic_lr,
                             with_cuda=True)
        if load_model:
            trainer.load_model(trainer.bert_model, dir_path=config["output_path"])
        return trainer


    start_epoch = 3
    train_epoches = 1
    trainer = init_trainer(config["lr"], load_model=True)
    # if train from scratch
    all_loss = []
    threshold = 0
    patient = 10
    best_f1 = 0
    dynamic_lr = config["lr"]
    for epoch in range(start_epoch, start_epoch + train_epoches):
        print("train with learning rate {}".format(str(dynamic_lr)))
        trainer.train(epoch)

        trainer.save_state_dict(trainer.bert_model, epoch, dir_path=config["output_path"],
                                file_path="bert.model")
        trainer.test(epoch)

这就是预训练的代码了，init_trainer 实例化前面定义的 Pretraniner 类，然后就是每个 epoch 调用类内的 train 方法来进行训练了，每个 epoch 保存一个模型，并进行测试。

---

总结

以上就是 Pytorch 版本 Bert 模型的预训练部分的全部源码了，相较于 TensorFlow 版本还是略显冗长，另外这里的 Positional encoding 部分是使用的传统 Transformer 中，直接计算得出的方式，与 Google 官方给出的源码不同。官方源码中将 Positional encoding 也作为参数，加入了训练，这会导致一定量的参数增加，但在如此大规模的训练数据面前，一点点参数的增加也就不值一提了，而究竟哪种 encoding 方式带来的效果更好，我自己还没有进行调研，暂未得出结论。

 

如有问题欢迎指正，转载请注明出处。