?PaperWeekly 原創 ·?作者?|?李濼秋

學校?|?浙江大學碩士生

研究方向?|?自然語言處理、知識圖譜

本文記錄一下對 HuggingFace 開源的 Transformers 項目代碼的理解。

眾所周知，BERT 模型自 2018 年問世起就各種屠榜，開啟了 NLP 領域預訓練+微調的范式。到現在，BERT 的相關衍生模型層出不窮（XL-Net、RoBERTa、ALBERT、ELECTRA、ERNIE 等），要理解它們可以先從 BERT 這個始祖入手。

HuggingFace 是一家總部位于紐約的聊天機器人初創服務商，很早就捕捉到 BERT 大潮流的信號并著手實現基于 pytorch 的 BERT 模型。這一項目最初名為 pytorch-pretrained-bert，在復現了原始效果的同時，提供了易用的方法以方便在這一強大模型的基礎上進行各種玩耍和研究。

隨著使用人數的增加，這一項目也發展成為一個較大的開源社區，合并了各種預訓練語言模型以及增加了 Tensorflow 的實現，并且在 2019 年下半年改名為 Transformers。截止寫文章時（2021 年 3 月 30 日）這一項目已經擁有 43k+ 的star，可以說 Transformers 已經成為事實上的 NLP 基本工具。

詳見：transformers [1]?

本文基于 Transformers 版本 4.4.2（2021 年 3 月 19 日發布）項目中，pytorch 版的 BERT 相關代碼，從代碼結構、具體實現與原理，以及使用的角度進行分析，包含以下內容：

1. BERT Tokenization分詞模型（BertTokenizer）

2. BERT Model本體模型（BertModel）

2.1 BertEmbeddings

2.2 BertEncoder

2.2.1 BertLayer

2.2.1.1?BertAttention

2.2.1.1.1?BertSelfAttention

2.2.1.1.2 BertSelfOutput

2.2.1.2 BertIntermediate

2.2.1.3 BertOutput

2.2.3 BertPooler

3. BERT-based Models應用模型（請看下篇）

3.1 BertForPreTraining

3.2 BertForSequenceClassification

3.3 BertForMultiChoice

3.4 BertForTokenClassification

3.5 BertForQuestionAnswering

4. BERT訓練與優化（請看下篇）

4.1 Pre-Training

4.2 Fine-Tuning

4.2.1 AdamW

4.2.2 Warmup

Tokenization（BertTokenizer）

和 BERT 有關的 Tokenizer 主要寫在/models/bert/tokenization_bert.py和/models/bert/tokenization_bert_fast.py 中。

這兩份代碼分別對應基本的BertTokenizer，以及不進行 token 到 index 映射的BertTokenizerFast，這里主要講解第一個。

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 是基于BasicTokenizer和WordPieceTokenizer 的分詞器：

• 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 >>>?bt?=?BertTokenizer.from_pretrained('./bert-base-uncased/') >>>?bt('I?like?natural?language?progressing!') {'input_ids':?[101,?1045,?2066,?3019,?2653,?27673,?999,?102],?'token_type_ids':?[0,?0,?0,?0,?0,?0,?0,?0],?'attention_mask':?[1,?1,?1,?1,?1,?1,?1,?1]}

Model（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?ofcross-attention?is?added?between?the?self-attention?layers,?following?the?architecture?described?in?`Attention?isall?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?configurationset?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?aninput?to?the?forward?pass."""?

BertModel 主要為 transformer encoder 結構，包含三個部分：

embeddings，即BertEmbeddings類的實體，對應詞嵌入；

encoder，即BertEncoder類的實體；

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_outputsif?not?return_dict:return?tuple(vfor?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方法。

2.1 BertEmbeddings

包含三個部分求和得到：

word_embeddings，上文中 subword 對應的嵌入。

token_type_embeddings，用于表示當前詞所在的句子，輔助區別句子與 padding、句子對間的差異。

position_embeddings，句子中每個詞的位置嵌入，用于區別詞的順序。和 transformer 論文中的設計不同，這一塊是訓練出來的，而不是通過 Sinusoidal 函數計算得到的固定嵌入。一般認為這種實現不利于拓展性（難以直接遷移到更長的句子中）。

三個 embedding 不帶權重相加，并通過一層 LayerNorm+dropout 后輸出，其大小為(batch_size, sequence_length, hidden_size)。

補充：這里為什么要用 LayerNorm+Dropout 呢？為什么要用 LayerNorm 而不是 BatchNorm？可以參考一個不錯的回答：transformer 為什么使用 layer normalization，而不是其他的歸一化方法？[2]

2.2 BertEncoder

包含多層 BertLayer，這一塊本身沒有特別需要說明的地方，不過有一個細節值得參考：

利用?gradient checkpointing?技術以降低訓練時的顯存占用。

補充：gradient checkpointing 即梯度檢查點，通過減少保存的計算圖節點壓縮模型占用空間，但是在計算梯度的時候需要重新計算沒有存儲的值，參考論文《Training Deep Nets with Sublinear Memory Cost》，過程如下示意圖：

在 BertEncoder 中，gradient checkpoint 是通過 torch.utils.checkpoint.checkpoint 實現的，使用起來比較方便，可以參考文檔：torch.utils.checkpoint - PyTorch 1.8.1 documentation [3]

這一機制的具體實現比較復雜（沒看懂），在此不作展開。

再往深一層走，就進入了 Encoder 的某一層：

2.2.1 BertLayer

這一層包裝了 BertAttention 和 BertIntermediate+BertOutput（即 Attention 后的 FFN 部分），以及這里直接忽略的 cross-attention 部分（將 BERT 作為 Decoder 時涉及的部分）。

理論上，這里順序調用三個子模塊就可以，沒有什么值得說明的地方。

然而這里又出現了一個細節：

????????#?這是forward的一部分self_attention_outputs?=?self.attention(hidden_states,attention_mask,head_mask,output_attentions=output_attentions,past_key_value=self_attn_past_key_value,)outputs?=?self_attention_outputs[1:]??#?add?self?attentions?if?we?output?attention?weights#?中間省略一部分……layer_output?=?apply_chunking_to_forward(self.feed_forward_chunk,?self.chunk_size_feed_forward,?self.seq_len_dim,?attention_output)outputs?=?(layer_output,)?+?outputs#?省略一部分……return?outputs#?這是feed_forward_chunk的部分def?feed_forward_chunk(self,?attention_output):intermediate_output?=?self.intermediate(attention_output)layer_output?=?self.output(intermediate_output,?attention_output)return?layer_output

看到上面那個apply_chunking_to_forward和feed_forward_chunk了嗎（為什么要整這么復雜，直接調用它不香嗎）？

那么這個apply_chunking_to_forward到底是啥？深入看看：

def?apply_chunking_to_forward(forward_fn:?Callable[...,?torch.Tensor],?chunk_size:?int,?chunk_dim:?int,?*input_tensors )?->?torch.Tensor:"""This?function?chunks?the?:obj:`input_tensors`?into?smaller?input?tensor?parts?of?size?:obj:`chunk_size`?over?thedimension?:obj:`chunk_dim`.?It?then?applies?a?layer?:obj:`forward_fn`?to?each?chunk?independently?to?save?memory.If?the?:obj:`forward_fn`?is?independent?across?the?:obj:`chunk_dim`?this?function?will?yield?the?same?result?asdirectly?applying?:obj:`forward_fn`?to?:obj:`input_tensors`...."""

原來又是一個節約顯存的技術——包裝了一個切分小 batch 或者低維數操作的功能：這里參數chunk_size其實就是切分的 batch 大小，而chunk_dim就是一次計算維數的大小，最后拼接起來返回。

不過，在默認操作中不會特意設置這兩個值（在源代碼中默認為 0 和 1），所以會直接等效于正常的 forward 過程。

繼續往下深入，就是 Transformer 的核心：BertAttention 部分，以及緊隨其后的 FFN 部分。

2.2.1.1 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:returnheads,?index?=?find_pruneable_heads_and_indices(heads,?self.self.num_attention_heads,?self.self.attention_head_size,?self.pruned_heads)#?Prune?linear?layersself.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?headsself.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_headsself.pruned_heads?=?self.pruned_heads.union(heads)?

這里的具體實現概括如下：

• find_pruneable_heads_and_indices是定位需要剪掉的 head，以及需要保留的維度下標 index；

• prune_linear_layer則負責將 Wk/Wq/Wv 權重矩陣（連同 bias）中按照 index 保留沒有被剪枝的維度后轉移到新的矩陣。

接下來就到重頭戲——Self-Attention 的具體實現。

2.2.1.1.1 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_headsself.attention_head_size?=?int(config.hidden_size?/?config.num_attention_heads)self.all_head_size?=?self.num_attention_heads?*?self.attention_head_sizeself.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_embeddingsself.distance_embedding?=?nn.Embedding(2?*?config.max_position_embeddings?-?1,?self.attention_head_size)self.is_decoder?=?config.is_decoder

• 除掉熟悉的 query、key、value 三個權重和一個 dropout，這里還有一個謎一樣的 position_embedding_type，以及 decoder 標記（當然，我不打算介紹 cross-attenton 部分）；

• 注意，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；

• position_embedding_type 是什么？繼續往下看就知道了……

然后是重點，也就是前向傳播過程。

首先回顧一下 multi-head self-attention 的基本公式：

其中 表示注意力頭的個數， 表示向量拼接，。

而這些注意力頭，眾所周知是并行計算的，所以上面的 query、key、value 三個權重是唯一的——這并不是所有 heads 共享了權重，而是“拼接”起來了。

補充：原論文中多頭的理由為 Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this. 而另一個比較靠譜的分析有：為什么 Transformer 需要進行 Multi-head Attention？[4]?

看看 forward 方法：

????def?transpose_for_scores(self,?x):new_x_shape?=?x.size()[:-1]?+?(self.num_attention_heads,?self.attention_head_size)x?=?x.view(*new_x_shape)return?x.permute(0,?2,?1,?3)def?forward(self,hidden_states,attention_mask=None,head_mask=None,encoder_hidden_states=None,encoder_attention_mask=None,past_key_value=None,output_attentions=False,):mixed_query_layer?=?self.query(hidden_states)#?省略一部分cross-attention的計算key_layer?=?self.transpose_for_scores(self.key(hidden_states))value_layer?=?self.transpose_for_scores(self.value(hidden_states))query_layer?=?self.transpose_for_scores(mixed_query_layer)#?Take?the?dot?product?between?"query"?and?"key"?to?get?the?raw?attention?scores.attention_scores?=?torch.matmul(query_layer,?key_layer.transpose(-1,?-2))#?...

這里的 transpose_for_scores 用來把 hidden_size 拆成多個頭輸出的形狀，并且將中間兩維轉置以進行矩陣相乘；

這里 key_layer/value_layer/query_layer 的形狀為：(batch_size, num_attention_heads, sequence_length, attention_head_size)；

這里 attention_scores 的形狀為：(batch_size, num_attention_heads, sequence_length, sequence_length)，符合多個頭單獨計算獲得的 attention map 形狀。

到這里實現了 K 與 Q 相乘，獲得 raw attention scores 的部分，按公式接下來應該是按 dk 進行 scaling 并做 softmax 的操作。然而——

先出現在眼前的是一個奇怪的positional_embedding，以及一堆愛因斯坦求和：

????????#?...if?self.position_embedding_type?==?"relative_key"?or?self.position_embedding_type?==?"relative_key_query":seq_length?=?hidden_states.size()[1]position_ids_l?=?torch.arange(seq_length,?dtype=torch.long,?device=hidden_states.device).view(-1,?1)position_ids_r?=?torch.arange(seq_length,?dtype=torch.long,?device=hidden_states.device).view(1,?-1)distance?=?position_ids_l?-?position_ids_rpositional_embedding?=?self.distance_embedding(distance?+?self.max_position_embeddings?-?1)positional_embedding?=?positional_embedding.to(dtype=query_layer.dtype)??#?fp16?compatibilityif?self.position_embedding_type?==?"relative_key":relative_position_scores?=?torch.einsum("bhld,lrd->bhlr",?query_layer,?positional_embedding)attention_scores?=?attention_scores?+?relative_position_scoreselif?self.position_embedding_type?==?"relative_key_query":relative_position_scores_query?=?torch.einsum("bhld,lrd->bhlr",?query_layer,?positional_embedding)relative_position_scores_key?=?torch.einsum("bhrd,lrd->bhlr",?key_layer,?positional_embedding)attention_scores?=?attention_scores?+?relative_position_scores_query?+?relative_position_scores_key#?...

補充：關于愛因斯坦求和約定，參考以下文檔：torch.einsum - PyTorch 1.8.1 documentation [5]?

補充：這里的positional_embedding引入了 attention map 中的位置嵌入——為什么要這么做呢？我目前還沒搞明白……

對于不同的positional_embedding_type，有三種操作：

• absolute：默認值，這部分就不用處理；

• relative_key：對 key_layer 作處理，將其與這里的positional_embedding和 key 矩陣相乘作為 key 相關的位置編碼；

• relative_key_query：對 key 和 value 都進行相乘以作為位置編碼。

暫時跳過這一迷惑的部分，回到正常 attention 的流程：

????????#?...attention_scores?=?attention_scores?/?math.sqrt(self.attention_head_size)if?attention_mask?is?not?None:#?Apply?the?attention?mask?is?(precomputed?for?all?layers?in?BertModel?forward()?function)attention_scores?=?attention_scores?+?attention_mask??#?這里為什么是+而不是*？#?Normalize?the?attention?scores?to?probabilities.attention_probs?=?nn.Softmax(dim=-1)(attention_scores)#?This?is?actually?dropping?out?entire?tokens?to?attend?to,?which?might#?seem?a?bit?unusual,?but?is?taken?from?the?original?Transformer?paper.attention_probs?=?self.dropout(attention_probs)#?Mask?heads?if?we?want?toif?head_mask?is?not?None:attention_probs?=?attention_probs?*?head_maskcontext_layer?=?torch.matmul(attention_probs,?value_layer)context_layer?=?context_layer.permute(0,?2,?1,?3).contiguous()new_context_layer_shape?=?context_layer.size()[:-2]?+?(self.all_head_size,)context_layer?=?context_layer.view(*new_context_layer_shape)outputs?=?(context_layer,?attention_probs)?if?output_attentions?else?(context_layer,)#?省略decoder返回值部分……return?outputs

重大疑問：這里的attention_scores = attention_scores + attention_mask是在做什么？難道不應該是乘 mask 嗎？

因為這里的 attention_mask 已經【被動過手腳】，將原本為 1 的部分變為 0，而原本為 0 的部分（即 padding）變為一個較大的負數，這樣相加就得到了一個較大的負值：

至于為什么要用【一個較大的負數】？因為這樣一來經過 softmax 操作以后這一項就會變成接近 0 的小數。

(Pdb)?attention_mask tensor([[[[????-0.,?????-0.,?????-0.,??...,?-10000.,?-10000.,?-10000.]]],[[[????-0.,?????-0.,?????-0.,??...,?-10000.,?-10000.,?-10000.]]],[[[????-0.,?????-0.,?????-0.,??...,?-10000.,?-10000.,?-10000.]]],...,[[[????-0.,?????-0.,?????-0.,??...,?-10000.,?-10000.,?-10000.]]],[[[????-0.,?????-0.,?????-0.,??...,?-10000.,?-10000.,?-10000.]]],[[[????-0.,?????-0.,?????-0.,??...,?-10000.,?-10000.,?-10000.]]]],device='cuda:0')

那么，這一步是在哪里執行的呢？

我在modeling_bert.py中沒有找到答案，但是在modeling_utils.py中找到了一個特別的類：class ModuleUtilsMixin，在它的get_extended_attention_mask方法中發現了端倪：

????def?get_extended_attention_mask(self,?attention_mask:?Tensor,?input_shape:?Tuple[int],?device:?device)?->?Tensor:"""Makes?broadcastable?attention?and?causal?masks?so?that?future?and?masked?tokens?are?ignored.Arguments:attention_mask?(:obj:`torch.Tensor`):Mask?with?ones?indicating?tokens?to?attend?to,?zeros?for?tokens?to?ignore.input_shape?(:obj:`Tuple[int]`):The?shape?of?the?input?to?the?model.device:?(:obj:`torch.device`):The?device?of?the?input?to?the?model.Returns::obj:`torch.Tensor`?The?extended?attention?mask,?with?a?the?same?dtype?as?:obj:`attention_mask.dtype`."""#?省略一部分……#?Since?attention_mask?is?1.0?for?positions?we?want?to?attend?and?0.0?for#?masked?positions,?this?operation?will?create?a?tensor?which?is?0.0?for#?positions?we?want?to?attend?and?-10000.0?for?masked?positions.#?Since?we?are?adding?it?to?the?raw?scores?before?the?softmax,?this?is#?effectively?the?same?as?removing?these?entirely.extended_attention_mask?=?extended_attention_mask.to(dtype=self.dtype)??#?fp16?compatibilityextended_attention_mask?=?(1.0?-?extended_attention_mask)?*?-10000.0return?extended_attention_mask

那么，這個函數是在什么時候被調用的呢？和BertModel有什么關系呢？

OK，這里涉及到 BertModel 的繼承細節了：BertModel繼承自BertPreTrainedModel，后者繼承自PreTrainedModel，而PreTrainedModel繼承自[nn.Module, ModuleUtilsMixin, GenerationMixin]三個基類。——好復雜的封裝！

這也就是說，BertModel必然在中間的某個步驟對原始的attention_mask調用了get_extended_attention_mask，導致attention_mask從原始的[1, 0]變為[0, -1e4]的取值。

最終在 BertModel 的前向傳播過程中找到了這一調用（第 944 行）：

????????#?We?can?provide?a?self-attention?mask?of?dimensions?[batch_size,?from_seq_length,?to_seq_length]#?ourselves?in?which?case?we?just?need?to?make?it?broadcastable?to?all?heads.extended_attention_mask:?torch.Tensor?=?self.get_extended_attention_mask(attention_mask,?input_shape,?device)

問題解決了：這一方法不但實現了改變 mask 的值，還將其廣播（broadcast）為可以直接與 attention map 相加的形狀。

不愧是你，HuggingFace（抱臉蟲可不是亂叫的）。

除此之外，值得注意的細節有：

• 按照每個頭的維度進行縮放，對于 bert-base 就是 64 的平方根即 8；

• attention_probs 不但做了 softmax，還用了一次 dropout，這是擔心 attention 矩陣太稠密嗎…… 這里也提到很不尋常，但是原始 Transformer 論文就是這么做的；

• head_mask 就是之前提到的對多頭計算的 mask，如果不設置默認是全 1，在這里就不會起作用；

• context_layer 即 attention 矩陣與 value 矩陣的乘積，原始的大小為：(batch_size, num_attention_heads, sequence_length, attention_head_size) ；

• context_layer 進行轉置和 view 操作以后，形狀就恢復了(batch_size, sequence_length, hidden_size)。

OK, that's all for attention.

2.2.1.1.2 BertSelfOutput

這一塊操作略多但不復雜，一目了然：

class?BertSelfOutput(nn.Module):def?__init__(self,?config):super().__init__()self.dense?=?nn.Linear(config.hidden_size,?config.hidden_size)self.LayerNorm?=?nn.LayerNorm(config.hidden_size,?eps=config.layer_norm_eps)self.dropout?=?nn.Dropout(config.hidden_dropout_prob)def?forward(self,?hidden_states,?input_tensor):hidden_states?=?self.dense(hidden_states)hidden_states?=?self.dropout(hidden_states)hidden_states?=?self.LayerNorm(hidden_states?+?input_tensor)return?hidden_states

補充：這里又出現了 LayerNorm 和 Dropout 的組合，只不過這里是先 Dropout，進行殘差連接后再進行 LayerNorm。至于為什么要做殘差連接，最直接的目的就是降低網絡層數過深帶來的訓練難度，對原始輸入更加敏感～

2.2.1.2 BertIntermediate

看完了 BertAttention，在 Attention 后面還有一個全連接+激活的操作：

class?BertIntermediate(nn.Module):def?__init__(self,?config):super().__init__()self.dense?=?nn.Linear(config.hidden_size,?config.intermediate_size)if?isinstance(config.hidden_act,?str):self.intermediate_act_fn?=?ACT2FN[config.hidden_act]else:self.intermediate_act_fn?=?config.hidden_actdef?forward(self,?hidden_states):hidden_states?=?self.dense(hidden_states)hidden_states?=?self.intermediate_act_fn(hidden_states)return?hidden_states

• 這里的全連接做了一個擴展，以 bert-base 為例，擴展維度為 3072，是原始維度 768 的 4 倍之多；

補充：為什么要過一個 FFN？不知道…… 谷歌最近的論文貌似說明只有 attention 的模型什么用都沒有：

Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Deptharxiv.org

• 這里的激活函數默認實現為 gelu（Gaussian Error Linerar Units(GELUS）： ；當然，它是無法直接計算的，可以用一個包含tanh的表達式進行近似（略)。

作為參考（圖源網絡）：

至于為什么在 transformer 中要用這個激活函數……

補充：看了一些研究，應該是說 GeLU 比 ReLU 這些表現都好，以至于后續的語言模型都沿用了這一激活函數。

2.2.1.3 BertOutput

在這里又是一個全連接 +dropout+LayerNorm，還有一個殘差連接 residual connect：

class?BertOutput(nn.Module):def?__init__(self,?config):super().__init__()self.dense?=?nn.Linear(config.intermediate_size,?config.hidden_size)self.LayerNorm?=?nn.LayerNorm(config.hidden_size,?eps=config.layer_norm_eps)self.dropout?=?nn.Dropout(config.hidden_dropout_prob)def?forward(self,?hidden_states,?input_tensor):hidden_states?=?self.dense(hidden_states)hidden_states?=?self.dropout(hidden_states)hidden_states?=?self.LayerNorm(hidden_states?+?input_tensor)return?hidden_states

這里的操作和 BertSelfOutput 不能說沒有關系，只能說一模一樣…… 非常容易混淆的兩個組件。

以下內容還包含基于 BERT 的應用模型，以及 BERT 相關的優化器和用法，將在下一篇文章作詳細介紹。

2.2.3 BertPooler

這一層只是簡單地取出了句子的第一個token，即[CLS]對應的向量，然后過一個全連接層和一個激活函數后輸出：（這一部分是可選的，因為pooling有很多不同的操作）

class?BertPooler(nn.Module):def?__init__(self,?config):super().__init__()self.dense?=?nn.Linear(config.hidden_size,?config.hidden_size)self.activation?=?nn.Tanh()def?forward(self,?hidden_states):#?We?"pool"?the?model?by?simply?taking?the?hidden?state?corresponding#?to?the?first?token.first_token_tensor?=?hidden_states[:,?0]pooled_output?=?self.dense(first_token_tensor)pooled_output?=?self.activation(pooled_output)return?pooled_output

Takeaways·小結

在 HuggingFace 實現的 Bert 模型中，使用了多種節約顯存的技術：

• gradient checkpoint，不保留前向傳播節點，只在用時計算；

• apply_chunking_to_forward，按多個小批量和低維度計算 FFN 部

BertModel 包含復雜的封裝和較多的組件。以 bert-base 為例，主要組件如下：

• 總計Dropout出現了1+(1+1+1)x12=37次；

• 總計LayerNorm出現了1+(1+1)x12=25次；

• 總計dense全連接層出現了(1+1+1)x12+1=37次，并不是每個dense都配了激活函數……

BertModel 有極大的參數量。以 bert-base 為例，其參數量為 109M，具體計算過程可以參考：HiroLin：小白 Bert 系列-參數計算?[6]。

參考文獻

[1] https://github.com/huggingface/transformers

[2] https://www.zhihu.com/question/395811291/answer/1260290120

[3] https://pytorch.org/docs/stable/checkpoint.html

[4] https://www.zhihu.com/question/341222779/answer/814111138

[5] https://pytorch.org/docs/stable/generated/torch.einsum.html

[6] https://zhuanlan.zhihu.com/p/144582114

更多閱讀

#投 稿?通 道#

?讓你的文字被更多人看到?

如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。

總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個，讓知識真正流動起來。

?????稿件基本要求：

? 文章確系個人原創作品，未曾在公開渠道發表，如為其他平臺已發表或待發表的文章，請明確標注?

? 稿件建議以?markdown?格式撰寫，文中配圖以附件形式發送，要求圖片清晰，無版權問題

? PaperWeekly 尊重原作者署名權，并將為每篇被采納的原創首發稿件，提供業內具有競爭力稿酬，具體依據文章閱讀量和文章質量階梯制結算

?????投稿通道：

? 投稿郵箱：hr@paperweekly.site?

? 來稿請備注即時聯系方式（微信），以便我們在稿件選用的第一時間聯系作者

? 您也可以直接添加小編微信（pwbot02）快速投稿，備注：姓名-投稿

△長按添加PaperWeekly小編

????

現在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關注」訂閱我們的專欄吧

關于PaperWeekly

PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術平臺。如果你研究或從事 AI 領域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。

總結

以上是生活随笔為你收集整理的HuggingFace BERT源码详解：基本模型组件实现的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。