从 11 月初开始,google-research 就陆续开源了 BERT 的各个版本。google 此次开源的 BERT 是通过 tensorflow 高级 API—— tf.estimator 进行封装 (wrapper) 的。因此对于不同数据集的适配,只需要修改代码中的 processor 部分,就能进行代码的训练、交叉验证和测试。
在自己的数据集上运行 BERT
BERT 的代码同论文里描述的一致,主要分为两个部分。一个是训练语言模型(language model)的预训练(pretrain)部分。另一个是训练具体任务 (task) 的 fine-tune 部分。在开源的代码中,预训练的入口是在 run_pretraining.py 而 fine-tune 的入口针对不同的任务分别在 run_classifier.py 和 run_squad.py。其中 run_classifier.py 适用的任务为分类任务。如 CoLA、MRPC、MultiNLI 这些数据集。而 run_squad.py 适用的是阅读理解 (MRC) 任务,如 squad2.0 和 squad1.1。
预训练是 BERT 很重要的一个部分,与此同时,预训练需要巨大的运算资源。按照论文里描述的参数,其 Base 的设定在消费级的显卡 Titan x 或 Titan 1080ti(12GB RAM) 上,甚至需要近几个月的时间进行预训练,同时还会面临显存不足的问题。不过所幸的是谷歌满足了 Issues#2 里各国开发者的请求,针对大部分语言都公布了 BERT 的预训练模型。因此在我们可以比较方便地在自己的数据集上进行 fine-tune。
下载预训练模型
对于中文而言,google 公布了一个参数较小的 BERT 预训练模型。具体参数数值如下所示:
Chinese Simplified and Traditional, 12-layer, 768-hidden, 12-heads, 110M parameters
模型的下载链接可以在 github 上 google 的开源代码里找到。对下载的压缩文件进行解压,可以看到文件里有五个文件,其中 bert_model.ckpt 开头的文件是负责模型变量载入的,而 vocab.txt 是训练时中文文本采用的字典,最后 bert_config.json 是 BERT 在训练时,可选调整的一些参数。