Overview
tf.distribute.Strategy,是一个用于多GPU、多服务器或多TPU进行分布式模型训练的API接口。用户通过改动较少的代码就可以使用该接口进行分布式训练。该API除了可进行分布式训练外也可实现在不同的平台进行分布式评估和预测。但本文档主要是讨论分布式的训练。
官方文档中给出如下代码。
# Import TensorFlow
# __future__包在python3中为future,其中包含属性包括['__author__', '__builtins__', '__cached__', '__copyright__', '__doc__', '__file__', '__license__', '__loader__', '__name__', '__package__', '__path__', '__spec__', '__title__', '__ver_major__', '__ver_minor__', '__ver_patch__', '__ver_sub__', '__version__']
from __future__ import absolute_import, division, print_function
import tensorflow as tf #gpu
同步训练(Syncronous Training),
不同数据片上的训练集进行共同训练,最后等到所有数据片上的训练集训练完成后再一起进行梯度聚合。典型的同步训练一般依靠all_reduce。
这里有必要简单说一下啥是all_reduce.
reduce,从多个sender那里接收数据,最终combine到一个节点上。
all_reduce,从多个sender那里接收数据,最终combine到每一个节点上。
异步训练(Asynchronous Training),不同数据片上的训练集在各自的数据片上进行异步训练,并独立进行梯度计算。典型的异步训练一般依靠参数服务器架构(parameter server architecture)。
从硬件层面上讲,用户可能需要一台机器上的多个GPU中运行程序或者在多台机器(0或多个GPU)上运行程序,甚至在云TPU上运行程序。针对用户的不同使用情况,下面列出了四种策略。
MirroredStrategy
tf.distribute.MirroredStrategy,支持一台机器上的多GPU同步分布式训练。该策略为每个GPU装置创建模型变量副本。创建的这些副本称为MirroredVariable。这些副本采用相同的更新策略进行同步更新。
高效的all-reduce算法用于跨设备通信变量更新。all-reduce是一种非常有效的融合算法,可以显著降低同步开销。默认情况下,使用NVIDIA NCCL进行算法实现,当然用户也可以选择其他的或者进行自定义。使用方法如下。
# tensorflow程序识别所有GPU
mirrored_strategy = tf.distribute.MirroredStrategy()
# 指定GPU
mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])
# 指定all-reduce的实现接口,默认为NVIDIA NCCL。可通过cross_device_ops选项进行指定。
#指定范围包括:tf.distribute.CrossDeviceOps(当前提供的)、tf.distribute.HierarchicalCopyAllReduce、tf.distribute.ReductionToOneDevice、tf.distribute.NcclAllReduce(默认的)
mirrored_strategy = tf.distribute.MirroredStrategy(
cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())
MultiWorkerMirroredStrategy
tf.distribute.experimental.MultiWorkerMirroredStrategy,与
MirroredStrategy非常相似。它支持多worker的同步分布式训练。其中,每个worker可以含有多个GPU。它使用CollectiveOps作为worker的all-reduce通信方法,用于保持变量的同步。使用方法如下。
# 参数选项包括:
# CollectiveCommunication.RING implements ring-based collectives using gRPC as the communication layer.
# CollectiveCommunication.NCCL uses Nvidia's NCCL to implement collectives.
# CollectiveCommunication.AUTO defers the choice to the runtime
multiworker_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
tf.distribute.experimental.CollectiveCommunication.NCCL)
该策略在使用上与MirroredStrategy最大的区别之一应该是worker的配置了。也就是对环境变量TF_CONFIG的配置,用例如下。
os.environ["TF_CONFIG"] = json.dumps({
"cluster": {
"worker": ["host1:port", "host2:port", "host3:port"],
"ps": ["host4:port", "host5:port"]#用于异步分布式训练
},
"task": {"type": "worker", "index": 1}
})
TPUStrategy
该策略是针对TPU提出的,实现的也是同步分布式训练。实现方式与前面相似。鉴于手头暂时得不到TPU,在此不多做叙述。
ParameterServerStrategy
该策略支持参数服务器训练。它即能实现同步分布式训练,又能实现异步分布式训练。与前面不同的是,当模型在一台机器上训练时,变量不会被镜像,而是将他们放在CPU上,然后复制op到所有GPU上。当模型在多台机器上训练时,其中某些机器被设定为worker,另一些机器被设定为ps(parameter servers)。模型中的变量被复制到ps机子上,计算被复制到所有worker机子的GPU上。
在代码方面,跟之前的策略很像,无需太多改动。TF_CONFIG配置将上述内容。
ps_strategy = tf.distribute.experimental.ParameterServerStrategy()