BERT源码实现与解读(Pytorch)本文内容环境准备BERT模型定义预训练BERT参考资料

63 阅读 0 评论 42 点赞

我是靠谱客的博主唠叨大侠，最近开发中收集的这篇文章主要介绍BERT源码实现与解读(Pytorch)本文内容环境准备BERT模型定义预训练BERT参考资料，觉得挺不错的，现在分享给大家，希望可以做个参考。

概述

文章目录

本文内容
环境准备
BERT模型定义
- BERT Embedding
- BERT
预训练BERT
参考资料

本文源码地址: 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

torch.__version__

'1.12.1+cpu'

torchtext.__version__

'0.13.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)

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)]

Position Embedding和Transformer也一致。

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

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)

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，
# 例如句子：<bos>唱跳Rap篮球<eos><pad><pad>
# 对应的key_padding_mask为：[T, T, T, T, T, T, False, False]
# 最后两个False是对<pad>进行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

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

$bf{BERT_{BASE}}$ ：12层，hidden size为768，Attention Head数为12，共110M个参数。
$bf{BERT_{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()]))

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

可以看到，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())

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

预训练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篮球，接下来我会为大家带来一首鸡你太美。"

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

接下来定义数据集，这里我只定义两条数据，仅仅是为了模拟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]

接下来将上面的文本数据变成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())

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

定义好数据集后，我们需要定义出针对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))

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]))

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)

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())

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

接下来开始训练网络：

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

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))

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

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

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

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

nsp_outputs, mlm_outputs = bert_mlm(inputs, segment_label)

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

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

参考资料

(原论文)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