计算机视觉(十四):Tensorflow分布式训练

Posted GeniusAng

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了计算机视觉(十四):Tensorflow分布式训练相关的知识,希望对你有一定的参考价值。

计算机视觉笔记总目录


当我们拥有大量计算资源时,通过使用合适的分布式策略,我们可以充分利用这些计算资源,从而大幅压缩模型训练的时间。针对不同的使用场景,TensorFlow 在 tf.distribute.Strategy中为我们提供了若干种分布式策略,使得我们能够更高效地训练模型。

1 TensorFlow 分布式的分类

图间并行(又称数据并行)

  • 每个机器上都会有一个完整的模型,将数据分散到各个机器,分别计算梯度。

图内并行(又称模型并行)

  • 每个机器分别负责整个模型的一部分计算任务。

1.1 图间并行用的非常多,会包含两种方式进行更新

同步:收集到足够数量的梯度,一同更新,下图上方

异步: 即同步方式下,更新需求的梯度数量为1,下图下方

注:PS(parameter server):维护全局共享的模型参数的服务器

1.2 实现方式

  • 单机多卡
  • 多机单卡

TensorFlow多机多卡实现思路。多机多卡的分布式有很多实现方式,比如:

  1. 将每个GPU当做一个worker
  2. 同一个机器的各个GPU进行图内并行
  3. 同一个机器的各个GPU进行图间并行

比如说第三种:模型实现封装成函数、将数据分成GPU数量的份数、在每个GPU下,进行一次模型forward计算,并使用优化器算出梯度、reduce每个GPU下的梯度,并将梯度传入到分布式中的优化器中

2 单机多卡训练: MirroredStrategy

tf.distribute.MirroredStrategy 是一种简单且高性能的,数据并行的同步式分布式策略,主要支持多个 GPU 在同一台主机上训练。

MirroredStrategy运行原理:

  • 1、训练开始前,该策略在所有 N 个计算设备(GPU)上均各复制一份完整的模型
  • 2、每次训练传入一个批次的数据时,将数据分成 N 份,分别传入 N 个计算设备(即数据并行)
  • 3、使用分布式计算的 All-reduce 操作,在计算设备间高效交换梯度数据并进行求和,使得最终每个设备都有了所有设备的梯度之和,使用梯度求和的结果更新本地变量
    • 当所有设备均更新本地变量后,进行下一轮训练(即该并行策略是同步的)。默认情况下,TensorFlow 中的 MirroredStrategy 策略使用 NVIDIA NCCL 进行 All-reduce 操作。

构建代码步骤:

使用这种策略时,我们只需实例化一个 MirroredStrategy 策略:

strategy = tf.distribute.MirroredStrategy()

并将模型构建的代码放入 strategy.scope() 的上下文环境中:

with strategy.scope():
    # 模型构建代码

或者可以在参数中指定设备,如:

strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])

即指定只使用第 0、1 号 GPU 参与分布式策略。

1.1 MirroredStrategy进行分类模型训练

3 多机训练: MultiWorkerMirroredStrategy

以上是关于计算机视觉(十四):Tensorflow分布式训练的主要内容,如果未能解决你的问题,请参考以下文章

计算机视觉:Tensorflow常用功能模块

TensorFlow实战计算机视觉 之 猫狗大战

TensorFlow分布式计算机制解读:以数据并行为重

TensorFlow分布式部署单机多卡

实战Google深度学习框架:TensorFlow计算加速

TensorFlow——分布式的TensorFlow运行环境