BERT

模型概述

BERT(Bidirectional Encoder Representations from Transformers)是NLP领域的一种预训练模型。本案例中,基于论文BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding实现了BERT模型的OneFlow版本。

模型架构

ModelHidden layersHidden unit sizeAttention headsFeedforward filter sizeMax sequence lengthParameters
BERTBASE12 encoder768124 x 768512110M

BERT 在实际应用中往往分为两步:

  • 首先,预训练得到 BERT 语言模型;

  • 然后,为满足下游应用,在得到的 BERT 语言模型的基础上,多加一层网络,并进行微调,得到下游应用。

快速开始

获取相关数据集

我们提供了完成 BERT 预训练及 SQuAD 微调的 OFRecord 数据集及相关数据文件,可以通过以下命令下载并解压:

  1. wget https://oneflow-static.oss-cn-beijing.aliyuncs.com/oneflow-tutorial-attachments/bert_squad_dataset.zip
  2. unzip bert_squad_dataset.zip

解压后的文件目录清单如下:

  • bert_config.json、vocab.txt:制作 prediction json 文件需要的文件,来自google bert

  • dev-v1.1/、dev-v1.1.json:SQuAD 检验集,用于打分

  • part-0:预训练集样本(40个样本)

  • train-v1.1:SQuAD 训练集,已经转为 ofrecord 数据集格式

以上各个文件将在下文的预训练任务、SQuAD 微调中使用到。

训练 BERT 模型

首先,克隆 OneFlow-Benchmark 仓库。

  1. git clone https://github.com/Oneflow-Inc/OneFlow-Benchmark.git
  2. cd OneFlow-Benchmark/LanguageModeling/BERT/

然后,通过以下命令,使用我们预训练好的 pretrain 模型以及小型样本集合,开始 BERT 预训练查看效果:

  1. python ./run_pretraining.py\
  2. --gpu_num_per_node=1 \
  3. --learning_rate=3e-5 \
  4. --batch_size_per_device=1 \
  5. --iter_num=3 \
  6. --loss_print_every_n_iter=50 \
  7. --seq_length=128 \
  8. --max_predictions_per_seq=20 \
  9. --num_hidden_layers=12 \
  10. --num_attention_heads=12 \
  11. --max_position_embeddings=512 \
  12. --type_vocab_size=2 \
  13. --vocab_size=30522 \
  14. --attention_probs_dropout_prob=0.0 \
  15. --hidden_dropout_prob=0.0 \
  16. --hidden_size_per_head=64 \
  17. --use_boxing_v2=True \
  18. --data_dir=./dataset/ \
  19. --data_part_num=1 \
  20. --log_dir=./bert_regresssioin_test/of \
  21. --loss_print_every_n_iter=5 \
  22. --model_save_dir=./bert_regresssioin_test/of \
  23. --warmup_batches 831 \
  24. --save_last_snapshot True

我们将获得类似以下输出:

  1. ==================================================================
  2. Running bert: num_gpu_per_node = 1, num_nodes = 1.
  3. ==================================================================
  4. gpu_num_per_node = 1
  5. node_num = 1
  6. node_list = None
  7. learning_rate = 3e-05
  8. weight_decay_rate = 0.01
  9. batch_size_per_device = 1
  10. iter_num = 20
  11. warmup_batches = 831
  12. log_every_n_iter = 1
  13. data_dir = ./dataset/
  14. data_part_num = 1
  15. use_fp16 = None
  16. use_boxing_v2 = True
  17. loss_print_every_n_iter = 5
  18. model_save_every_n_iter = 10000
  19. model_save_dir = ./bert_regresssioin_test/of
  20. save_last_snapshot = True
  21. model_load_dir = None
  22. log_dir = ./bert_regresssioin_test/of
  23. seq_length = 128
  24. max_predictions_per_seq = 20
  25. num_hidden_layers = 12
  26. num_attention_heads = 12
  27. max_position_embeddings = 512
  28. type_vocab_size = 2
  29. vocab_size = 30522
  30. attention_probs_dropout_prob = 0.0
  31. hidden_dropout_prob = 0.0
  32. hidden_size_per_head = 64
  33. ------------------------------------------------------------------
  34. Time stamp: 2020-07-06-19:09:29
  35. I0706 19:09:29.605840639 34801 ev_epoll_linux.c:82] Use of signals is disabled. Epoll engine will not be used
  36. Init model on demand
  37. iter 4, total_loss: 11.032, mlm_loss: 10.281, nsp_loss: 0.751, speed: 33.086(sec/batch), 0.151(sentences/sec)
  38. iter 9, total_loss: 11.548, mlm_loss: 10.584, nsp_loss: 0.965, speed: 0.861(sec/batch), 5.806(sentences/sec)
  39. iter 14, total_loss: 10.697, mlm_loss: 10.249, nsp_loss: 0.448, speed: 0.915(sec/batch), 5.463(sentences/sec)
  40. iter 19, total_loss: 10.685, mlm_loss: 10.266, nsp_loss: 0.419, speed: 1.087(sec/batch), 4.602(sentences/sec)
  41. Saving model to ./bert_regresssioin_test/of/last_snapshot.
  42. ------------------------------------------------------------------
  43. average speed: 0.556(sentences/sec)
  44. ------------------------------------------------------------------

详细说明

脚本说明

分类说明所属
pretrain.py、bert.py定义了 BERT 网络模型;BERT
run_pretraining.py启动BERT训练的用户脚本,用户通过命令行参数进行BERT训练的训练环境及超参配置,各个参数的具体作用将在下文 脚本参数 中说明。BERT
squad.py定义了squad网络;SQuAD
run_squad.py用于启动SQuAD的训练SQuAD
run_squad_predict.py使用训练好的SQuAD模型进行预测SQuAD
npy2json.py将OneFlow的预测结果转化为prediction json格式的必要脚本SQuAD
convert_tf_ckpt_to_of.py将TensorFlow模型转为OneFlow的模型格式BERT/SQuAD

脚本参数

run_pretraining.py通过命令行参数配置包括超参在内的训练环境,可以通过 run_pretraining.py --help查看,以下是这些参数作用的具体说明:

  • gpu_num_per_node: 每个节点上 GPU 的数目,OneFlow 要求每个节点的 GPU 数目必须一致

  • node_num: 节点数目,即分布式训练时的主机数目

  • node_list: 节点列表,如果节点数大于1,则需要通过 node_list 指定节点列表,节点列表为字符串形式,采用逗号分隔,如--node_num=2 --node_list="192.168.1.12,192.168.1.14"

  • learning_rate: Learning rate

  • weight_decay_rate:设置权重衰减率

  • batch_size_per_device: 分布式训练时每个设备上的batch大小

  • iter_num ITER_NUM: 训练的总轮数

  • warmup_batches: 预热轮数,默认值为10000

  • data_dir: OFRecord数据集的路径

  • data_part_num:OFRecord数据集目录下的数据文件数目

  • use_fp16: 是否使用fp16

  • use_boxing_v2: 是否使用boxing v2

  • loss_print_every_n_iter:训练中每隔多少轮打印一次训练信息(loss信息)

  • model_save_every_n_iter: 训练中每隔多少轮保存一次模型

  • model_save_dir: 模型存储路径

  • save_last_snapshot:指定最后一轮训练完成后,模型保存路径

  • model_load_dir:指定模型加载路径

  • log_dir LOG_DIR:指定日志路径

  • seq_length: 指定BERT句子长度,默认值为512

  • max_predictions_per_seq: 默认值为80

  • num_hidden_layers:隐藏层数目,默认值为24

  • num_attention_heads: Attention头数目,默认值为16

使用完整的 Wikipedia + BookCorpus 数据集

如果需要从无到有进行 BERT 的 pretrain 训练,则需要使用较大的训练集。

如果感兴趣,可以通过 google-research BERT 的页面,下载 tfrecord 格式的数据集。再根据加载与准备OFRecord数据集中的方法,将 TFRecord 数据转为 OFRecord 数据集使用。

将 Tensorflow 的 BERT 模型转为 OneFlow 模型格式

如果想直接使用已经训练好的 pretrained 模型做 fine-tune 任务(如以下将展示的SQuAD),可以考虑直接从 google-research BERT 页面下载已经训练好的 BERT 模型。

再利用我们提供的 convert_tf_ckpt_to_of.py 脚本,将其转为 OneFlow 模型格式。转换过程如下:

首先,下载并解压某个版本的 BERT 模型,如 uncased_L-12_H-768_A-12

  1. wget https://storage.googleapis.com/bert_models/2020_02_20/uncased_L-12_H-768_A-12.zip
  2. unzip uncased_L-12_H-768_A-12.zip -d uncased_L-12_H-768_A-12

然后,运行以下命令:

  1. cd uncased_L-12_H-768_A-12/
  2. cat > checkpoint <<ONEFLOW
  3. model_checkpoint_path: "bert_model.ckpt"
  4. all_model_checkpoint_paths: "bert_model.ckpt"
  5. ONEFLOW

该命令将在解压目录下创建一个 checkpoint 文件,并写入以下内容:

  1. model_checkpoint_path: "bert_model.ckpt"
  2. all_model_checkpoint_paths: "bert_model.ckpt"

此时,已经准备好待转化的 TensorFlow 模型目录,整个模型目录的结构如下:

  1. uncased_L-12_H-768_A-12
  2. ├── bert_config.json
  3. ├── bert_model.ckpt.data-00000-of-00001
  4. ├── bert_model.ckpt.index
  5. ├── checkpoint
  6. └── vocab.txt

我们接着使用 convert_tf_ckpt_to_of.py 将 TensorFlow 模型转为 OneFlow 模型:

  1. python convert_tf_ckpt_to_of.py \
  2. --tf_checkpoint_path ./uncased_L-12_H-768_A-12 \
  3. --of_dump_path ./uncased_L-12_H-768_A-12-oneflow

以上命令,将转化好的 OneFlow 格式的模型保存在 ./uncased_L-12_H-768_A-12-oneflow 目录下,供后续微调训练(如:SQuAD)使用。

微调:SQuAD 问答任务

将 pretrained 模型修改为 SQuAD 模型

我们只需要在 BERT 的 backbone 基础上,加上一层 output 层,并修改 loss 的表达式即可,完整的代码可以查看 squad.py 脚本,以下是几处关键修改:

  1. def SQuADTrain():
  2. #...
  3. backbone = bert_util.BertBackbone()
  4. #在BERT的基础上加上一个全连接层
  5. with flow.name_scope("cls-squad"):
  6. final_hidden = backbone.sequence_output()
  7. final_hidden_matrix = flow.reshape(final_hidden, [-1, hidden_size])
  8. logits = bert_util._FullyConnected(
  9. final_hidden_matrix,
  10. hidden_size,
  11. units=2,
  12. weight_initializer=bert_util.CreateInitializer(initializer_range),
  13. name='output')
  14. logits = flow.reshape(logits, [-1, seq_length, 2])
  15. start_logits = flow.slice(logits, [None, None, 0], [None, None, 1])
  16. end_logits = flow.slice(logits, [None, None, 1], [None, None, 1])
  17. #重新定义SQuAD任务的loss
  18. start_loss = _ComputeLoss(start_logits, start_positions_blob, seq_length)
  19. end_loss = _ComputeLoss(end_logits, end_positions_blob, seq_length)
  20. total_loss = 0.5*(start_loss + end_loss)
  21. return total_loss

为了得到一个初始化的 squad 模型,我们通过以下脚本启动 squad 训练,并保存模型。

  1. python ./run_squad.py\
  2. --gpu_num_per_node=1\
  3. --learning_rate=3e-5\
  4. --batch_size_per_device=2\
  5. --iter_num=50\
  6. --loss_print_every_n_iter=50\
  7. --seq_length=384\
  8. --max_predictions_per_seq=20\
  9. --num_hidden_layers=12\
  10. --num_attention_heads=12\
  11. --max_position_embeddings=512\
  12. --type_vocab_size=2\
  13. --vocab_size=30522\
  14. --attention_probs_dropout_prob=0.0\
  15. --hidden_dropout_prob=0.0\
  16. --hidden_size_per_head=64\
  17. --use_boxing_v2=True\
  18. --data_dir=./dataset/train-v1.1\
  19. --data_part_num=1\
  20. --log_dir=./bert_regresssioin_test/of\
  21. --model_save_dir=./bert_regresssioin_test/of\
  22. --warmup_batches 831\
  23. --save_last_snapshot True

完成训练后,在 ./bert_regresssioin_test/of/last_snapshot 中保存有初始化的 SQuAD 模型,我们将其与训练好的 BERT 合并后,进行微调(fine-tune)训练。

合并 pretrained 模型为 SQuAD 模型

SQuAD 模型是在 pretrained 模型基础上的扩充,我们需要参照模型的加载与保存中的“模型部分初始化和部分导入”方法,将训练好的 BERT pretrained 模型与初始化的 SQuAD 模型合并。

  1. cp -R ./bert_regresssioin_test/of/last_snapshot ./squadModel
  2. cp -R --remove-destination ./dataset/uncased_L-12_H-768_A-12_oneflow/* ./squadModel/

OneFlow 预训练模型的训练次数问题

OneFlow 生成的模型目录中,会有一个名为 System-Train-TrainStep-xxx 的子目录(xxx为作业函数的函数名),该子目录下的 out 文件中,保存有训练总迭代数,并且这个迭代数会用于动态调节训练过程的learning rate

为了防止保存的迭代数影响到微调的训练,应该将out文件中的二进制数据清零:

  1. cd System-Train-TrainStep-xxx
  2. xxd -r > out <<ONEFLOW
  3. 00000000: 0000 0000 0000 0000
  4. ONEFLOW

如果你使用的是由 TensorFlow 转过来的预训练模型,则可以省去这个步骤。

开始 SQuAD 训练

通过 run_suqad.py 脚本,开始训练 SQuAD 模型,主要配置如下:

  • 使用以上合并得到的 SQuAD 模型 ./squadModel

  • 采用 SQuAD v1.1 作为训练集

  • epoch = 3 (iternum = 88641*3/(4*8) = 8310)

  • learning rate = 3e-5

  1. python ./run_squad.py\
  2. --gpu_num_per_node=4\
  3. --learning_rate=3e-5\
  4. --batch_size_per_device=8\
  5. --iter_num=8310\
  6. --loss_print_every_n_iter=50\
  7. --seq_length=384\
  8. --max_predictions_per_seq=20\
  9. --num_hidden_layers=12\
  10. --num_attention_heads=12\
  11. --max_position_embeddings=512\
  12. --type_vocab_size=2\
  13. --vocab_size=30522\
  14. --attention_probs_dropout_prob=0.0\
  15. --hidden_dropout_prob=0.0\
  16. --hidden_size_per_head=64\
  17. --use_boxing_v2=True\
  18. --data_dir=./dataset/train-v1.1\
  19. --data_part_num=8\
  20. --log_dir=./bert_regresssioin_test/of\
  21. --model_save_dir=./bert_regresssioin_test/of\
  22. --warmup_batches 831\
  23. --save_last_snapshot True\
  24. --model_load_dir=./squadModel

预测及打分

生成为了生成 Preidiction File 格式的 json 文件,我们先将预测结果保存为 npy 文件,再使用 google BERT的run_squad.py 中的 write_predictions 函数,转化为 json 格式。

利用 run_squad_predict.py 生成 all_results.npy 文件:

  1. python run_squad_predict.py \
  2. --gpu_num_per_node=1 \
  3. --batch_size_per_device=4 \
  4. --iter_num=2709 \
  5. --seq_length=384 \
  6. --max_predictions_per_seq=20 \
  7. --num_hidden_layers=12 \
  8. --num_attention_heads=12 \
  9. --max_position_embeddings=512 \
  10. --type_vocab_size=2 \
  11. --vocab_size=30522 \
  12. --attention_probs_dropout_prob=0.0 \
  13. --hidden_dropout_prob=0.0 \
  14. --hidden_size_per_head=64 \
  15. --use_boxing_v2=True \
  16. --data_part_num=1 \
  17. --data_dir=./dataset/dev-v1.1 \
  18. --log_dir=./bert_regresssioin_test/of \
  19. --model_load_dir=path/to/squadModel \
  20. --warmup_batches 831

注意将以上 model_load_dir 修改为 训练好的 squadModel。

得到 all_results.npy 文件后,在google bert仓库目录下(注意该仓库的 tensorflow 版本为 tensorflow v1 ),运行我们提供的 npy2json.py (由 google bert 中的 run_squand.py 修改得来):

  1. python npy2json.py\
  2. --vocab_file=./dataset/vocab.txt \
  3. --bert_config_file=./dataset/bert_config.json \
  4. --do_train=False \
  5. --do_predict=True \
  6. --all_results_file=./all_results.npy \
  7. --predict_file=./dataset/dev-v1.1.json \
  8. --max_seq_length=384 \
  9. --doc_stride=128 \
  10. --output_dir=./squad_base/

注意将 all_results_file 修改为上一步得到的 all_results.npy 的路径。

最终,得到 predictions.json 文件,可以使用 evaluate-v1.1.py 进行打分。

  1. python evaluate-v1.1.py \
  2. ./dataset/dev-v1.1.json \
  3. path/to/squad_base/predictions.json

分布式训练

如之前介绍脚本参数时描述:进行分布式训练,只需要在启动训练脚本式加入 node_num 选项指定主机数目及 node_list 选项即可:

  1. python run_squad_predict.py \
  2. --gpu_num_per_node=1 \
  3. --batch_size_per_device=4 \
  4. --iter_num=2709 \
  5. --seq_length=384 \
  6. --max_predictions_per_seq=20 \
  7. --num_hidden_layers=12 \
  8. --num_attention_heads=12 \
  9. --max_position_embeddings=512 \
  10. --type_vocab_size=2 \
  11. --vocab_size=30522 \
  12. --attention_probs_dropout_prob=0.0 \
  13. --hidden_dropout_prob=0.0 \
  14. --hidden_size_per_head=64 \
  15. --use_boxing_v2=True \
  16. --data_part_num=1 \
  17. --data_dir=./dataset/dev-v1.1 \
  18. --log_dir=./bert_regresssioin_test/of \
  19. --model_load_dir=path/to/squadModel \
  20. --warmup_batches 831 \
  21. --node_num=2 \
  22. --node_list="192.168.1.12,192.168.1.14"