如何实现多GPU并发: goole提供的BERT源码使用TPUEstimator实现的,TPUEstimator支持多TPU并发但不支持多GPU并发。 为了能支持多GPU并发充分利用GPU资源提高效率,对BERT做一下修改: 1、使用optimization_multigpu.py替换optimization.py,执行一下命令 cp optimization.py optimization.bak cp optimization_multigpu.py optimization.py 2、修改训练代码,预训练代码在run_pretrain.py,句子分类run_classifier.py,Squad任务run_squad.py。 以Squad为例,其他文件可以用相同的方式修改。修改后的文件为run_squad_multigpu.py。 1) 添加num_gpus入参 flags.DEFINE_integer("num_gpus", 1, "Total number of GPUs to use.") 2)在model_fn_builder函数中修改EstimatorSpec,修改点有三处,修改成: a)添加scafflod: if init_checkpoint: (assignment_map, initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint) if use_tpu: def tpu_scaffold(): tf.train.init_from_checkpoint(init_checkpoint, assignment_map) return tf.train.Scaffold() scaffold_fn = tpu_scaffold else: tf.train.init_from_checkpoint(init_checkpoint, assignment_map) scaffold_fn = tf.train.Scaffold() b)修改Train的Spec: output_spec = tf.estimator.EstimatorSpec( mode=mode, loss=total_loss, train_op=train_op, scaffold=scaffold_fn) c)修改predict的Spec: output_spec = tf.estimator.EstimatorSpec(mode=mode, predictions=predictions, scaffold=scaffold_fn) 3)在main函数创建RunConfig地方修改RunConfig,RunConfig添加分布式策略 distribute = tf.contrib.distribute.MirroredStrategy(num_gpus=FLAGS.num_gpus) run_config = tf.estimator.RunConfig( model_dir=FLAGS.output_dir, save_checkpoints_steps=FLAGS.save_checkpoints_steps, train_distribute=distribute) 注意:这里可以删除TPUEstimator的RunConfig,也可以保留,保留的话添加分支判断。 4)修改main函数中调用model_fn_builder的地方,修改成: model_fn = model_fn_builder( bert_config=bert_config, init_checkpoint=FLAGS.init_checkpoint, learning_rate=FLAGS.learning_rate, num_train_steps=num_train_steps, num_warmup_steps=num_warmup_steps, use_tpu=False, use_one_hot_embeddings=False) 5)main函数中将TPUEstimator修改成tf.estimator.Estimator params = {"batch_size": FLAGS.train_batch_size} estimator = tf.estimator.Estimator( model_fn=model_fn, config=run_config, params = params) 6)修改后可以不覆盖run_squad.py,直接使用run_squad_multigpu.py运行 python run_squad_multigpu.py --vocab_file=./uncased_L-12_H-768_A-12/vocab.txt --bert_config_file=./uncased_L-12_H-768_A-12/bert_config.json --init_checkpoint=./uncased_L-12_H-768_A-12/bert_model.ckpt --do_train=True --train_file=./data/train-v2.0.json --do_predict=True --predict_file=./data/dev-v2.0.json --learning_rate=3e-5 --num_train_epochs=2.0 --max_seq_length=384 --doc_stride=128 --train_batch_size=12 --num_gpus=2 --output_dir=./output/result_dir --version_2_with_negative=True 通过num_gpus指定GPU个数,train_batch_size修改batch_size,V100的显存大可以改成24,修改学习率learning_rate和num_train_epochs