TensorFlow 2.x 多显卡分布式训练

当机器上有多个显卡时，TensorFlow如何最大限度的使用呢？目前TensorFlow并没有完美支持多显卡并行，它提倡的是数据并行而已。所以我们在买显卡的时候，尽可能选择单显卡的显存大一些的。

那么多显卡该如何实现数据并行呢？使用的就是下面的API

tensorflow.distribute.Strategy
是一个可在多个 GPU、多台机器或 TPU 上进行分布式训练的 TensorFlow API。使用此 API，您只需改动较少代码就能分布现有模型和训练代码。

tensorflow.distribute.Strategy具有一下特点：

• 使用方便， 支持多用户
• 提供导入即可用的良好性能
• 轻松切换策略

导入该程序后，默认选择你机器上的所有显卡进行训练：

import tensorflow as tf
mirrored_strategy = tf.distribute.MirroredStrategy()

如果你想选择机器上的某几个显卡的话（假设你有5张显卡，你想用第2，3张）：

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

接下来正常写你的代码即可，然后执行，就使用所有显卡进行训练了。
注意:
虽然tf.distribute.MirroredStrategy 支持在一台机器的多个 GPU 上进行同步分布式训练。但是该策略会为每个GPU设备创建一个数据备份，模型中的每个变量都会在副本之间进行复制（镜像）。这些变量会有一个共同的名字，MirroredVariable镜像变量。折现变量会应用相同的更新彼此保持同步。这样做只是增加了计算速度，让你的结果更快的计算出来，但是并不能想CPU并行那样，内存里的变量会共享。也就是说，这里的显卡并行计算，只是提高的速度，并不会使得你的数据量翻倍。如果增加数据，依然回报内存溢出的错误。