BERT源码实现与解读(Pytorch)

本文源码地址: https://github.com/iioSnail/chaotic-transformer-tutorials

本文内容

论文地址: https://arxiv.org/abs/1810.04805

在BERT的论文中，描述的基本的都是模型如何训练等，对于本身的模型架构并没有过多的说明。这也可以理解，因为BERT架构本身也确实比较简单，就是一些TransformerEncoder的堆叠。虽然这么说，但不看代码很多人还是无法具体知道BERT是怎么样的，所以本文就来搭建一个BERT模型，并使用论文中提到的MLM任务和NSP任务对模型进行训练。

本篇需要大家有Transformer的基础，默认你已经熟悉Transformer，所以本篇会直接使用Pytorch中的nn.Transformer进行实现。

本篇参考了BERT-pytorch的部分代码，如果感兴趣，可以看看。

如果不想看原论文，可以直接参考我的笔记BERT论文阅读

相关文章：

Pytorch中 nn.Transformer的使用详解与Transformer的黑盒讲解: https://blog.csdn.net/zhaohongfei_358/article/details/126019181

万字逐行解析与实现Transformer，并进行德译英实战: https://blog.csdn.net/zhaohongfei_358/article/details/126085246

环境准备

导入本文需要的包：

import math
import copy

import torch
import torchtext
from torch import nn
import torch.nn.functional as F
from torchtext.vocab import build_vocab_from_iterator
1
2
3
4
5
6
7
8

torch.__version__
1

'1.12.1+cpu'
1

torchtext.__version__
1

'0.13.1'
1

BERT模型定义

原始BERT使用了两种任务对BERT进行预训练，所以我们将数据集和训练放在后面，先把BERT模型定义出来。

BERT Embedding

在Transformer中，对token的编码使用的是token embedding+position embedding，而在BERT中，增加了segment embedding，即文本的段落信息，在原论文中，bert的inputs是两句话，该embedding用于区分这是第一句话还是第二句话。我们先将这3中Embedding定义出来：

class TokenEmbedding(nn.Embedding):
    def __init__(self, vocab_size, embed_size):
        super().__init__(vocab_size, embed_size, padding_idx=0)
1
2
3

Token Embedding就是一个nn.Emebdding，和Transformer一致。

class PositionalEmbedding(nn.Module):

    def __init__(self, d_model, max_len=512):
        super().__init__()

        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model).float()
        pe.require_grad = False

        position = torch.arange(0, max_len).float().unsqueeze(1)
        div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return self.pe[:, :x.size(1)]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

Position Embedding和Transformer也一致。

class SegmentEmbedding(nn.Embedding):
    def __init__(self, embed_size=512):
        super().__init__(3, embed_size, padding_idx=0)
1
2
3

Segment Embedding也是一个nn.Embedding，但需要注意的是其词典大小只有3，其中0是填充，1代表第一句话，2代表第二句话。

定义完上面三个Embedding类，就可以把BERT Embedding类定义出来了，其比较简单，就是它们三个相加：

class BERTEmbedding(nn.Module):

    def __init__(self, vocab_size, embed_size, dropout=0.1):
        """
        :param vocab_size: token的词典大小
        :param embed_size: 词向量大小
        """
        super().__init__()
        self.token = TokenEmbedding(vocab_size=vocab_size, embed_size=embed_size)
        self.position = PositionalEmbedding(d_model=self.token.embedding_dim)
        self.segment = SegmentEmbedding(embed_size=self.token.embedding_dim)
        self.dropout = nn.Dropout(p=dropout)
        self.embed_size = embed_size

    def forward(self, sequence, segment_label):
        x = self.token(sequence) + self.position(sequence) + self.segment(segment_label)
        return self.dropout(x)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

BERT

class BERT(nn.Module):

    def __init__(self, vocab_size, hidden=768, n_layers=12, attn_heads=12, dropout=0.1):
        """
        :param vocab_size: 词典大小
        :param hidden: 隐状态大小，即词向量大小
        :param n_layers: TransformerEncoderLayer的层数
        :param attn_heads: Multi-head Attention的head数
        :param dropout: dropout rate
        """

        super().__init__()
        self.hidden = hidden
        self.n_layers = n_layers
        self.attn_heads = attn_heads

        # 论文中提到它们使用的feed_forward_hidden的大小为hidde_size*4
        feed_forward_hidden = hidden * 4

        # 定义BERT的embedding
        self.embedding = BERTEmbedding(vocab_size=vocab_size, embed_size=hidden)

        # 在论文中提到，BERT中Transformer的激活函数使用的是GELU
        transformer_encoder = nn.TransformerEncoderLayer(d_model=hidden, nhead=attn_heads, dim_feedforward=feed_forward_hidden, dropout=dropout, activation=F.gelu, batch_first=True)

        # 多层TransformerEncoder堆叠
        self.transformer_blocks = nn.ModuleList([copy.deepcopy(transformer_encoder) for _ in range(n_layers)])

    def forward(self, x, segment_info):
        """
        BERT前向传递
        :param x: 要被bert编码的向量，例如[[1,2,3,4,5,5,4,3,2,1,0,0]]，
                  即一对儿句子，包含有两句话(1,2,3,4,5)和(5,4,3,2,1),其中0是填充
        :param segment_info: 句子的段落信息，例如[1,1,1,1,1,2,2,2,2,2,0,0]，
                             即前5个token属于第一句话，接下来5个token是第二句话，
                             0是填充，不属于任何话。
        :return: 所有token经过bert后包含上下文的隐状态，例如Shape为(1, 12, 768)，
                 即1个句子，12个token，每个token被编码成了768维的向量
        """

        # 定义key_padding_mask，
        # 例如句子：唱跳Rap篮球
        # 对应的key_padding_mask为：[F, F, F, F, F, F, T, T] # 注意False为不遮掩，True为遮掩。搞反的话容易出现nan问题
        # 最后两个False是对进行mask
        key_padding_mask = x <= 0

        # 将index编码成向量
        x = self.embedding(x, segment_info)

        # 将编码后的向量经过TransformerEncoder一层一层传递
        for transformer in self.transformer_blocks:
            x = transformer.forward(x, src_key_padding_mask=key_padding_mask)

        return x
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

定义好BERT模型类的代码后，我们来初始化一下。在原论文中，作者提出了两种不同大小的BERT，分别为：

• $\mathbf{BER}{\mathbf{T}}_{\mathbf{BASE}}$：12层，hidden size为768，Attention Head数为12，共110M个参数。
• $\mathbf{BER}{\mathbf{T}}_{\mathbf{LARGE}}$：24层，hidden size为1024，Attention Head数为16，共340M个参数。

我们接下来就将两种BERT定义出来（词向量是参考了Hugging Face上bert_base_uncased的bert模型）：

bert_base = BERT(vocab_size=30522, hidden=768, n_layers=12, attn_heads=12)
bert_large = BERT(vocab_size=30522, hidden=1024, n_layers=24, attn_heads=16)
print("bert_base参数量: ", sum([param.nelement() for param in bert_base.parameters()]))
print("bert_large参数量: ", sum([param.nelement() for param in bert_large.parameters()]))
1
2
3
4

bert_base参数量:  108497664
bert_large参数量:  333566976
1
2

可以看到，bert_base和bert_large的参数量与原文的描述基本一致。接下来简单尝试使用一下BERT:

x = torch.LongTensor([[1,2,3,4,5,5,4,3,2,1,0,0]])
segment_info = torch.LongTensor([[1,1,1,1,1,2,2,2,2,2,0,0]])
print("bert_base outputs size:", bert_base(x, segment_info).size())
print("bert_large outputs size:", bert_large(x, segment_info).size())
1
2
3
4

bert_base outputs size: torch.Size([1, 12, 768])
bert_large outputs size: torch.Size([1, 12, 1024])
1
2

预训练BERT

在BERT原论文中，作者使用了两种任务对BERT进行预训练，分别是MLM任务(masked language model)和NSP(Next Sentence Prediction)，我们这里也使用这两种任务对BERT进行简单的预训练。

MLM任务(masked language model):

MLM任务简介：MLM任务就是把一个句子中的部分token给替换掉，然后让bert去结合上下文来预测被替换掉的词是什么。

在BERT原论文中，作者是将句子中15%的token给替换掉，而对于被替换的15%的token使用的规则为：选择80%的token使用[MASK]这个特殊的token进行替换，10%的token随机替换成其他token，10%不做任何变化。

这里我们简单一点，对一个句子只选择2个token进行替换，并且只替换成[MASK]

Next Sentence Prediction(NSP)任务：

NSP任务简介：NSP任务就是预测传给BERT的两句话是不是一对儿，是一个二分类任务。预测方式就是使用输入的第一个token[CLS]的输出接全连接网络进行二分类。

在开始前，我们先准备一下词典，这里就简单的弄几个词：

sentence = "大家好，我是练习时长两年半的个人练习生蔡徐坤，喜欢唱跳RAP篮球，接下来我会为大家带来一首鸡你太美。"
1

vocab = build_vocab_from_iterator(sentence, specials=['[PAD]', '[CLS]', '[SEP]', '[MASK]'])
1

接下来定义数据集，这里我只定义两条数据，仅仅是为了模拟BERT训练过程：

# BERT的输入，以[CLS]开头，两句话中间以[SEP]分割，长度都为24。
inputs = [
    '[CLS] 我 是 [MASK] 习 时 长 两 年 半 [SEP] 的 个 人 练 习 生 [MASK] 徐 坤 [SEP] [PAD] [PAD] [PAD]',
    '[CLS] 喜 欢 [MASK] 跳 R A P 篮 [MASK] [SEP] 鸡 你 太 [MASK] [SEP] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD]'
]

# 段落信息，表示该token属于哪句话
segment_label = [
    [1,1,1,1,1,1,1,1,1,1,1,2,2,2,2,2,2,2,2,2,2,0,0,0],
    [1,1,1,1,1,1,1,1,1,1,1,2,2,2,2,2,0,0,0,0,0,0,0,0]
]

# 定义MLM任务的targets，对于[CLS]和[SEP]不需要预测，所以标签中用[PAD]代替
mlm_targets = [
    '[PAD] 我 是 练 习 时 长 两 年 半 [PAD] 的 个 人 练 习 生 蔡 徐 坤 [PAD] [PAD] [PAD] [PAD]',
    '[PAD] 喜 欢 唱 跳 R A P 篮 球 [PAD] 鸡 你 太 美 [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD]'
]

# 定义NSP任务的targets
nsp_targets = [1, 0]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

接下来将上面的文本数据变成token，并将其变成tensor

inputs = torch.LongTensor([vocab(input.split(' ')) for input in inputs])
segment_label = torch.LongTensor(segment_label)
mlm_targets = torch.LongTensor([vocab(target.split(' ')) for target in mlm_targets])
nsp_targets = torch.LongTensor(nsp_targets)
print("inputs.shape:", inputs.size())
print("targets.shape:", mlm_targets.size())
1
2
3
4
5
6

inputs.shape: torch.Size([2, 24])
targets.shape: torch.Size([2, 24])
1
2

定义好数据集后，我们需要定义出针对MLM和NSP任务的网络，其实就是后面再接一个全连接层就行了：

class MaskedLanguageModel(nn.Module):
    """
    predicting origin token from masked input sequence
    n-class classification problem, n-class = vocab_size
    """

    def __init__(self, hidden, vocab_size):
        """
        :param hidden: output size of BERT model
        :param vocab_size: total vocab size
        """
        super().__init__()
        self.linear = nn.Linear(hidden, vocab_size)
        self.softmax = nn.LogSoftmax(dim=-1)

    def forward(self, x):
        return self.softmax(self.linear(x))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

class NextSentencePrediction(nn.Module):
    """
    2-class classification model : is_next, is_not_next
    """

    def __init__(self, hidden):
        """
        :param hidden: BERT model output size
        """
        super().__init__()
        self.linear = nn.Linear(hidden, 2)
        self.softmax = nn.LogSoftmax(dim=-1)

    def forward(self, x):
        return self.softmax(self.linear(x[:, 0]))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

class BERTLM(nn.Module):

    def __init__(self, vocab_size):
        super(BERTLM, self).__init__()
        self.vocab_size = vocab_size

        # 这里就使用bert_base吧
        self.bert = BERT(vocab_size=vocab_size, hidden=768, n_layers=12, attn_heads=12)

        self.next_sentence = NextSentencePrediction(self.bert.hidden)
        self.mask_lm = MaskedLanguageModel(self.bert.hidden, vocab_size)

    def forward(self, x, segment_label):
        x = self.bert(x, segment_label)
        return self.next_sentence(x), self.mask_lm(x)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

bert_mlm = BERTLM(len(vocab))
nsp_outputs, mlm_outputs = bert_mlm(inputs, segment_label)
print("nsp_outputs shape:", nsp_outputs.size())
print("mlm_outputs shape:", mlm_outputs.size())
1
2
3
4

nsp_outputs shape: torch.Size([2, 2])
mlm_outputs shape: torch.Size([2, 24, 46])
1
2

接下来开始训练网络：

criterion = nn.NLLLoss(ignore_index = 0)
optimizer = torch.optim.Adam(bert_mlm.parameters(), lr=3e-5)
1
2

for epoch in range(300):
    nsp_outputs, mlm_outputs = bert_mlm(inputs, (inputs>0).int())
    nsp_loss = criterion(nsp_outputs, nsp_targets)
    mlm_loss = criterion(mlm_outputs.view(-1, 46), mlm_targets.view(-1))
    loss = nsp_loss + mlm_loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

    print("loss {:.4}, nsp loss: {:.4}, mlm_loss {:.4}".format(loss, nsp_loss, mlm_loss))
1
2
3
4
5
6
7
8
9
10

loss 4.49, nsp loss: 0.4757, mlm_loss 4.014
loss 3.851, nsp loss: 0.007476, mlm_loss 3.843
loss 3.67, nsp loss: 0.002525, mlm_loss 3.667
...
loss 0.004217, nsp loss: 4.911e-05, mlm_loss 0.004168
loss 0.004457, nsp loss: 8.523e-05, mlm_loss 0.004372
1
2
3
4
5
6

这里虽然只有一个样本，但收敛的并不是很快，因为网络太深了，如果你把BERT改浅一点，就会发现收敛特别快。

训练好之后，我们来试一下:

inputs = '我 是 练 习 时 长 两 年 半 的 个 人 练 习 生 [MASK] 徐 坤'
inputs = torch.LongTensor([vocab(inputs.split(' '))])
segment_label = torch.ones(inputs.size()).long()
1
2
3

nsp_outputs, mlm_outputs = bert_mlm(inputs, segment_label)
1

print(vocab.lookup_tokens(mlm_outputs.argmax(-1)[0].tolist()))
1

['我', '是', '练', '习', '时', '长', '两', '年', '半', '的', '个', '人', '练', '习', '生', '蔡', '徐', '坤']
1

参考资料

(原论文)BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding: https://arxiv.org/abs/1810.04805

(论文阅读)BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding: https://blog.csdn.net/zhaohongfei_358/article/details/126838417

Pytorch中 nn.Transformer的使用详解与Transformer的黑盒讲解: https://blog.csdn.net/zhaohongfei_358/article/details/126019181

万字逐行解析与实现Transformer，并进行德译英实战: https://blog.csdn.net/zhaohongfei_358/article/details/126085246

BERT-Pytorch实现: https://github.com/codertimo/BERT-pytorch