Huggingface BERT源码详解：应用模型与训练优化
?PaperWeekly 原創(chuàng) ·?作者｜李濼秋
學(xué)校｜浙江大學(xué)碩士生
研究方向｜自然語(yǔ)言處理、知識(shí)圖譜
接上篇，記錄一下對(duì) HuggingFace 開源的 Transformers 項(xiàng)目代碼的理解。
本文基于 Transformers 版本 4.4.2（2021 年 3 月 19 日發(fā)布）項(xiàng)目中，pytorch 版的 BERT 相關(guān)代碼，從代碼結(jié)構(gòu)、具體實(shí)現(xiàn)與原理，以及使用的角度進(jìn)行分析，包含以下內(nèi)容：
1. BERT Tokenization 分詞模型（BertTokenizer）
2. BERT Model 本體模型（BertModel）
3.?1.?BertEmbeddings
? ? 2. BertEncoder
? ? 3.1.?BertLayer
? ? ? ? 2.1.?BertAttention
? ? ? ? ? ??2.1.?BertIntermediate
? ? ? ? ? ? ? ?2. BertOutput
? ? ? ? ? ? 3. BertEmbeddings
? ? ? ?? ? ?4. BertEncoder
? ? ? ? 3. BERT-based Models應(yīng)用模型
4. BertForPreTraining
5. 1. BertForSequenceClassification
? ? 2. BertForMultiChoice
? ? 3. BertForTokenClassification
? ? 4. BertForQuestionAnswering
? ? 5. BERT訓(xùn)練與優(yōu)化
6. BERT訓(xùn)練與優(yōu)化
7. 1. Pre-Training
? ? 2. Fine-Tuning
? ? 3. 1. AdamW
? ? ? ? 2.?Warmup
BERT-based Models
基于 BERT 的模型都寫在/models/bert/modeling_bert.py里面，包括 BERT 預(yù)訓(xùn)練模型和 BERT 分類模型，UML 圖如下：
BERT模型一圖流（建議保存后放大查看）：
▲ 畫圖工具：Pyreverse
首先，以下所有的模型都是基于BertPreTrainedModel這一抽象基類的，而后者則基于一個(gè)更大的基類PreTrainedModel。這里我們關(guān)注BertPreTrainedModel的功能：
用于初始化模型權(quán)重，同時(shí)維護(hù)繼承自PreTrainedModel的一些標(biāo)記身份或者加載模型時(shí)的類變量。
下面，首先從預(yù)訓(xùn)練模型開始分析。
3.1 BertForPreTraining
眾所周知，BERT 預(yù)訓(xùn)練任務(wù)包括兩個(gè)：
Masked Language Model（MLM）：在句子中隨機(jī)用[MASK]替換一部分單詞，然后將句子傳入 BERT 中編碼每一個(gè)單詞的信息，最終用[MASK]的編碼信息預(yù)測(cè)該位置的正確單詞，這一任務(wù)旨在訓(xùn)練模型根據(jù)上下文理解單詞的意思；
Next Sentence Prediction（NSP）：將句子對(duì) A 和 B 輸入 BERT，使用[CLS]的編碼信息進(jìn)行預(yù)測(cè) B 是否 A 的下一句，這一任務(wù)旨在訓(xùn)練模型理解預(yù)測(cè)句子間的關(guān)系。
▲ 圖源網(wǎng)絡(luò)
而對(duì)應(yīng)到代碼中，這一融合兩個(gè)任務(wù)的模型就是BertForPreTraining，其中包含兩個(gè)組件：
class?BertForPreTraining(BertPreTrainedModel):def?__init__(self,?config):super().__init__(config)self.bert?=?BertModel(config)self.cls?=?BertPreTrainingHeads(config)self.init_weights()#?...這里的BertModel在上一篇文章中已經(jīng)詳細(xì)介紹了（注意，這里設(shè)置的是默認(rèn)add_pooling_layer=True，即會(huì)提取[CLS]對(duì)應(yīng)的輸出用于 NSP 任務(wù)），而BertPreTrainingHeads則是負(fù)責(zé)兩個(gè)任務(wù)的預(yù)測(cè)模塊：
class?BertPreTrainingHeads(nn.Module):def?__init__(self,?config):super().__init__()self.predictions?=?BertLMPredictionHead(config)self.seq_relationship?=?nn.Linear(config.hidden_size,?2)def?forward(self,?sequence_output,?pooled_output):prediction_scores?=?self.predictions(sequence_output)seq_relationship_score?=?self.seq_relationship(pooled_output)return?prediction_scores,?seq_relationship_score?又是一層封裝：BertPreTrainingHeads包裹了BertLMPredictionHead 和一個(gè)代表 NSP 任務(wù)的線性層。這里不把 NSP 對(duì)應(yīng)的任務(wù)也封裝一個(gè)BertXXXPredictionHead，估計(jì)是因?yàn)樗?jiǎn)單了，沒有必要……
補(bǔ)充：其實(shí)是有封裝這個(gè)類的，不過(guò)它叫做BertOnlyNSPHead，在這里用不上……
繼續(xù)下探BertPreTrainingHeads ：
class?BertLMPredictionHead(nn.Module):def?__init__(self,?config):super().__init__()self.transform?=?BertPredictionHeadTransform(config)#?The?output?weights?are?the?same?as?the?input?embeddings,?but?there?is#?an?output-only?bias?for?each?token.self.decoder?=?nn.Linear(config.hidden_size,?config.vocab_size,?bias=False)self.bias?=?nn.Parameter(torch.zeros(config.vocab_size))#?Need?a?link?between?the?two?variables?so?that?the?bias?is?correctly?resized?with?`resize_token_embeddings`self.decoder.bias?=?self.biasdef?forward(self,?hidden_states):hidden_states?=?self.transform(hidden_states)hidden_states?=?self.decoder(hidden_states)return?hidden_states這個(gè)類用于預(yù)測(cè)[MASK]位置的輸出在每個(gè)詞作為類別的分類輸出，注意到：
該類重新初始化了一個(gè)全 0 向量作為預(yù)測(cè)權(quán)重的 bias；
該類的輸出形狀為[batch_size, seq_length, vocab_size]，即預(yù)測(cè)每個(gè)句子每個(gè)詞是什么類別的概率值（注意這里沒有做 softmax）；
又一個(gè)封裝的類：BertPredictionHeadTransform，用來(lái)完成一些線性變換：
補(bǔ)充：感覺這一層去掉也行？輸出的形狀也沒有發(fā)生變化。我個(gè)人的理解是和 Pooling 那里做一個(gè)對(duì)稱的操作，同樣過(guò)一層 dense 再接分類器……
回到BertForPreTraining，繼續(xù)看兩塊 loss 是怎么處理的。它的前向傳播和BertModel的有所不同，多了labels和next_sentence_label 兩個(gè)輸入：
labels：形狀為[batch_size, seq_length] ，代表 MLM 任務(wù)的標(biāo)簽，注意這里對(duì)于原本未被遮蓋的詞設(shè)置為 -100，被遮蓋詞才會(huì)有它們對(duì)應(yīng)的 id，和任務(wù)設(shè)置是反過(guò)來(lái)的。
例如，原始句子是I want to [MASK] an apple，這里我把單詞eat給遮住了輸入模型，對(duì)應(yīng)的label設(shè)置為[-100, -100, -100, 【eat對(duì)應(yīng)的id】, -100, -100]；
為什么要設(shè)置為 -100 而不是其他數(shù)？因?yàn)閠orch.nn.CrossEntropyLoss默認(rèn)的ignore_index=-100，也就是說(shuō)對(duì)于標(biāo)簽為 100 的類別輸入不會(huì)計(jì)算 loss。
next_sentence_label：這一個(gè)輸入很簡(jiǎn)單，就是 0 和 1 的二分類標(biāo)簽。
OK，接下來(lái)兩部分 loss 的組合：
直接相加，就是這么單純的策略。
當(dāng)然，這份代碼里面也包含了對(duì)于只想對(duì)單個(gè)目標(biāo)進(jìn)行預(yù)訓(xùn)練的 BERT 模型（具體細(xì)節(jié)不作展開）：
BertForMaskedLM：只進(jìn)行 MLM 任務(wù)的預(yù)訓(xùn)練；
基于BertOnlyMLMHead，而后者也是對(duì)BertLMPredictionHead的另一層封裝；
BertLMHeadModel：這個(gè)和上一個(gè)的區(qū)別在于，這一模型是作為 decoder 運(yùn)行的版本；
同樣基于BertOnlyMLMHead；
BertForNextSentencePrediction：只進(jìn)行 NSP 任務(wù)的預(yù)訓(xùn)練。
基于BertOnlyNSPHead，內(nèi)容就是一個(gè)線性層……
接下來(lái)介紹的是各種 Fine-tune 模型，基本都是分類任務(wù)：
▲ 圖源：原始BERT論文附錄
3.2 BertForSequenceClassification
這一模型用于句子分類（也可以是回歸）任務(wù)，比如 GLUE benchmark 的各個(gè)任務(wù)。
句子分類的輸入為句子（對(duì)），輸出為單個(gè)分類標(biāo)簽。
結(jié)構(gòu)上很簡(jiǎn)單，就是BertModel（有 pooling）過(guò)一個(gè) dropout 后接一個(gè)線性層輸出分類：
?class?BertForSequenceClassification(BertPreTrainedModel):def?__init__(self,?config):super().__init__(config)self.num_labels?=?config.num_labelsself.bert?=?BertModel(config)self.dropout?=?nn.Dropout(config.hidden_dropout_prob)self.classifier?=?nn.Linear(config.hidden_size,?config.num_labels)self.init_weights()#?...在前向傳播時(shí)，和上面預(yù)訓(xùn)練模型一樣需要傳入labels輸入。
如果初始化的num_labels=1，那么就默認(rèn)為回歸任務(wù)，使用 MSELoss；
否則認(rèn)為是分類任務(wù)。
3.3 BertForMultipleChoice
這一模型用于多項(xiàng)選擇，如 RocStories/SWAG 任務(wù)。
多項(xiàng)選擇任務(wù)的輸入為一組分次輸入的句子，輸出為選擇某一句子的單個(gè)標(biāo)簽。
結(jié)構(gòu)上與句子分類相似，只不過(guò)線性層輸出維度為 1，即每次需要將每個(gè)樣本的多個(gè)句子的輸出拼接起來(lái)作為每個(gè)樣本的預(yù)測(cè)分?jǐn)?shù)。
實(shí)際上，具體操作時(shí)是把每個(gè) batch 的多個(gè)句子一同放入的，所以一次處理的輸入為[batch_size, num_choices]數(shù)量的句子，因此相同 batch 大小時(shí)，比句子分類等任務(wù)需要更多的顯存，在訓(xùn)練時(shí)需要小心。
3.4 BertForTokenClassification
這一模型用于序列標(biāo)注（詞分類），如 NER 任務(wù)。
序列標(biāo)注任務(wù)的輸入為單個(gè)句子文本，輸出為每個(gè) token 對(duì)應(yīng)的類別標(biāo)簽。
由于需要用到每個(gè) token對(duì)應(yīng)的輸出而不只是某幾個(gè)，所以這里的BertModel不用加入 pooling 層；
同時(shí)，這里將_keys_to_ignore_on_load_unexpected這一個(gè)類參數(shù)設(shè)置為[r"pooler"]，也就是在加載模型時(shí)對(duì)于出現(xiàn)不需要的權(quán)重不發(fā)生報(bào)錯(cuò)。
3.5 BertForQuestionAnswering
這一模型用于解決問(wèn)答任務(wù)，例如 SQuAD 任務(wù)。
問(wèn)答任務(wù)的輸入為問(wèn)題 +（對(duì)于 BERT 只能是一個(gè)）回答組成的句子對(duì)，輸出為起始位置和結(jié)束位置用于標(biāo)出回答中的具體文本。
這里需要兩個(gè)輸出，即對(duì)起始位置的預(yù)測(cè)和對(duì)結(jié)束位置的預(yù)測(cè)，兩個(gè)輸出的長(zhǎng)度都和句子長(zhǎng)度一樣，從其中挑出最大的預(yù)測(cè)值對(duì)應(yīng)的下標(biāo)作為預(yù)測(cè)的位置。
對(duì)超出句子長(zhǎng)度的非法 label，會(huì)將其壓縮（torch.clamp_）到合理范圍。
作為一個(gè)遲到的補(bǔ)充，這里稍微介紹一下ModelOutput這個(gè)類。它作為上述各個(gè)模型輸出包裝的基類，同時(shí)支持字典式的存取和下標(biāo)順序的訪問(wèn)，繼承自python原生的OrderedDict 類。
以上就是關(guān)于 BERT 源碼的介紹，下面介紹一些關(guān)于 BERT 模型實(shí)用的訓(xùn)練細(xì)節(jié)。
BERT訓(xùn)練和優(yōu)化
4.1 Pre-Training
預(yù)訓(xùn)練階段，除了眾所周知的 15%、80% mask 比例，有一個(gè)值得注意的地方就是參數(shù)共享。
不止 BERT，所有 huggingface 實(shí)現(xiàn)的 PLM 的 word embedding 和 masked language model 的預(yù)測(cè)權(quán)重在初始化過(guò)程中都是共享的：
class?PreTrainedModel(nn.Module,?ModuleUtilsMixin,?GenerationMixin):#?...def?tie_weights(self):"""Tie?the?weights?between?the?input?embeddings?and?the?output?embeddings.If?the?:obj:`torchscript`?flag?is?set?in?the?configuration,?can't?handle?parameter?sharing?so?we?are?cloningthe?weights?instead."""output_embeddings?=?self.get_output_embeddings()if?output_embeddings?is?not?None?and?self.config.tie_word_embeddings:self._tie_or_clone_weights(output_embeddings,?self.get_input_embeddings())if?self.config.is_encoder_decoder?and?self.config.tie_encoder_decoder:if?hasattr(self,?self.base_model_prefix):self?=?getattr(self,?self.base_model_prefix)self._tie_encoder_decoder_weights(self.encoder,?self.decoder,?self.base_model_prefix)#?...至于為什么，應(yīng)該是因?yàn)?word_embedding 和 prediction 權(quán)重太大了，以 bert-base 為例，其尺寸為(30522, 768)，降低訓(xùn)練難度。
4.2 Fine-Tuning
微調(diào)也就是下游任務(wù)階段，也有兩個(gè)值得注意的地方。
4.2.1 AdamW
首先介紹一下 BERT 的優(yōu)化器：AdamW（AdamWeightDecayOptimizer）。
這一優(yōu)化器來(lái)自 ICLR 2017 的 Best Paper：《Fixing Weight Decay Regularization in Adam》中提出的一種用于修復(fù) Adam 的權(quán)重衰減錯(cuò)誤的新方法。論文指出，L2 正則化和權(quán)重衰減在大部分情況下并不等價(jià)，只在 SGD 優(yōu)化的情況下是等價(jià)的；而大多數(shù)框架中對(duì)于 Adam+L2 正則使用的是權(quán)重衰減的方式，兩者不能混為一談。
AdamW 是在 Adam+L2 正則化的基礎(chǔ)上進(jìn)行改進(jìn)的算法，與一般的 Adam+L2 的區(qū)別如下：
關(guān)于 AdamW 的分析可以參考：
AdamW and Super-convergence is now the fastest way to train neural nets [1]
paperplanet：都 9102 年了，別再用 Adam + L2 regularization了 [2]
ICLR 2018 有什么值得關(guān)注的亮點(diǎn)？[3]
話說(shuō)，《STABLE WEIGHT DECAY REGULARIZATION》這篇好像吐槽AdamW 的 Weight Decay 實(shí)現(xiàn)還是有問(wèn)題…… 有空整整優(yōu)化器相關(guān)的內(nèi)容。
通常，我們會(huì)選擇模型的 weight 部分參與 decay 過(guò)程，而另一部分（包括 LayerNorm 的 weight）不參與（代碼最初來(lái)源應(yīng)該是 Huggingface 的示例）：
補(bǔ)充：關(guān)于這么做的理由，我暫時(shí)沒有找到合理的解答，但是找到了一些相關(guān)的討論：https://forums.fast.ai/t/is-weight-decay-applied-to-the-bias-term/73212/4forums.fast.ai
4.2.2 Warmup
BERT 的訓(xùn)練中另一個(gè)特點(diǎn)在于 Warmup，其含義為：
在訓(xùn)練初期使用較小的學(xué)習(xí)率（從 0 開始），在一定步數(shù)（比如 1000 步）內(nèi)逐漸提高到正常大小（比如上面的 2e-5），避免模型過(guò)早進(jìn)入局部最優(yōu)而過(guò)擬合；
在訓(xùn)練后期再慢慢將學(xué)習(xí)率降低到 0，避免后期訓(xùn)練還出現(xiàn)較大的參數(shù)變化。
在 Huggingface 的實(shí)現(xiàn)中，可以使用多種 warmup 策略：
TYPE_TO_SCHEDULER_FUNCTION?=?{SchedulerType.LINEAR:?get_linear_schedule_with_warmup,SchedulerType.COSINE:?get_cosine_schedule_with_warmup,SchedulerType.COSINE_WITH_RESTARTS:?get_cosine_with_hard_restarts_schedule_with_warmup,SchedulerType.POLYNOMIAL:?get_polynomial_decay_schedule_with_warmup,SchedulerType.CONSTANT:?get_constant_schedule,SchedulerType.CONSTANT_WITH_WARMUP:?get_constant_schedule_with_warmup, }具體而言：
CONSTANT：保持固定學(xué)習(xí)率不變；
CONSTANT_WITH_WARMUP：在每一個(gè) step 中線性調(diào)整學(xué)習(xí)率；
LINEAR：上文提到的兩段式調(diào)整；
COSINE：和兩段式調(diào)整類似，只不過(guò)采用的是三角函數(shù)式的曲線調(diào)整；
COSINE_WITH_RESTARTS：訓(xùn)練中將上面 COSINE 的調(diào)整重復(fù) n 次；
POLYNOMIAL：按指數(shù)曲線進(jìn)行兩段式調(diào)整。
具體使用參考transformers/optimization.py：
最常用的還是get_linear_scheduler_with_warmup即線性兩段式調(diào)整學(xué)習(xí)率的方案……
以上即為關(guān)于 transformers 庫(kù)（4.4.2 版本）中 BERT 相關(guān)代碼的具體實(shí)現(xiàn)分析，歡迎與讀者共同交流探討。
參考文獻(xiàn)
[1] https://www.fast.ai/2018/07/02/adam-weight-decay/
[2] https://zhuanlan.zhihu.com/p/63982470
[3] https://www.zhihu.com/question/67335251/answer/262989932
特別鳴謝
感謝 TCCI 天橋腦科學(xué)研究院對(duì)于 PaperWeekly 的支持。TCCI 關(guān)注大腦探知、大腦功能和大腦健康。
更多閱讀
#投 稿?通 道#
?讓你的文字被更多人看到?
如何才能讓更多的優(yōu)質(zhì)內(nèi)容以更短路徑到達(dá)讀者群體，縮短讀者尋找優(yōu)質(zhì)內(nèi)容的成本呢？答案就是：你不認(rèn)識(shí)的人。
總有一些你不認(rèn)識(shí)的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學(xué)者和學(xué)術(shù)靈感相互碰撞，迸發(fā)出更多的可能性。?
PaperWeekly 鼓勵(lì)高校實(shí)驗(yàn)室或個(gè)人，在我們的平臺(tái)上分享各類優(yōu)質(zhì)內(nèi)容，可以是最新論文解讀，也可以是學(xué)術(shù)熱點(diǎn)剖析、科研心得或競(jìng)賽經(jīng)驗(yàn)講解等。我們的目的只有一個(gè)，讓知識(shí)真正流動(dòng)起來(lái)。
?????稿件基本要求：
? 文章確系個(gè)人原創(chuàng)作品，未曾在公開渠道發(fā)表，如為其他平臺(tái)已發(fā)表或待發(fā)表的文章，請(qǐng)明確標(biāo)注?
? 稿件建議以?markdown?格式撰寫，文中配圖以附件形式發(fā)送，要求圖片清晰，無(wú)版權(quán)問(wèn)題
? PaperWeekly 尊重原作者署名權(quán)，并將為每篇被采納的原創(chuàng)首發(fā)稿件，提供業(yè)內(nèi)具有競(jìng)爭(zhēng)力稿酬，具體依據(jù)文章閱讀量和文章質(zhì)量階梯制結(jié)算
?????投稿通道：
? 投稿郵箱：hr@paperweekly.site?
? 來(lái)稿請(qǐng)備注即時(shí)聯(lián)系方式（微信），以便我們?cè)诟寮x用的第一時(shí)間聯(lián)系作者
? 您也可以直接添加小編微信（pwbot02）快速投稿，備注：姓名-投稿
△長(zhǎng)按添加PaperWeekly小編
????
現(xiàn)在，在「知乎」也能找到我們了
進(jìn)入知乎首頁(yè)搜索「PaperWeekly」
點(diǎn)擊「關(guān)注」訂閱我們的專欄吧
關(guān)于PaperWeekly
PaperWeekly 是一個(gè)推薦、解讀、討論、報(bào)道人工智能前沿論文成果的學(xué)術(shù)平臺(tái)。如果你研究或從事 AI 領(lǐng)域，歡迎在公眾號(hào)后臺(tái)點(diǎn)擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。
總結(jié)
以上是生活随笔為你收集整理的Huggingface BERT源码详解：应用模型与训练优化的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。
- 上一篇： ACL 2021 | SimCLS: 概
- 下一篇： 一步搞定模型训练和商品召回：京东全新索引