多GPU（2卡）TensorFlow训练梯度平均成瓶颈且速度波动问题

我来帮你捋捋这个用2块GPU训练神经网络时，梯度平均成为性能瓶颈、训练速度忽快忽慢的问题，结合你提供的代码片段，给你几个实用的优化方向：

先确认问题根源 #

你通过TensorFlow Timeline定位到梯度平均是瓶颈，这在多GPU同步训练里太常见了——每块GPU各自算完梯度后，得等所有GPU都完成才能做梯度平均，要是某块GPU因为数据加载慢、计算负载不均拖了后腿，整个批次的训练时间就会被拉长，自然就出现速度波动了。

先看看你现有代码的潜在问题 #

你给出的训练循环片段是这样的：

for i in range(resume_epoch, c.num_epochs):
    print("Epoch %d" % i)
    sess.run(train_itr.initializer)
    num_batches = num_egs // c.batch_size
    for batch in range(num_batches):
        start_time = time.time()
        _, loss_value = sess.run([train_op, loss])
        duration = time.time() - start_time
        examples_per_sec = c.batch_size / float(duration)
        print('step %d, loss = %.2f (%.1f examples/sec)' % (batch, loss_value, examples_per_sec))

这里的train_op如果是手动实现的跨GPU梯度平均，很大概率是用了阻塞式的同步梯度聚合——必须等所有GPU梯度都计算完才能进行平均，一旦某块GPU慢半拍，整个流程就得等它，这就是速度波动的核心原因之一。

针对性的优化方案 #

1. 换用更高效的梯度聚合方式 #

• 如果业务允许，试试异步梯度更新：要是你的任务对梯度同步的严格性要求不高（比如一些图像分类任务），可以把同步训练改成异步训练——每块GPU算完梯度直接更新参数，不用等其他GPU，这样彻底避免了梯度同步的等待开销，速度波动自然就没了。不过要注意异步训练会带来参数更新的噪声，得适当调小学习率或者调整其他超参数。
• 用TensorFlow内置的分布式策略：如果必须用同步训练，别自己手动写梯度平均了，直接用tf.distribute.MirroredStrategy这类官方分布式策略。它内部用了更高效的集体通信框架（比如NCCL）来做梯度聚合，还能自动处理GPU负载均衡，比手动实现的效率高太多，能大幅降低梯度平均的耗时和波动。

2. 优化数据加载，消除GPU等待 #

很多时候梯度计算的延迟是因为数据加载跟不上，导致某块GPU没事干，拖慢了整个梯度同步的节奏：

• 把数据加载改成tf.data.Dataset的并行加载模式，用prefetch(tf.data.AUTOTUNE)让数据预加载到GPU，map的时候设置num_parallel_calls=tf.data.AUTOTUNE并行处理数据，确保GPU一直有数据可算。
• 检查下两块GPU的batch分配是否均匀，别出现某块GPU要处理更多数据的情况，负载均衡了，梯度计算的耗时才会稳定。

3. 调整梯度更新的粒度，减少同步次数 #

可以试试梯度累积：比如每累积4个batch的梯度再做一次平均和参数更新，这样梯度平均的频率就降低了，单次同步的开销对整体速度的影响也会变小，训练速度会更平稳。

验证优化效果 #

改完之后，再用TensorFlow Timeline跑一遍，看看梯度平均的耗时是不是降下来了，同时监控两块GPU的利用率，如果利用率变得更稳定，说明优化起作用了。

内容的提问来源于stack exchange，提问作者Yi Bill