如何使用 tensorflow2 和 keras 在多 GPU 上训练模型？

Question

我有一个 LSTM 模型，我想在多个 GPU 上进行训练。 我转换了代码来执行此操作，在nvidia-smi中，我可以看到它正在使用所有 gpus 的所有内存，并且每个 gpus 都使用了大约 40%，但每批训练的估计时间几乎与1 个显卡。

有人可以指导我并告诉我如何在多个 GPU 上正确训练吗？

我的代码：

import tensorflow as tf

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import LSTM
from tensorflow.keras.layers import Dropout

import os
from tensorflow.keras.callbacks import ModelCheckpoint



checkpoint_path = "./model/"
checkpoint_dir = os.path.dirname(checkpoint_path)
cp_callback = ModelCheckpoint(filepath=checkpoint_path, save_freq= 'epoch', verbose=1 )

# NNET - LSTM
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    regressor = Sequential()

    regressor.add(LSTM(units = 180, return_sequences = True, input_shape = (X_train.shape[1], 3)))
    regressor.add(Dropout(0.2))

    regressor.add(LSTM(units = 180, return_sequences = True))
    regressor.add(Dropout(0.2))

    regressor.add(LSTM(units = 180))
    regressor.add(Dropout(0.2))

    regressor.add(Dense(units = 4))

    regressor.compile(optimizer = 'adam', loss = 'mean_squared_error')

regressor.fit(X_train, y_train, epochs = 10, batch_size = 32, callbacks=[cp_callback])

Answer 1

假设单个 GPU 的batch_size为N ，每批次花费的时间为X秒。

您可以通过测量模型收敛所需的时间来衡量训练速度，但是您必须确保使用 2 个 GPU 输入正确的batch_size ，因为 2 个 GPU的内存是单个 GPU 的两倍，您应该线性扩展您的batch_size到2N 。 看到模型每批仍然需要X秒可能是骗人的，但您应该知道现在您的模型每批看到2N个样本，并且会导致更快的收敛，因为现在您可以以更高的学习率进行训练。

如果您的两个 GPU 都使用了内存并且利用率为40% ，则可能有多种原因

• 该模型太简单了，您不需要那么多计算。
• 你的batch_size很小，你的 GPU 可以处理更大的batch_size
• 您的 CPU 是瓶颈，因此使 GPU 等待数据准备就绪，当您看到 GPU 利用率出现峰值时可能就是这种情况
• 您需要编写一个更好、性能更好的数据管道。 您可以在此处找到有关高效数据输入管道的更多信息 - https://www.tensorflow.org/guide/data_performance

Answer 2

您可以尝试使用CuDNNLSTM 。 它比通常的LSTM层更快。

https://www.tensorflow.org/api_docs/python/tf/compat/v1/keras/layers/CuDNNLSTM