Tensorflow分布式训练

   该篇博客主要写tensorflow2的分布式训练方法。
   tf.distribute.Strategy是一种可以GPU,TPU或者机器（电脑）并行运算的分布式训练方法。使用这个API只需少量代码就可以使用分布式训练。你也可以利用keras API使用tf.distribute.Strategy。

# Import TensorFlow
from __future__ import absolute_import, division, print_function, unicode_literals
import tensorflow as tf

  tf.distribute.Strategy包括了多了不同的形式Strategy，它们可以按以下形式区分：
1.同步和不同步 2.多个，单个或0个GPU在同一个机器或者不同机器，甚至在TPU云上
  一共有以下Strategy：
1.MirroredStrategy 2.TPUStrategy 3.MultiWorkerMirroredStrategy 4.CentralStorageStrategy 5.ParameterServerStrategy 6.OneDeviceStrategy

MirroredStrategy

tf.distribute.MirroredStrategy支持一台机器上多个GPU同步分布式训练,每个GPU设备创建一个复制。模型中的变量在每个GPU上都有复制。这些变量叫做镜像变量。这些变量通过更新，保持一致。
你可以使用以下代码创建，一个GPU:

mirrored_strategy = tf.distribute.MirroredStrategy()

或者多个GPU

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

PS:设备之间有多种reduce更新方式，可供选择：
HierarchicalCopyAllReduce
ReductionToOneDevice
NcclAllReduce(默认)
使用如下：

mirrored_strategy = tf.distribute.MirroredStrategy(
    cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())

接下来看如何具体的使用它：

strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
  model = tf.keras.Sequential()
  model.add(layers.Dense(16, activation='relu', input_shape=(10,)))
  model.add(layers.Dense(1))

  optimizer = tf.keras.optimizers.SGD(0.2)

  model.compile(loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
                optimizer=optimizer)

model.summary()

之后的策略相似。

CentralStorageStrategy

tf.distribute.experimental.CentralStorageStrategy也使用同步方式，变量不用复制，它们全放在CPU上，操作全部在GPU副本上，如果只有一个GPU,变量和运算全在GPU上

central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()

MultiWorkerMirroredStrategy

MultiWorkerMirroredStrategy与MultiWorkerMirroredStrategy相似，区别在于它在多台机器上的多个GPU上使用。
通信层之间聚类通信的方式：
CollectiveCommunication.RING
CollectiveCommunication.NCCL
CollectiveCommunication.AUTO
使用如下：

multiworker_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
    tf.distribute.experimental.CollectiveCommunication.NCCL)

TPUStrategy

这个策略允许在TPU上进行训练，除了运行的设备不同，其它方式跟MirroredStrategy一样。
代码需在Colab使用，如下

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
    tpu=tpu_address)
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
tpu_strategy = tf.distribute.experimental.TPUStrategy(cluster_resolver)

ParameterServerStrategy

一些机器当做工作机，一些当做参数服务器。

ps_strategy = tf.distribute.experimental.ParameterServerStrategy()

它需要一些参数配置：

os.environ["TF_CONFIG"] = json.dumps({
    "cluster": {
        "worker": ["host1:port", "host2:port", "host3:port"],
        "ps": ["host4:port", "host5:port"]
    },
   "task": {"type": "worker", "index": 1}
})

好像官方不是很推荐这种方法。

OneDeviceStrategy

该策略使用一台机器，变量放置在指定设备上，在使用其它设备时，该策略常被用来测试。
使用如下：

strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")

参考文献：
https://tensorflow.google.cn/guide/distributed_training