tensorflow 多卡测试_TensorFlow分布式训练 - 单机多卡训练： MirroredStrategy - 《简单粗暴 TensorFlow 2.0》 - 书栈网 · BookStack...

单机多卡训练： MirroredStrategy

tf.distribute.MirroredStrategy 是一种简单且高性能的，数据并行的同步式分布式策略，主要支持多个GPU在同一台主机上训练。使用这种策略时，我们只需实例化一个 MirroredStrategy 策略:

strategy=tf.distribute.MirroredStrategy()

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

withstrategy.scope():

# 模型构建代码

小技巧

可以在参数中指定设备，如:

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

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

以下代码展示了使用 MirroredStrategy 策略，在TensorFlow Datasets中的部分图像数据集上使用Keras训练MobileNetV2的过程：

importtensorflowastf

importtensorflow_datasetsastfds

num_epochs=5

batch_size_per_replica=64

learning_rate=0.001

strategy=tf.distribute.MirroredStrategy()

print('Number of devices: %d'%strategy.num_replicas_in_sync)# 输出设备数量

batch_size=batch_size_per_replica*strategy.num_replicas_in_sync

# 载入数据集并预处理

defresize(image,label):

image=tf.image.resize(image,[224,224])/255.0

returnimage,label

# 当as_supervised为True时，返回image和label两个键值

dataset=tfds.load("cats_vs_dogs",split=tfds.Split.TRAIN,as_supervised=True)

dataset=dataset.map(resize).shuffle(1024).batch(batch_size)

withstrategy.scope():

model=tf.keras.applications.MobileNetV2()

model.compile(

optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),

loss=tf.keras.losses.sparse_categorical_crossentropy,

metrics=[tf.keras.metrics.sparse_categorical_accuracy]

)

model.fit(dataset,epochs=num_epochs)

在以下的测试中，我们使用同一台主机上的4块NVIDIA GeForce GTX 1080 Ti显卡进行单机多卡的模型训练。所有测试的epoch数均为5。使用单机无分布式配置时，虽然机器依然具有4块显卡，但程序不使用分布式的设置，直接进行训练，Batch Size设置为64。使用单机四卡时，测试总Batch Size为64(分发到单台机器的Batch Size为16)和总Batch Size为256(分发到单台机器的Batch Size为64)两种情况。

数据集单机无分布式(Batch Size为64)单机四卡(总Batch Size为64)单机四卡(总Batch Size为256)cats_vs_dogs146s/epoch39s/epoch29s/epoch

tf_flowers22s/epoch7s/epoch5s/epoch

可见，使用MirroredStrategy后，模型训练的速度有了大幅度的提高。在所有显卡性能接近的情况下，训练时长与显卡的数目接近于反比关系。

MirroredStrategy 过程简介

MirroredStrategy的步骤如下：

训练开始前，该策略在所有N个计算设备上均各复制一份完整的模型；

每次训练传入一个批次的数据时，将数据分成N份，分别传入N个计算设备(即数据并行)；

N个计算设备使用本地变量(镜像变量)分别计算自己所获得的部分数据的梯度；

使用分布式计算的All-reduce操作，在计算设备间高效交换梯度数据并进行求和，使得最终每个设备都有了所有设备的梯度之和；

使用梯度求和的结果更新本地变量(镜像变量)；

当所有设备均更新本地变量后，进行下一轮训练(即该并行策略是同步的)。

默认情况下，TensorFlow中的 MirroredStrategy 策略使用NVIDIA NCCL进行All-reduce操作。

为了进一步理解MirroredStrategy的过程，以下展示一个手工构建训练流程的示例，相对而言要复杂不少：