新智元報道
來源:GitHub
作者:huggingface 編譯:肖琴
【新智元導讀】谷歌NLP模型的官方TensorFlow實現很強,現在,它的PyTorch版本來了!只需簡單運行一次轉換腳本,就可得到一個PyTorch模型,且結果與原始版本相近,甚至更好。
上週,谷歌最強NLP模型BERT開源了官方TensorFlow代碼和預訓練模型,引起大量關注。
現在,PyTorch用戶的福利來了:一個名爲Hugging Face的團隊近日公開了BERT模型的谷歌官方TensorFlow庫的op-for-op PyTorch重新實現:
https://github.com/huggingface/pytorch-pretrained-BERT
這個實現可以爲BERT加載任何預訓練的TensorFlow checkpoint(特別是谷歌的官方預訓練模型),並提供一個轉換腳本。
BERT-base和BERT-large模型的參數數量分別爲110M和340M,爲了獲得良好的性能,很難使用推薦的batch size在單個GPU上對其進行微調。爲了幫助微調模型,這個repo還提供了3種可以在微調腳本中激活技術:梯度累積(gradient-accumulation)、 multi-GPU 和分佈式訓練。
其結果如下:
- 在序列級MRPC分類任務上,該實現使用小型BERT-base模型再現了原始實現的84%-88%的準確率。
- 在token級的SQuAD 任務上,該個實現使用小型BERT-base模型再現了原始實現的88.52 F1的結果。
作者表示,正致力於在其他任務以及更大的BERT模型上重現結果。
BERT模型的PyTorch實現
這個存儲庫包含了谷歌BERT模型的官方TensorFlow存儲庫的op-for-op PyTorch重新實現。谷歌的官方存儲庫是與BERT論文一起發佈的:BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding,作者是Jacob Devlin、Ming-Wei Chang、Kenton Lee和Kristina Toutanova。
這個實現可以爲BERT加載任何預訓練的TensorFlow checkpoint(特別是谷歌的預訓練模型),並提供了一個轉換腳本(見下文)。
此外,我們將在本週晚些時候添加多語言版本和中文版本的模型代碼。
腳本:加載任何TensorFlow檢查點
使用convert_tf_checkpoint_to_pytorch.py腳本,你可以在PyTorch保存文件中轉換BERT的任何TensorFlow檢查點(尤其是谷歌發佈的官方預訓練模型)。
這個腳本將TensorFlow checkpoint(以bert_model.ckpt開頭的三個文件)和相關的配置文件(bert_config.json)作爲輸入,併爲此配置創建PyTorch模型,從PyTorch模型的TensorFlow checkpoint加載權重並保存生成的模型在一個標準PyTorch保存文件中,可以使用 torch.load() 導入(請參閱extract_features.py,run_classifier.py和run_squad.py中的示例)。
只需要運行一次這個轉換腳本,就可以得到一個PyTorch模型。然後,你可以忽略TensorFlow checkpoint(以bert_model.ckpt開頭的三個文件),但是一定要保留配置文件(bert_config.json)和詞彙表文件(vocab.txt),因爲PyTorch模型也需要這些文件。
要運行這個特定的轉換腳本,你需要安裝TensorFlow和PyTorch。該庫的其餘部分只需要PyTorch。
下面是一個預訓練的BERT-Base Uncased 模型的轉換過程示例:
export BERT_BASE_DIR=/path/to/bert/uncased_L-12_H-768_A-12 python convert_tf_checkpoint_to_pytorch.py \ --tf_checkpoint_path $BERT_BASE_DIR/bert_model.ckpt \ --bert_config_file $BERT_BASE_DIR/bert_config.json \ --pytorch_dump_path $BERT_BASE_DIR/pytorch_model.bin
你可以在這裏下載Google的預訓練轉換模型:
https://github.com/google-research/bert#pre-trained-models
BERT的PyTorch模型
在這個庫裏,我們提供了三個PyTorch模型,你可以在modeling.py中找到:
- BertModel - 基本的BERT Transformer 模型
- BertForSequenceClassification - 頂部帶有sequence classification head的BERT模型
- BertForQuestionAnswering - 頂部帶有token classification head 的BERT模型,
以下是每類模型的一些細節。
1 . BertModel
BertModel是一個基本的BERT Transformer模型,包含一個summed token、位置和序列嵌入層,然後是一系列相同的self-attention blocks(BERT-base是12個blocks, BERT-large是24個blocks)。
輸入和輸出與TensorFlow 模型的輸入和輸出相同。
具體來說,該模型的輸入是:
- input_ids:一個形狀爲[batch_size, sequence_length]的torch.LongTensor,在詞彙表中包含單詞的token索引
- token_type_ids:形狀[batch_size, sequence_length]的可選torch.LongTensor,在[0,1]中選擇token類型索引。類型0對應於句子A,類型1對應於句子B。
- attention_mask:一個可選的torch.LongTensor,形狀爲[batch_size, sequence_length],索引在[0,1]中選擇。
模型的輸出是由以下內容組成的一個元組:
- all_encoder_layers:一個大小爲[batch_size, sequence_length,hidden_size]的torch.FloatTensor列表,它是每個注意塊末端隱藏狀態的完整序列列表(即BERT-base的12個完整序列,BERT-large的24個完整序列)
- pooled_output:一個大小爲[batch_size, hidden_size]的torch.FloatTensor,它是在與輸入(CLF)的第一個字符相關聯的隱藏狀態之上預訓練的分類器的輸出,用於訓練Next-Sentence任務(參見BERT的論文)。
extract_features.py腳本提供了有關如何使用這類模型的示例,該腳本可用於爲給定輸入提取模型的隱藏狀態。
2 . BertForSequenceClassification
BertForSequenceClassification是一個fine-tuning 模型,包括BertModel,以及BertModel頂部的一個序列級分類器(sequence-level classifier)。
序列級分類器是一個線性層,它將輸入序列中第一個字符的最後隱藏狀態作爲輸入(參見BERT論文中的圖3a和3b)。
run_classifier.py腳本提供了關於如何使用此類模型的示例,該腳本可用於使用BERT微調單個序列(或序列對)分類器,例如用於MRPC任務。
3. BertForQuestionAnswering
BertForQuestionAnswering是一個fine-tuning 模型,包括BertModel,它在最後隱藏狀態的完整序列之上具有token級分類器(token-level classifiers)。
token-level 分類器將最後隱藏狀態的完整序列作爲輸入,併爲每個token計算得分,(參見BERT論文的圖3c和3d)。
run_squad.py腳本提供了有關如何使用此類模型的示例,該腳本可用於使用BERT微調token分類器,例如用於SQuAD任務。
安裝、要求、測試
這段代碼在Python 3.5+上進行了測試。必備條件是:
- PyTorch (> = 0.4.1)
- tqdm
安裝dependencies:
pip install -r ./requirements.txt
測試文件夾中包含一系列測試,可以使用pytest運行(如果需要,請安裝pytest: pip install pytest)。
你可以使用以下命令運行測試:
python -m pytest -sv tests/
大批量訓練:梯度積累、多GPU、分佈式訓練
BERT-base和BERT-large的模型參數分別是110M和340M,爲了獲得良好的性能(大多數情況下批大小是32),很難在單個GPU上對它們進行微調。
爲了幫助對這些模型進行微調,我們介紹了在微調腳本run_classifier.py和run_squad中可以激活的四種技術:優化CPU、梯度積累、multi-gpu和分佈式訓練。
有關如何使用這些技術的更多細節,你可以閱讀這篇關於PyTorch批量訓練技巧的文章:
https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255
BERT的微調:運行示例
我們展示了與原始實現相同的示例:在MRPC分類語料庫上微調sequence級分類器和在問題回答數據集SQuAD上微調token級分類器。
在運行這些示例之前,應該先下載GLUE數據,並將其解壓縮到某個目錄$GLUE_DIR。還需下載BERT-Base checkpoint,將其解壓縮到某個目錄$BERT_BASE_DIR,並將其轉換爲上一節所述的PyTorch版本。
這個示例代碼基於微軟研究意譯語料庫(MRPC)調優了BERT-Base,在單個K-80上運行只需不到10分鐘。
export GLUE_DIR=/path/to/glue python run_classifier.py \ --task_name MRPC \ --do_train \ --do_eval \ --do_lower_case \ --data_dir $GLUE_DIR/MRPC/ \ --vocab_file $BERT_BASE_DIR/vocab.txt \ --bert_config_file $BERT_BASE_DIR/bert_config.json \ --init_checkpoint $BERT_PYTORCH_DIR/pytorch_model.bin \ --max_seq_length 128 \ --train_batch_size 32 \ --learning_rate 2e-5 \ --num_train_epochs 3.0 \ --output_dir /tmp/mrpc_output/
基於原始實現的超參數進行測試,評估結果達到84%到88%。
第二個示例是基於SQuAD問題回答任務微調BERT-Base。
export SQUAD_DIR=/path/to/SQUAD python run_squad.py \ --vocab_file $BERT_BASE_DIR/vocab.txt \ --bert_config_file $BERT_BASE_DIR/bert_config.json \ --init_checkpoint $BERT_PYTORCH_DIR/pytorch_model.bin \ --do_train \ --do_predict \ --do_lower_case --train_file $SQUAD_DIR/train-v1.1.json \ --predict_file $SQUAD_DIR/dev-v1.1.json \ --train_batch_size 12 \ --learning_rate 3e-5 \ --num_train_epochs 2.0 \ --max_seq_length 384 \ --doc_stride 128 \ --output_dir ../debug_squad/
使用之前的超參數進行訓練,得到如下結果:
{"f1": 88.52381567990474, "exact_match": 81.22043519394512}
在GPU上微調BERT-large
上面列出的選項允許在GPU上很容易地對BERT-large進行微調,而不是像原始實現那樣使用TPU。
例如,針對SQuAD任務微調BERT-large模型,可以在服務器上用4個k-80在18個小時內完成。我們的結果與TensorFlow的實現結果相似(實際上是略高):
{"exact_match": 84.56953642384106, "f1": 91.04028647786927}
爲了得到這些結果,我們使用了以下組合:
- 多GPU訓練(在多GPU服務器上自動激活),
- 梯度累積
- 在CPU上執行優化步驟,將Adam的平均值存儲在RAM中。
以下是我們在此次運行中使用的超參數的完整列表:
python ./run_squad.py --vocab_file $BERT_LARGE_DIR/vocab.txt --bert_config_file $BERT_LARGE_DIR/bert_config.json --init_checkpoint $BERT_LARGE_DIR/pytorch_model.bin --do_lower_case --do_train --do_predict --train_file $SQUAD_TRAIN --predict_file $SQUAD_EVAL --learning_rate 3e-5 --num_train_epochs 2 --max_seq_length 384 --doc_stride 128 --output_dir $OUTPUT_DIR/bert_large_bsz_24 --train_batch_size 24 --gradient_accumulation_steps 2 --optimize_on_cpu