華為諾亞方舟開源哪吒、TinyBERT模型，可直接下載使用

杜佳豪發表於2019-12-05
模型

昨日，華為諾亞方舟實驗室的 NLP 團隊開源了兩個重要的預訓練語言模型——哪吒和 TinyBERT。這兩個模型可以直接下載、預訓練和微調。華為語音語義首席科學家劉群在微博上轉發了這一訊息。

華為諾亞方舟開源哪吒、TinyBERT模型，可直接下載使用

專案地址：https://github.com/huawei-noah/Pretrained-Language-Model

根據 GitHub 上的介紹，這一專案是諾亞方舟實驗室用來開源各種預訓練模型的專案，目前有兩個，日後不排除有更多模型加入進來。

哪吒模型

該專案中第一個開源的模型是哪吒（NEZHA：NEural contextualiZed representation for CHinese lAnguage understanding），是華為諾亞方舟實驗室自研的預訓練語言模型，在一些 NLP 任務上取得了 SOTA 的表現。這一模型基於 BERT，可以在普通的 GPU 叢集上進行訓練，同時融合了英偉達和谷歌程式碼的早期版本。

華為諾亞方舟開源哪吒、TinyBERT模型，可直接下載使用在今年的智源大會上，劉群介紹了哪吒模型的相關工作。

哪吒是一個基於 BERT 進行最佳化和改進的預訓練語言模型，中文采用的語料是 Wikipedia 和 Baike 和 News，而谷歌的中文語料只用了 Wikipedia。英文的哪吒採用的是 Wikipedia 和 BookCorpus，跟中文一樣。

哪吒在訓練過程中採用了華為雲，能夠實現多卡多機訓練。當然，本次釋出的哪吒和華為的版本有些不同，可以在一般的 GPU 上訓練。此外，哪吒還採用了英偉達的混合精度訓練方法，本次專案開源的時候也公開了。

預訓練過程

模型的使用方法和 BERT 類似。首先，使用者需要準備資料，方法和 BERT 準備資料的過程類似：

python utils/create_pretraining_data.py \ --input_file=./sample_text.txt \ --output_file=/tmp/tf_examples.tfrecord \ --vocab_file=./your/path/vocab.txt \ --do_lower_case=True \ --max_seq_length=128 \ --max_predictions_per_seq=20 \ --masked_lm_prob=0.15 \ --random_seed=12345 \ --dupe_factor=5

在訓練過程中，使用者需要準備 horovod 分散式訓練環境，然後執行 script 資料夾中的 run_pretraining.sh 檔案即可。

微調過程

目前，哪吒模型支援三種任務上的微調：文字分類、序列標註和類似 SQuAD 的機器閱讀理解，微調的相關程式碼基於谷歌的相關程式碼。

預訓練過程如下：

  • 下載預訓練模型，解壓模型檔案、詞彙檔案和配置檔案到資料夾「nezha」中；

  • 建立微調任務，根據任務不同執行不同的指令碼：

    • scripts/run_clf.sh 指令碼用於文字分類任務，如 LCQMC、ChnSenti 和 XNLI；

    • scripts/run_seq_labelling.sh 指令碼用於序列標註任務，如 Peoples-daily-NER；

    • scripts/run_reading.sh 指令碼用於機器閱讀理解任務，如 CMRC 2018；

  • 從相關輸出位置獲得結果。

需要注意的是，CMRC 任務的評估稍有不同，需要執行以下程式碼：

python cmrc2018_evaluate.py data/cmrc/cmrc2018_dev.json output/cmrc/dev_predictions.json output/cmrc/metric.txt.

以上任務需要的檔案都儲存在了專案資料夾中。

哪吒模型已經可以下載，已有四種中文的預訓練模型，分別是 base、large 和對應的 mask 和全詞 mask 型別。

下載地址：https://pan.baidu.com/s/1V7btNIDqBHvz4g9LOPLeeg

密碼：x3qk

TinyBERT

本專案開源的另一個模型是 TinyBERT，這是一個透過蒸餾方法獲得的 BERT 模型。

華為諾亞方舟開源哪吒、TinyBERT模型，可直接下載使用TinyBERT 的建模過程。

相比原版的 BERT-base，TinyBERT 比它小了 7.5 倍，推理速度則快了 9.4 倍。無論是在預訓練階段還是特定任務學習階段，TinyBERT 的效能都更好。

使用方法

TinyBERT 的使用方法和哪吒專案有所不同，它分為三個步驟：蒸餾、資料增強以及特定任務蒸餾。

1. 蒸餾

在第一個蒸餾階段，研究者需要使用原版的 BERT-base（沒有微調的版本）作為教師模型、透過 Transformer 蒸餾的方法，可以獲得一個泛化的 TinyBERT 模型。之後再進行特定任務蒸餾。

泛化蒸餾有兩個步驟：生成 json 形式的語料，然後進行蒸餾。

第一步，使用 use pregenerate_training_data.py 檔案，生成語料。

# ${BERT_BASE_DIR}$ includes the BERT-base teacher model.python pregenerate_training_data.py --train_corpus ${CORPUS_RAW} \  --bert_model ${BERT_BASE_DIR}$ \ --reduce_memory --do_lower_case \ --epochs_to_generate 3 \ --output_dir ${CORPUS_JSON_DIR}$

第二步，使用 use general_distill.py 進行蒸餾。

# ${STUDENT_CONFIG_DIR}$ includes the config file of student_model. python general_distill.py --pregenerated_data ${CORPUS_JSON}$ \  --teacher_model ${BERT_BASE}$ \ --student_model ${STUDENT_CONFIG_DIR}$ \ --reduce_memory --do_lower_case \ --train_batch_size 256 \ --output_dir ${GENERAL_TINYBERT_DIR}$

在專案中，研究者提供了已蒸餾過的 TinyBERT 模型。

2. 資料增強

資料增強是一個能夠擴充套件特定任務中訓練集的方法。透過學習更多工相關的樣本，學生模型的泛化能力可以進一步提升。研究者在資料增強階段結合了 BERT 和 GloVe 的方法，進行詞級別的替換，以此來增強資料。

在這裡，使用者可以使用 data_augmentation.py 檔案進行資料增強。增強後的資料集檔案會自動儲存在對應的 ${GLUE_DIR/TASK_NAME}$ 中。

python data_augmentation.py --pretrained_bert_model ${BERT_BASE_DIR}$ \ --glove_embs ${GLOVE_EMB}$ \ --glue_dir ${GLUE_DIR}$ \  --task_name ${TASK_NAME}$

在執行 GLUE 任務中的資料增強程式碼前，你需要下載 GLUE 資料，並解壓檔案到 GLUE_DIR 中。在命令列中，TASK_NAME 可以是 GLUE 任務中的 CoLA、SST-2、MRPC、STS-B、QQP、MNLI、QNLI 和 RTE。

3. 特定任務蒸餾

這一階段主要對 transformer 進行進一步蒸餾，以提升在特定任務上的效能。這一階段分兩步：首先進行中間層蒸餾，然後進行預測層蒸餾。

執行 task_distill.py 即可進行中間層蒸餾。

# ${FT_BERT_BASE_DIR}$ contains the fine-tuned BERT-base model.python task_distill.py --teacher_model ${FT_BERT_BASE_DIR}$ \ --student_model ${GENERAL_TINYBERT_DIR}$ \ --data_dir ${TASK_DIR}$ \ --task_name ${TASK_NAME}$ \  --output_dir ${TMP_TINYBERT_DIR}$ \ --max_seq_length 128 \ --train_batch_size 32 \ --num_train_epochs 10 \ --aug_train \ --do_lower_case

之後執行 task_distill.py 即可進行預測層蒸餾。

python task_distill.py --pred_distill \ --teacher_model ${FT_BERT_BASE_DIR}$ \ --student_model ${TMP_TINYBERT_DIR}$ \ --data_dir ${TASK_DIR}$ \ --task_name ${TASK_NAME}$ \ --output_dir ${TINYBERT_DIR}$ \ --aug_train \  --do_lower_case \ --learning_rate 3e-5 \ --num_train_epochs 3 \ --eval_step 100 \ --max_seq_length 128 \ --train_batch_size 32

模型效能評估

TinyBERT 的模型效能評估使用的是單獨的程式碼，執行 task_distill.py 即可。

${TINYBERT_DIR}$ includes the config file, student model and vocab file.python task_distill.py --do_eval \ --student_model ${TINYBERT_DIR}$ \ --data_dir ${TASK_DIR}$ \ --task_name ${TASK_NAME}$ \ --output_dir ${OUTPUT_DIR}$ \ --do_lower_case \ --eval_batch_size 32 \ --max_seq_length 128

安裝方法

對於 TinyBERT，使用者需要下載專門的模型，安裝相關依賴（基於 Python3）：

pip install -r requirements.txt

模型分為以下幾種：

1. 沒有特定任務蒸餾的泛化模型：

論文復現版本：

General_TinyBERT(4layer-312dim)

下載地址：https://drive.google.com/uc?export=download&id=1dDigD7QBv1BmE6pWU71pFYPgovvEqOOj

General_TinyBERT(6layer-768dim)

下載地址：https://drive.google.com/uc?export=download&id=1wXWR00EHK-Eb7pbyw0VP234i2JTnjJ-x

第二版 (2019/11/18) 使用更多資料、無 [MASK] 語料版本：

General_TinyBERT_v2(4layer-312dim)

下載地址：https://drive.google.com/open?id=1PhI73thKoLU2iliasJmlQXBav3v33-8z

General_TinyBERT_v2(6layer-768dim)

下載地址：https://drive.google.com/open?id=1r2bmEsQe4jUBrzJknnNaBJQDgiRKmQjF

2. GLUE 任務上蒸餾後的模型

每個任務對應一個模型，這些模型可以使用之前提供的評估程式碼進行評估：

TinyBERT(4layer-312dim)

下載地址：https://drive.google.com/uc?export=download&id=1_sCARNCgOZZFiWTSgNbE7viW_G5vIXYg

TinyBERT(6layer-768dim)

下載地址：https://drive.google.com/uc?export=download&id=1Vf0ZnMhtZFUE0XoD3hTXc6QtHwKr_PwS

專案作者表示，之後他們會測試 TinyBERT 在中文任務上的表現、嘗試使用哪吒或 ALBERT 作為教師模型進行蒸餾，以及公開效能更好的模型等。

相關文章