在本篇文章中，我们将介绍并实现BERT（Bidirectional Encoder Representations from Transformers）。与传统的Transformer模型相比，BERT的主要区别在于如何处理数据，尤其是引入了掩码语言建模（Masked Language Modeling）和下一句预测（Next Sentence Prediction）。在这篇文章中，我们将从零开始构建BERT，并演示其训练和推理过程。
1. 数据预处理
我们将使用一个简单的语料库，首先对文本进行预处理，包括将文本转换为小写、去掉标点符号等。
import spacy
# 加载简单的文本数据
with open("data/wiki_king.txt", "r") as f:
raw_text = f.read()
nlp = spacy.load("en_core_web_sm")
doc = nlp(raw_text)
sentences = list(doc.sents)
# 处理文本：转为小写并去除标点符号
text = [x.text.lower() for x in sentences]
text = [re.sub("[.,!?\\-]", '', x) for x in text]
print(text)
我们可以看到，所有的句子都已经被处理成小写，并且去除了标点符号。
2. 词汇表生成
在生成词汇表之前，我们首先需要将文本拆分为单词。我们为每个单词生成唯一的ID，并为特殊标记（如[PAD]、[CLS]、[SEP]、[MASK]）预留ID。
# 生成词汇表
word_list = list(set(" ".join(text).split()))
word2id = {'[PAD]': 0, '[CLS]': 1, '[SEP]': 2, '[MASK]': 3} # 特殊标记
# 为每个单词生成ID
for i, w in enumerate(word_list):
word2id[w] = i + 4 # 特殊标记占用了0-3
id2word = {i: w for w, i in word2id.items()}
vocab_size = len(word2id)
# 将文本转换为ID序列
token_list = [[word2id[word] for word in sentence.split()] for sentence in text]
print(token_list)
3. 数据加载器
BERT模型需要处理两种嵌入：Token嵌入和Segment嵌入，并且要对输入句子进行随机掩码处理。我们将实现一个生成批处理数据的函数，该函数包含以下步骤：
- Token嵌入：在句子开头添加
[CLS]标记，两个句子之间添加[SEP]标记。 - Segment嵌入：用0和1来区分两个句子。
- 掩码语言建模：随机掩盖15%的单词，其中80%替换为
[MASK]标记。 - 填充：将所有序列填充到相同长度。
batch_size = 6
max_mask = 5
max_len = 1000
def make_batch():
batch = []
positive = negative = 0
while positive != batch_size/2 or negative != batch_size/2:
tokens_a_index, tokens_b_index = randrange(len(sentences)), randrange(len(sentences))
tokens_a, tokens_b = token_list[tokens_a_index], token_list[tokens_b_index]
# 1. Token嵌入
input_ids = [word2id['[CLS]']] + tokens_a + [word2id['[SEP]']] + tokens_b + [word2id['[SEP]']]
# 2. Segment嵌入
segment_ids = [0] * (1 + len(tokens_a) + 1) + [1] * (len(tokens_b) + 1)
# 3. 掩码语言建模
n_pred = min(max_mask, max(1, int(len(input_ids) * 0.15)))
cand_maked_pos = [i for i, token in enumerate(input_ids) if token not in (word2id['[CLS]'], word2id['[SEP]'])]
shuffle(cand_maked_pos)
masked_tokens, masked_pos = [], []
for pos in cand_maked_pos[:n_pred]:
masked_pos.append(pos)
masked_tokens.append(input_ids[pos])
if random() < 0.1:
input_ids[pos] = randint(0, vocab_size - 1)
elif random() < 0.9:
input_ids[pos] = word2id['[MASK]']
# 4. 填充
n_pad = max_len - len(input_ids)
input_ids.extend([0] * n_pad)
segment_ids.extend([0] * n_pad)
if max_mask > n_pred:
n_pad = max_mask - n_pred
masked_tokens.extend([0] * n_pad)
masked_pos.extend([0] * n_pad)
if tokens_a_index + 1 == tokens_b_index and positive < batch_size / 2:
batch.append([input_ids, segment_ids, masked_tokens, masked_pos, True])
positive += 1
elif tokens_a_index + 1 != tokens_b_index and negative < batch_size / 2:
batch.append([input_ids, segment_ids, masked_tokens, masked_pos, False])
negative += 1
return batch
batch = make_batch()
input_ids, segment_ids, masked_tokens, masked_pos, isNext = map(torch.LongTensor, zip(*batch))
print(input_ids.shape, segment_ids.shape, masked_tokens.shape, masked_pos.shape, isNext.shape)
4. BERT模型实现
4.1 嵌入层
在BERT模型中，嵌入层负责将输入的Token、位置和句子段嵌入整合在一起。我们使用LayerNorm来标准化输出。
class Embedding(nn.Module):
def __init__(self):
super(Embedding, self).__init__()
self.tok_embed = nn.Embedding(vocab_size, d_model)
self.pos_embed = nn.Embedding(max_len, d_model)
self.seg_embed = nn.Embedding(n_segments, d_model)
self.norm = nn.LayerNorm(d_model)
def forward(self, x, seg):
pos = torch.arange(x.size(1), dtype=torch.long).unsqueeze(0).expand_as(x)
embedding = self.tok_embed(x) + self.pos_embed(pos) + self.seg_embed(seg)
return self.norm(embedding)
4.2 多头注意力机制
BERT的多头注意力机制允许模型在不同的注意力头上并行关注不同部分的输入。
class MultiHeadAttention(nn.Module):
def __init__(self):
super(MultiHeadAttention, self).__init__()
self.W_Q = nn.Linear(d_model, d_k * n_heads)
self.W_K = nn.Linear(d_model, d_k * n_heads)
self.W_V = nn.Linear(d_model, d_v * n_heads)
def forward(self, Q, K, V, attn_mask):
q_s = self.W_Q(Q).view(Q.size(0), -1, n_heads, d_k).transpose(1, 2)
k_s = self.W_K(K).view(K.size(0), -1, n_heads, d_k).transpose(1, 2)
v_s = self.W_V(V).view(V.size(0), -1, n_heads, d_v).transpose(1, 2)
attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)
context, attn = ScaledDotProductAttention()(q_s, k_s, v_s, attn_mask)
context = context.transpose(1, 2).contiguous().view(Q.size(0), -1, n_heads * d_v)
return nn.LayerNorm(d_model)(nn.Linear(n_heads * d_v, d_model)(context) + Q), attn
4.3 BERT模型
最后我们定义完整的BERT模型，其中包括嵌入层、编码层以及用于掩码语言模型和下一句预测的分类器。
class BERT(nn.Module):
def __init__(self):
super(BERT, self).__init__()
self.embedding = Embedding()
self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])
self.fc = nn.Linear(d_model, d_model)
self.activ = nn.Tanh()
self.classifier = nn.Linear(d_model, 2)
embed_weight = self.embedding.tok_embed.weight
self.decoder = nn.Linear(embed_weight.size(1), vocab_size, bias=False)
self.decoder.weight = embed_weight
self.decoder_bias = nn.Parameter(torch.zeros(vocab_size))
def forward(self, input_ids, segment_ids, masked_pos):
output = self.embedding(input_ids, segment_ids)
enc_self_attn_mask = get_attn_pad_mask(input_ids, input_ids)
for layer in self.layers:
output, _ = layer(output, enc_self_attn_mask)
h_pooled = self.activ(self.fc(output[:, 0]))
logits_nsp = self.classifier(h_pooled)
masked_pos = masked_pos[:, :, None].expand(-1, -1, output.size(-1))
h_masked = torch.gather(output, 1, masked_pos)
logits_lm = self.decoder(h_masked) + self.decoder_bias
return logits_lm, logits_nsp
5. 模型训练
我们将BERT模型进行训练，使用交叉熵损失函数对掩码语言建模和下一句预测进行优化。
num_epoch = 500
model = BERT()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
for epoch in range(num_epoch):
optimizer.zero_grad()
logits_lm, logits_nsp = model(input_ids, segment_ids, masked_pos)
loss_lm = criterion(logits_lm.transpose(1, 2), masked_tokens).mean()
loss_nsp = criterion(logits_nsp, isNext)
loss = loss_lm + loss_nsp
if epoch % 100 == 0:
print(f'Epoch: {epoch}, Loss: {loss.item():.6f}')
loss.backward()
optimizer.step()
6. 推理
最后，我们演示如何使用训练好的BERT模型进行推理。
logits_lm, logits_nsp = model(input_ids, segment_ids, masked_pos)
logits_lm = logits_lm.data.max(2)[1][0].data.numpy()
print('预测的掩码单词:', [id2word[pos] for pos in logits_lm])
logits_nsp = logits_nsp.data.max(1)[1][0].data.numpy()
print('是否为下一句:', '是' if logits_nsp else '否')
结语
在本篇文章中，我们详细探讨了BERT模型的构建与实现，尤其是在数据处理阶段的独特之处，如掩码语言模型（Masked Language Model, MLM）和下一句预测（Next Sentence Prediction, NSP）的结合。通过逐步解析BERT的编码器结构、注意力机制以及多层自注意力网络，我们展示了如何在实际项目中应用BERT模型进行预训练任务。
BERT的强大在于其双向编码的特性，使得模型能够充分利用上下文信息，特别适合解决文本分类、句子配对等复杂的自然语言理解任务。在这个实现中，我们使用了一个简单的数据集来演示BERT的核心功能。虽然实际性能受到数据量的限制，但通过扩展到更大的数据集和更复杂的任务，BERT的能力会得到充分发挥。
在下一篇文章中，我们将探讨Pruning技术，它是一种模型压缩方法，可以通过减少模型中的冗余参数来提高计算效率和内存使用率。这对于部署深度学习模型至关重要，尤其是在资源有限的设备上。敬请期待！
如果你觉得这篇博文对你有帮助，请点赞、收藏、关注我，并且可以打赏支持我！
欢迎关注我的后续博文，我将分享更多关于人工智能、自然语言处理和计算机视觉的精彩内容。
谢谢大家的支持！
扫一扫
3158
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
