BERT源码解析-#notebook-code

回顾:BERT结构

BERT模型结构基本上就是TransformerEncoder部分,BERT-base对应的是12层encoder,BERT-large对应的是24层encoder.

BERT模型结构

在这里插入图片描述

  1. BERT Tokenization 分词模型(BertTokenizer)
  2. BERT Model 本体模型(BertModel)
  • BertEmbeddings
  • BertEncoder
    • BertLayer
      • BertAttention
      • BertIntermediate
      • BertOutput
  • BertPooler

BertTokenizer

Tokenizer是在自然语言处理(NLP)中的一个关键组件,它负责将文本转换成一种格式,以便机器学习模型能够理解和处理
作用直白地说就是
我们可以把Tokenizer比作是将一长串文字“切割”成有意义的小块(比如单词或字符)的工具。

Tokenizer的输入和输出

  • 输入:Tokenizer的输入通常是原始文本。这可以是一句话、一个段落或者一个完整的文档。输入文本通常包含了自然语言的所有复杂性,如不同的语言、方言、专业术语、俚语等。

  • 输出:Tokenizer的输出是一系列tokens,这些tokens已经被转换成了一种结构化的格式。输出的形式依赖于

Tokenizer的设计和目的,常见的有:

  • Token序列:一系列分割后的文本单元(如单词或字符)。
  • 数字ID序列:如果Tokenizer包括编码过程,每个token会被映射到一个唯一的数字ID,这些ID对应于模型中的词汇表。
  • 向量序列:在某些高级应用中,每个token可能直接被转换为一个稠密向量,这通常发生在使用预训练的嵌入模型时。
class BertTokenizer(PreTrainedTokenizer):
    """
    Construct a BERT tokenizer. Based on WordPiece.

    This tokenizer inherits from :class:`~transformers.PreTrainedTokenizer` which contains most of the main methods.
    Users should refer to this superclass for more information regarding those methods.
    ...
    """

BertTokenizer 是基于BasicTokenizerWordPieceTokenizer的分词器:

  • BasicTokenizer负责处理的第一步——按标点、空格等分割句子,并处理是否统一小写,以及清理非法字符
    • 对于中文字符,通过预处理(加空格)来按字分割;
    • 同时可以通过never_split指定对某些词不进行分割;
    • 这一步是可选的(默认执行)。
  • WordPieceTokenizer在词的基础上,进一步将词分解为子词(subword)。
    • subword 介于 char 和 word 之间,既在一定程度保留了词的含义,又能够照顾到英文中单复数、时态导致的词表爆炸和未登录词的 OOV(Out-Of-Vocabulary)问题,将词根与时态词缀等分割出来,从而减小词表,也降低了训练难度;
    • 例如,tokenizer 这个词就可以拆解为“token”和“##izer”两部分,注意后面一个词的“##”表示接在前一个词后面。

BertTokenizer 有以下常用方法

  • from_pretrained:从包含词表文件(vocab.txt)的目录中初始化一个分词器;
  • tokenize:将文本(词或者句子)分解为子词列表;
  • convert_tokens_to_ids:将子词列表转化为子词对应下标的列表;
  • convert_ids_to_tokens :与上一个相反;
  • convert_tokens_to_string:将 subword 列表按“##”拼接回词或者句子;
  • encode:对于单个句子输入,分解词并加入特殊词形成“[CLS], x, [SEP]”的结构并转换为词表对应下标的列表;对于两个句子输入(多个句子只取前两个),分解词并加入特殊词形成“[CLS], x1, [SEP], x2, [SEP]”的结构并转换为下标列表;
  • decode:可以将 encode 方法的输出变为完整句子。
    以及,类自身的方法:…
from transformers import BertTokenizer


tokenizer_dir = './pretrained_bert_models/bert_base_uncased/vocab.txt'
tokenizer = BertTokenizer(tokenizer_dir)
# 定义输入文本
text = "Hello, world! This is a test for the Tokenizer."
# 使用Tokenizer
encoded_input = tokenizer(text)
# 打印输出
print("原始文本:", text)
print("输出的内容:",encoded_input)
print("Tokenized 输出:", tokenizer.convert_ids_to_tokens(encoded_input['input_ids']))
print("数字ID序列:", encoded_input['input_ids'])
# 原始文本: Hello, world! This is a test for the Tokenizer.
# 输出的内容: {'input_ids': [101, 7592, 1010, 2088, 999, 2023, 2003, 1037, 3231, 2005, 1996, 19204, 17629, 1012, 102],
#             'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#             'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}
# Tokenized 输出: ['[CLS]', 'hello', ',', 'world', '!', 'this', 'is', 'a', 'test', 'for', 'the', 'token', '##izer', '.', '[SEP]']
# 数字ID序列: [101, 7592, 1010, 2088, 999, 2023, 2003, 1037, 3231, 2005, 1996, 19204, 17629, 1012, 102]
  • input_ids是模型理解文本的数字表示,列表中的每个整数代表词汇表中相应token的索引。例如,BERT模型中101102分别标识句子开始和结束的特殊token。
  • token_type_ids用于区分输入序列的不同部分,如BERT中使用01来标记句子对中的第一个和第二个句子。
  • attention_mask是一个与input_ids长度相同的数组,其值为1的元素表示对应token是有效且应被模型关注的,而0则表示该token是填充或不重要的。这对于处理不同长度的输入序列至关重要。

BertModel

和 BERT 模型有关的代码主要写在/models/bert/modeling_bert.py中,这一份代码有一千多行,包含 BERT 模型的基本结构和基于它的微调模型等。

下面从 BERT 模型本体入手分析:

class BertModel(BertPreTrainedModel):
    """

    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.

    To behave as an decoder the model needs to be initialized with the :obj:`is_decoder` argument of the configuration
    set to :obj:`True`. To be used in a Seq2Seq model, the model needs to initialized with both :obj:`is_decoder`
    argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
    input to the forward pass.
    """ 

BertModel 主要为 transformer encoder 结构,包含三个部分:

  1. embeddings,即BertEmbeddings类的实体,根据单词符号获取对应的向量表示;
  2. encoder,即BertEncoder类的实体;
  3. pooler,即BertPooler类的实体,这一部分是可选的。

注意 BertModel 也可以配置为 Decoder,不过下文中不包含对这一部分的讨论。

下面将介绍 BertModel 的前向传播过程中各个参数的含义以及返回值:

def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ): ...
  • input_ids:经过 tokenizer 分词后的 subword 对应的下标列表;
  • attention_mask:在 self-attention 过程中,这一块 mask 用于标记 subword 所处句子和 padding 的区别,将 padding 部分填充为 0;
  • token_type_ids:标记 subword 当前所处句子(第一句/第二句/ padding);
  • position_ids:标记当前词所在句子的位置下标;
  • head_mask:用于将某些层的某些注意力计算无效化;
  • inputs_embeds:如果提供了,那就不需要input_ids,跨过 embedding lookup 过程直接作为 Embedding 进入 Encoder 计算;
  • encoder_hidden_states:这一部分在 BertModel 配置为 decoder 时起作用,将执行 cross-attention 而不是 self-attention;
  • encoder_attention_mask:同上,在 cross-attention 中用于标记 encoder 端输入的 padding;
  • past_key_values:这个参数貌似是把预先计算好的 K-V 乘积传入,以降低 cross-attention 的开销(因为原本这部分是重复计算);
  • use_cache:将保存上一个参数并传回,加速 decoding;
  • output_attentions:是否返回中间每层的 attention 输出;
  • output_hidden_states:是否返回中间每层的输出;
  • return_dict:是否按键值对的形式(ModelOutput 类,也可以当作 tuple 用)返回输出,默认为真。

注意,这里的 head_mask 对注意力计算的无效化,和下文提到的注意力头剪枝不同,而仅仅把某些注意力的计算结果给乘以这一系数。

输出部分如下:

# BertModel的前向传播返回部分
if not return_dict:
    return (sequence_output, pooled_output) + encoder_outputs[1:]

return BaseModelOutputWithPoolingAndCrossAttentions(
    last_hidden_state=sequence_output,
    pooler_output=pooled_output,
    past_key_values=encoder_outputs.past_key_values,
    hidden_states=encoder_outputs.hidden_states,
    attentions=encoder_outputs.attentions,
    cross_attentions=encoder_outputs.cross_attentions,
)

可以看出,返回值不但包含了 encoder 和 pooler 的输出,也包含了其他指定输出的部分(hidden_states 和 attention 等,这一部分在encoder_outputs[1:])方便取用:

# BertEncoder的前向传播返回部分,即上面的encoder_outputs
if not return_dict:
    return tuple(
        v
        for v in [
            hidden_states,
            next_decoder_cache,
            all_hidden_states,
            all_self_attentions,
            all_cross_attentions,
        ]
        if v is not None
    )
return BaseModelOutputWithPastAndCrossAttentions(
    last_hidden_state=hidden_states,
    past_key_values=next_decoder_cache,
    hidden_states=all_hidden_states,
    attentions=all_self_attentions,
    cross_attentions=all_cross_attentions,
)

此外,BertModel 还有以下的方法,方便 BERT 玩家进行各种操作:

  • get_input_embeddings:提取 embedding 中的 word_embeddings 即词向量部分;
  • set_input_embeddings:为 embedding 中的 word_embeddings 赋值;
  • _prune_heads:提供了将注意力头剪枝的函数,输入为{layer_num: list of heads to prune in this layer}的字典,可以将指定层的某些注意力头剪枝。

剪枝是一个复杂的操作,需要将保留的注意力头部分的 Wq、Kq、Vq 和拼接后全连接部分的权重拷贝到一个新的较小的权重矩阵(注意先禁止 grad 再拷贝),并实时记录被剪掉的头以防下标出错。具体参考BertAttention部分的prune_heads方法.

BertEmbeddings


包含三个部分求和得到:
 图:Bert-embedding

  1. word_embeddings,上文中 subword 对应的嵌入。
  2. token_type_embeddings,用于表示当前词所在的句子,辅助区别句子与 padding、句子对间的差异。
  3. position_embeddings,句子中每个词的位置嵌入,用于区别词的顺序。和 transformer 论文中的设计不同,这一块是训练出来的,而不是通过 Sinusoidal 函数计算得到的固定嵌入。一般认为这种实现不利于拓展性(难以直接迁移到更长的句子中)。

三个 embedding 不带权重相加,并通过一层 LayerNorm+dropout 后输出,其大小为(batch_size, sequence_length, hidden_size)。

** 这里为什么要用 LayerNorm+Dropout 呢?为什么要用 LayerNorm 而不是 BatchNorm?可以参考一个不错的回答:transformer 为什么使用 layer normalization,而不是其他的归一化方法?**

要理解通过嵌入(embedding)层的输入和输出,我们可以使用一个预训练模型的嵌入层作为例子。这里,我们将继续使用Hugging Face的transformers库,并以BERT模型为例。嵌入层的主要作用是将输入的token ID转换成固定大小的向量,这些向量能够捕捉词汇的语义信息。

下面的代码将执行以下步骤:

  1. transformers库导入BERT模型和其Tokenizer。
  2. 初始化Tokenizer和模型。
  3. 定义一个文本输入,并用Tokenizer对其进行编码。
  4. 使用BERT模型的嵌入层对编码后的输入进行处理。
  5. 展示输入和输出。
from transformers import BertTokenizer, BertModel
import torch

# 初始化Tokenizer和模型
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

# 定义文本输入
text = "Here is some text to encode"
encoded_input = tokenizer(text, return_tensors='pt')

# 提取编码后的输入IDs
input_ids = encoded_input['input_ids']

# 使用模型的嵌入层
with torch.no_grad():
    outputs = model.embeddings(input_ids)

# 展示输入和嵌入层的输出
print("Input IDs:", input_ids)
print("Output of Embedding Layer (shape):", outputs.shape)

在这个例子中:

  • 输入(Input IDs)是文本经过Tokenizer处理后得到的input_ids,它是一个数字列表,每个数字代表文本中对应位置单词的唯一ID。
  • 输出(Output of Embedding Layer)是嵌入层处理后的结果,它是一个多维张量,其形状为(batch_size, sequence_length, hidden_size),其中:
    • batch_size是输入批次的大小(如果你只输入了一句话,这个值就是1)。
    • sequence_length是输入序列的长度。
    • hidden_size是嵌入向量的维度,对于bert-base-uncased模型,这个值是768,意味着每个token都被转换成了一个768维的向量。

这个多维张量捕捉了输入文本的丰富语义信息,每个单词(token)都通过嵌入向量在一个高维空间中被表示,这些向量将作为模型下游任务的输入。

如果你想要进一步地观察token_embeddings,segment_embeddings,postion_embeddings,你可以这样做:

# 提取段ID和位置ID
token_type_ids = encoded_input['token_type_ids']
position_ids = torch.arange(input_ids.size(1)).unsqueeze(0)

# 使用模型的嵌入层
token_embeddings = model.embeddings.word_embeddings(input_ids)
segment_embeddings = model.embeddings.token_type_embeddings(token_type_ids)
position_embeddings = model.embeddings.position_embeddings(position_ids)

# 打印嵌入的形状
print("Token Embeddings shape:", token_embeddings.shape)
print("Segment Embeddings shape:", segment_embeddings.shape)
print("Position Embeddings shape:", position_embeddings.shape)

这里:

  • Token Embeddings (word_embeddings):基于词汇表中每个token的ID,将每个token转换为一个固定大小的向量。
  • Segment Embeddings (token_type_embeddings):用于区分两个不同的句子或文本片段。在简单的单句子输入中,这个嵌入向量可能不会提供很多信息,但在处理成对的句子(如问答任务)时很有用。
  • Position Embeddings (position_embeddings):由于BERT使用的是Transformer架构,这个嵌入向量提供了每个token在句子中的位置信息,帮助模型理解单词顺序。

每种嵌入的形状都是(batch_size, sequence_length, hidden_size),其中hidden_size对于bert-base-uncased模型是768,表示每个嵌入向量的维度。

BertEncoder

包含多层 BertLayer,这一块本身没有特别需要说明的地方,不过有一个细节值得参考:利用 gradient checkpointing 技术以降低训练时的显存占用。

gradient checkpointing 即梯度检查点,通过减少保存的计算图节点压缩模型占用空间,但是在计算梯度的时候需要重新计算没有存储的值,参考论文《Training Deep Nets with Sublinear Memory Cost》,过程如下示意图

请添加图片描述

图:gradient-checkpointing

在 BertEncoder 中,gradient checkpoint 是通过 torch.utils.checkpoint.checkpoint 实现的,使用起来比较方便,可以参考文档:torch.utils.checkpoint - PyTorch 1.8.1 documentation,这一机制的具体实现比较复杂,在此不作展开。

再往深一层走,就进入了 Encoder 的某一层:

BertLayer

BertAttention

本以为 attention 的实现就在这里,没想到还要再下一层……其中,self 成员就是多头注意力的实现,而 output 成员实现 attention 后的全连接 +dropout+residual+LayerNorm 一系列操作。

class BertAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.self = BertSelfAttention(config)
        self.output = BertSelfOutput(config)
        self.pruned_heads = set()

首先还是回到这一层。这里出现了上文提到的剪枝操作,即 prune_heads 方法:

 def prune_heads(self, heads):
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
        )

        # Prune linear layers
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        # Update hyper params and store pruned heads
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads) 

这里的具体实现概括如下:

  • find_pruneable_heads_and_indices是定位需要剪掉的 head,以及需要保留的维度下标 index;

  • prune_linear_layer则负责将 Wk/Wq/Wv 权重矩阵(连同 bias)中按照 index 保留没有被剪枝的维度后转移到新的矩阵。
    接下来就到重头戏——Self-Attention 的具体实现。

BertSelfAttention

预警:这一块可以说是模型的核心区域,也是唯一涉及到公式的地方,所以将贴出大量代码。

初始化部分:

class BertSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
            )

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        self.key = nn.Linear(config.hidden_size, self.all_head_size)
        self.value = nn.Linear(config.hidden_size, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            self.max_position_embeddings = config.max_position_embeddings
            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)

        self.is_decoder = config.is_decoder
  • class BertSelfAttention(nn.Module): 定义了一个名为 BertSelfAttention 的类,这个类继承自 PyTorch 的 nn.Module。这表明 BertSelfAttention 是一个神经网络模块,可以集成到更大的神经网络架构中。

  • def __init__(self, config): 这是类的构造函数,用于初始化 BertSelfAttention 实例。它接收一个配置对象 config,这个对象包含了模型配置的各种参数。

  • super().__init__(): 调用父类 nn.Module 的构造函数,进行必要的初始化操作。

  • 检查 hidden_sizenum_attention_heads 的兼容性:这部分代码确保隐藏层的大小 (hidden_size) 是注意力头的数量 (num_attention_heads) 的倍数。这是因为在计算自注意力时,隐藏层的输出需要被平均分配到每个注意力头上。

  • self.num_attention_heads = config.num_attention_heads: 定义注意力头的数量。

  • self.attention_head_size = int(config.hidden_size / config.num_attention_heads): 计算每个注意力头的大小。

  • self.all_head_size = self.num_attention_heads * self.attention_head_size: 计算所有注意力头的总大小。

  • self.query, self.key, self.value: 分别定义了用于生成查询(query)、键(key)和值(value)向量的线性变换。这些向量是自注意力机制中的核心组成部分。

  • self.dropout = nn.Dropout(config.attention_probs_dropout_prob): 定义了一个 Dropout 层,用于在注意力概率上应用 dropout,以减少过拟合。

  • self.position_embedding_type = getattr(config, "position_embedding_type", "absolute"): 获取位置嵌入类型,它决定了如何对位置信息进行编码。默认为 “absolute”。

  • if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query": 这段条件语句检查是否使用相对位置编码,如果是,则进行相应的初始化。

  • self.max_position_embeddings = config.max_position_embeddings: 设置最大位置嵌入数,这是相对位置编码所需的。

  • self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size): 初始化用于相对位置编码的嵌入层。

  • self.is_decoder = config.is_decoder: 指示这个模块是否被用作解码器的一部分。在 Transformer 模型中,解码器用于生成输出序列,而编码器处理输入序列。

  • 注意,hidden_size 和 all_head_size 在一开始是一样的。至于为什么要看起来多此一举地设置这一个变量——显然是因为上面那个剪枝函数,剪掉几个 attention head 以后 all_head_size 自然就小了;

  • hidden_size 必须是 num_attention_heads 的整数倍,以 bert-base 为例,每个 attention 包含 12 个 head,hidden_size 是 768,所以每个 head 大小即 attention_head_size=768/12=64;

然后是重点,也就是前向传播过程。

首先回顾一下 multi-head self-attention 的基本公式:

D k D_k Dk表示key的维度; D m D_m Dm表示token维度,T表示序列长度.
多头自注意力(Multi-HeadSelf-Attention),在多个不同的投影空间中捕捉不同的交互信息.假设在 h h h 个投影空间中分别应用自注意力模型,有
MultiHead ⁡ ( Q , K , V ) = W o [ head ⁡ 1 ; ⋯   ; head ⁡ h ] , head ⁡ i = self-att ⁡ ( Q i , K i , V i ) , ∀ i ∈ { 1 , ⋯   , h } , Q i = W q i Q , K = W k i K , V = W v i V , \begin{aligned} \operatorname{MultiHead}(\mathbf{Q},\mathbf{K},\mathbf{V})&=\boldsymbol{W}_o[\operatorname{head}_1;\cdots;\operatorname{head}_h],\\ \operatorname{head}_i&=\operatorname{self-att}(\boldsymbol{Q}_i,\boldsymbol{K}_i,\boldsymbol{V}_i),\\ \forall i\in\{1,\cdots,h\},\quad \boldsymbol{Q}_i&=\boldsymbol{W}_q^i\boldsymbol{Q},\boldsymbol{K}=\boldsymbol{W}_k^i\boldsymbol{K},\boldsymbol{V}=\boldsymbol{W}_v^i\boldsymbol{V},\end{aligned} MultiHead(Q,K,V)headii{1,,h},Qi=Wo[head1;;headh],=self-att(Qi,Ki,Vi),=WqiQ,K=WkiK,V=WviV,
其中 W o ∈ R D m × h D v \boldsymbol{W}_o \in \mathbb{R}^{D_m \times h D_v} WoRDm×hD

上一篇:《量子十年》报告更新!IBM精研量子计算,助力行业优化转型


下一篇:C#,人工智能,机器学习,聚类算法,训练数据集生成算法、软件与源代码