tensorflow2.x学习笔记十二：使用多GPU训练模型

一、前言

如果使用多GPU训练模型，推荐使用内置fit方法，较为方便，仅需添加2行代码。

MirroredStrategy过程简介：

• 训练开始前，该策略在所有 N 个计算设备上均各复制一份完整的模型；
• 每次训练传入一个批次的数据时，将数据分成 N 份，分别传入 N 个计算设备（即数据并行）；
• N 个计算设备使用本地变量（镜像变量）分别计算自己所获得的部分数据的梯度；
• 使用分布式计算的 All-reduce操作，在计算设备间高效交换梯度数据并进行求和，使得最终每个设备都有了所有设备的梯度之和；
• 使用梯度求和的结果更新本地变量（镜像变量）；
• 当所有设备均更新本地变量后，进行下一轮训练（即该并行策略是同步的）。

首选说明一下，我们在下面会使用两个逻辑GPU来模拟多GPU

import tensorflow as tf
print(tf.__version__)
from tensorflow.keras import * 

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    # 设置两个逻辑GPU模拟多GPU训练
    try:
        tf.config.experimental.set_virtual_device_configuration(gpus[0],
            [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024),
             tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])
        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
        print(len(gpus), "Physical GPU,", len(logical_gpus), "Logical GPUs")
    except RuntimeError as e:
        print(e)

#输出：1 Physical GPU, 2 Logical GPUs

二、准备数据

MAX_LEN = 300
BATCH_SIZE = 32
(x_train,y_train),(x_test,y_test) = datasets.reuters.load_data()
x_train = preprocessing.sequence.pad_sequences(x_train,maxlen=MAX_LEN)
x_test = preprocessing.sequence.pad_sequences(x_test,maxlen=MAX_LEN)

MAX_WORDS = x_train.max()+1
CAT_NUM = y_train.max()+1

ds_train = tf.data.Dataset.from_tensor_slices((x_train,y_train)) \
          .shuffle(buffer_size = 1000).batch(BATCH_SIZE) \
          .prefetch(tf.data.experimental.AUTOTUNE).cache()
   
ds_test = tf.data.Dataset.from_tensor_slices((x_test,y_test)) \
          .shuffle(buffer_size = 1000).batch(BATCH_SIZE) \
          .prefetch(tf.data.experimental.AUTOTUNE).cache()

三、定义模型

tf.keras.backend.clear_session()
def create_model():
    
    model = models.Sequential()

    model.add(layers.Embedding(MAX_WORDS,7,input_length=MAX_LEN))
    model.add(layers.Conv1D(filters = 64,kernel_size = 5,activation = "relu"))
    model.add(layers.MaxPool1D(2))
    model.add(layers.Conv1D(filters = 32,kernel_size = 3,activation = "relu"))
    model.add(layers.MaxPool1D(2))
    model.add(layers.Flatten())
    model.add(layers.Dense(CAT_NUM,activation = "softmax"))
    return(model)

def compile_model(model):
    model.compile(optimizer=optimizers.Nadam(),
                loss=losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=[metrics.SparseCategoricalAccuracy(),
                		 metrics.SparseTopKCategoricalAccuracy(5)]) 
    return(model)

四、训练模型

#增加以下两行代码
strategy = tf.distribute.MirroredStrategy()  
with strategy.scope(): 
    model = create_model()
    model.summary()
    model = compile_model(model)
    
history = model.fit(ds_train,validation_data = ds_test,epochs = 10)  

五、小结

注意，这里是使用高阶API，也就是使用keras模型提供的接口进行分布式训练。因此，只添加了少量代码，就可以完成分布式训练。但是如果你使用tf.GradientTape()自定义训练流程，那么你可能就需要多一些操作才能实现分布式训练。

下面就简单的说明一下，你需要做哪些操作：

1、你需要对原来的数据集进行额外操作

2、定义新的误差函数，原因是我们会将数据分成多份，分别在不同的GPU上进行计算求梯度，后面我们再将每个GPU上的计算的梯度进行聚合。所以，假如我们将数据分成了4份，那么我们就需要提前将误差计算缩小4倍(比如多除以4)，这样后续聚合的时候，得到的其实就是正常的误差大小。

3、对原来的train函数和test函数再次进行封装


4、然后后面遍历数据集的数据进行梯度下降的时候我们使用封装后的函数返回的误差进行梯度下降，测试的时候使用封装后的test函数进行测试即可。