一、run_classifier.py文件
#chang 1:引入horovod
import horovod.tensorflow as hvd
...
def get_train_examples(self, data_dir):
"""See base class."""
file_path = os.path.join(data_dir, "train_bert.txt")
#chang 2 修改训练数据创建方式
return self._create_examples_train(file_path, "train")
#...
def get_labels(self):
"""See base class."""
return ["0", "1"]
#chang 3 新增训练数据创建代码
def _create_examples_train(self, file_path, set_type):
"""Creates examples for the training and dev sets."""
examples = []
for (i, line) in enumerate(codecs.open(file_path, "r", "utf-8")):
if i % hvd.size() == hvd.rank():
guid = "%s-%d" % (set_type, i)
split_line = line.strip('\r\n').split('\t')
text_a = tokenization.convert_to_unicode(split_line[1])
label = split_line[0]
examples.append(
InputExample(guid=guid, text_a=text_a, text_b=None, label=label))
return examples
def _create_examples(self, file_path, set_type):
...
def main(_):
#chang 4 初始化
hvd.init()
#chang 5 为不同的rank定义不同的输出文件名
FLAGS.output_dir = FLAGS.output_dir if hvd.rank() == 0 else os.path.join(FLAGS.output_dir, str(hvd.rank()))
...
#chang 6 确保一个GPU上一个进程
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.visible_device_list = str(hvd.local_rank())
run_config = tf.contrib.tpu.RunConfig(
...,
#chang 7 设置配置
session_config=config)
...
if FLAGS.do_train:
#change 8 定义训练记录文件
train_file = os.path.join(FLAGS.output_dir, "train.tf_record")
file_based_convert_examples_to_features(
...
#change 9 确保每个进程拥有同样的初始化权重
hooks = [hvd.BroadcastGlobalVariablesHook(0)]
estimator.train(input_fn=train_input_fn, max_steps=num_train_steps, hooks=hooks)
#change 10 #只在编号为0的GPU上进行评估
if FLAGS.do_eval and hvd.rank() == 0
...
#change 11#只在编号为0的GPU上进行预测
if FLAGS.do_predict and hvd.rank() == 0
二、optimization.py文件
#change 1
import horovod.tensorflow as hvd
#change 2
def create_optimizer(loss, init_lr, num_train_steps, num_warmup_steps, use_tpu):
...
#Change 3 (optional) 学习率乘以GPU个数
optimizer = AdamWeightDecayOptimizer(
learning_rate=learning_rate * hvd.size(),
...)
#Change 4 使用hvd分布式优化器
optimizer = hvd.DistributedOptimizer(optimizer)
#Change 5 使用分布式优化器计算梯度
grads_and_vars=optimizer.compute_gradients(loss, tvars)
#Change 6 根据分布式优化器的调试进行梯度裁剪
grads = [grad for grad,var in grads_and_vars]
tvars = [var for grad,var in grads_and_vars]
(grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0)
#Change 7 调整梯度裁剪进行优化
train_op = optimizer.apply_gradients(zip(grads, tvars), global_step=global_step)
参考文献:BERT Multi-GPU implementation using TensorFlow and Horovod with code