多GPU训练bert

概述

    做nlp的肯定都听说过bert了，它属于迁移学习，官方提供了通用语料bert预训练结果，通过微调，可以提高下游任务的准确率。但是bert使用的维基百科这类的通用语料，对特定领域的文本特征学的不是很好。例如我所在的医疗领域，bert并不能很好区分医学特征，此时就需要我们对自己的语料进行预训练。bert官方也提到利用fine-tune任务的语料进行预训练，能提高fine-tune任务的准确率。

训练bert

    官方提供的bert是使用tpu集群训练的代码，他可以兼容单核GPU、多核CPU，如果你的机器有多个GPU，直接运行python run_pretraining.py，发现只能利用一个GPU。其实git上已经有人提供了多GPU版本，亲证是可行的，只需要将他的run_pretraining_gpu.py替换run_pretraining.py、optimization_gpu.py替换optimization.py即可，然后将modeling.py大概344行：assignment_map[name] = name修改assignment_map[name] = name_to_variable[name]就可以了。

减小bert预训练结果大小

    bert的模型结果大概有1.2G，但是官方提供的预训练结果只有400多M，这是因为bert默认会保存adm优化器的中间参数，这个只是在训练时需要，在预测时只需要他的最终参数。网上提供了减小模型参数的方法：https://github.com/google-research/bert/issues/99，代码我不复制了。但是你多GPU的结果使用该方法发现并不可行，网上查了很多资料还是不行，最后想了一个笨办法。我在本地以多GPU的结果继续少量数据的run_pretraining.py训练（将init_checkpoint指定为GPU的ckpt），快速结束训练并保存结果数据。然后对该本地的结果数据使用前面的提到的代码，就能将模型压缩到400多M左右。该办法真的很暴力，大家如果发现了好的方法希望在留意区分享交流。