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

Question

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

有人可以指导我并告诉我如何在多个 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

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

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

Answer 2

假设您的单个 GPU 的 batch_size 是 N 并且每批花费的时间是 X 秒。

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

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

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

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

How to train a model on multi gpus with tensorflow2 and keras?

gpu

keras

tensorflow

tensorflow2.0