ROCm上来自Transformers的双向编码器表示（BERT）

14.8. 来自Transformers的双向编码器表示（BERT） — 动手学深度学习 2.0.0 documentation (d2l.ai)

代码

import torch
from torch import nn
from d2l import torch as d2l#@save
def get_tokens_and_segments(tokens_a, tokens_b=None):"""获取输入序列的词元及其片段索引"""tokens = ['<cls>'] + tokens_a + ['<sep>']# 0和1分别标记片段A和Bsegments = [0] * (len(tokens_a) + 2)if tokens_b is not None:tokens += tokens_b + ['<sep>']segments += [1] * (len(tokens_b) + 1)return tokens, segments#@save
class BERTEncoder(nn.Module):"""BERT编码器"""def __init__(self, vocab_size, num_hiddens, norm_shape, ffn_num_input,ffn_num_hiddens, num_heads, num_layers, dropout,max_len=1000, key_size=768, query_size=768, value_size=768,**kwargs):super(BERTEncoder, self).__init__(**kwargs)self.token_embedding = nn.Embedding(vocab_size, num_hiddens)self.segment_embedding = nn.Embedding(2, num_hiddens)self.blks = nn.Sequential()for i in range(num_layers):self.blks.add_module(f"{i}", d2l.EncoderBlock(key_size, query_size, value_size, num_hiddens, norm_shape,ffn_num_input, ffn_num_hiddens, num_heads, dropout, True))# 在BERT中，位置嵌入是可学习的，因此我们创建一个足够长的位置嵌入参数self.pos_embedding = nn.Parameter(torch.randn(1, max_len,num_hiddens))def forward(self, tokens, segments, valid_lens):# 在以下代码段中，X的形状保持不变：（批量大小，最大序列长度，num_hiddens）X = self.token_embedding(tokens) + self.segment_embedding(segments)X = X + self.pos_embedding.data[:, :X.shape[1], :]for blk in self.blks:X = blk(X, valid_lens)return Xvocab_size, num_hiddens, ffn_num_hiddens, num_heads = 10000, 768, 1024, 4
norm_shape, ffn_num_input, num_layers, dropout = [768], 768, 2, 0.2
encoder = BERTEncoder(vocab_size, num_hiddens, norm_shape, ffn_num_input,ffn_num_hiddens, num_heads, num_layers, dropout)tokens = torch.randint(0, vocab_size, (2, 8))
segments = torch.tensor([[0, 0, 0, 0, 1, 1, 1, 1], [0, 0, 0, 1, 1, 1, 1, 1]])
encoded_X = encoder(tokens, segments, None)
encoded_X.shapetokens = torch.randint(0, vocab_size, (2, 8))
segments = torch.tensor([[0, 0, 0, 0, 1, 1, 1, 1], [0, 0, 0, 1, 1, 1, 1, 1]])
encoded_X = encoder(tokens, segments, None)
encoded_X.shapemlm = MaskLM(vocab_size, num_hiddens)
mlm_positions = torch.tensor([[1, 5, 2], [6, 1, 5]])
mlm_Y_hat = mlm(encoded_X, mlm_positions)
mlm_Y_hat.shapemlm_Y = torch.tensor([[7, 8, 9], [10, 20, 30]])
loss = nn.CrossEntropyLoss(reduction='none')
mlm_l = loss(mlm_Y_hat.reshape((-1, vocab_size)), mlm_Y.reshape(-1))
mlm_l.shape#@save
class NextSentencePred(nn.Module):"""BERT的下一句预测任务"""def __init__(self, num_inputs, **kwargs):super(NextSentencePred, self).__init__(**kwargs)self.output = nn.Linear(num_inputs, 2)def forward(self, X):# X的形状：(batchsize,num_hiddens)return self.output(X)encoded_X = torch.flatten(encoded_X, start_dim=1)
# NSP的输入形状:(batchsize，num_hiddens)
nsp = NextSentencePred(encoded_X.shape[-1])
nsp_Y_hat = nsp(encoded_X)
nsp_Y_hat.shapensp_y = torch.tensor([0, 1])
nsp_l = loss(nsp_Y_hat, nsp_y)
nsp_l.shape#@save
class BERTModel(nn.Module):"""BERT模型"""def __init__(self, vocab_size, num_hiddens, norm_shape, ffn_num_input,ffn_num_hiddens, num_heads, num_layers, dropout,max_len=1000, key_size=768, query_size=768, value_size=768,hid_in_features=768, mlm_in_features=768,nsp_in_features=768):super(BERTModel, self).__init__()self.encoder = BERTEncoder(vocab_size, num_hiddens, norm_shape,ffn_num_input, ffn_num_hiddens, num_heads, num_layers,dropout, max_len=max_len, key_size=key_size,query_size=query_size, value_size=value_size)self.hidden = nn.Sequential(nn.Linear(hid_in_features, num_hiddens),nn.Tanh())self.mlm = MaskLM(vocab_size, num_hiddens, mlm_in_features)self.nsp = NextSentencePred(nsp_in_features)def forward(self, tokens, segments, valid_lens=None,pred_positions=None):encoded_X = self.encoder(tokens, segments, valid_lens)if pred_positions is not None:mlm_Y_hat = self.mlm(encoded_X, pred_positions)else:mlm_Y_hat = None# 用于下一句预测的多层感知机分类器的隐藏层，0是“<cls>”标记的索引nsp_Y_hat = self.nsp(self.hidden(encoded_X[:, 0, :]))return encoded_X, mlm_Y_hat, nsp_Y_hat

代码解析

这段代码是基于PyTorch框架实现的BERT（Bidirectional Encoder Representations from Transformers）模型。BERT是一种预训练语言表示模型，它可以用于各种自然语言处理（NLP）任务。下面是代码的中文解析：
1. get_tokens_and_segments(tokens_a, tokens_b=None) 函数用于获取输入句子的词元（tokens）及其对应的片段索引。如果有第二个句子 tokens_b，则会进行拼接，并用不同的索引来标识不同的句子。
2. BERTEncoder 类定义了BERT的编码器结构，它包含嵌入层（用于将词元转换为向量表示）、位置嵌入和多个Transformer编码块。
3. forward 方法定义了模型的前向传播逻辑。它将输入的词元和片段索引通过编码器进行编码，并返回编码后的向量表示。
4. 其中 tokens 是批量输入数据的词元索引，`segments` 是对应的片段索引，这里模拟了输入数据作为模型的示例。
5. 创建一个 BERTEncoder 实例，该实例就是BERT模型的编码器部分，类似于 Transformer 模型中的编码器层。
6. MaskLM 类未在代码中定义，通常用来实现BERT的掩码语言模型任务，它在一定比例的输入词元上应用掩码，并训练模型来预测这些被掩码的词元。
7. NextSentencePred 类定义了BERT的下一句预测（Next Sentence Prediction, NSP）任务，是一个简单的二分类器，用来预测给定的两个句子片段是否在原始文本中顺序相邻。
8. BERTModel 类将编码器、掩码语言模型（MaskLM），以及下一句预测（NSP）整合为完整的BERT模型。它通过前向传播来处理输入，同时能够根据需求进行掩码语言模型预测和下一句预测。
9. 模型实例化后，通过随机生成的 tokens 和 segments 调用其 forward 方法，得到编码后的向量 encoded_X，同时执行MLM和NSP任务，输出预测结果。
10. 最后计算MLM和NSP任务的损失，这些损失通常用于训练模型。`CrossEntropyLoss` 是在类别预测问题中经常使用的一个损失函数。
整体来看，这段代码展示了如何构建一个基于BERT结构的模型，其中涵盖了BERT的两个典型预训练任务：掩码语言模型和下一句预测。需要注意的是，这个代码片段作为一个解析，但实际中运行它需要额外的上下文（例如 MaskLM 类的实现）和适当的数据准备和预处理步骤。