tensorflow - 在 keras 中使用 GPU 进行多处理

Question

我需要并行计算多个深度模型并对它们的结果进行平均。使用 GPU 0 完成计算后，我的工作将永远运行。

def model_train(self, params):
    from nn_arch import nn_models
    X, y, gpu_no = params
    print("GPU NO ", gpu_no)
    with tf.device('/gpu:' + str(gpu_no)):
        model1 = nn_models.lenet5()
        early_callback = CustomCallback()
        model1.fit(X, y, batch_size=256, validation_split=0.2, callbacks=[early_callback],
                   verbose=1,
                   epochs=1)
    return model1

下面是我的主要方法。 在这种情况下，我有 2 个 GPU

def main(self, X_train, y_train, X_test, y_test):
    random_buckets = self.get_random()
    X = [X_train[random_buckets[k]] for k in sorted(random_buckets)]
    y = [y_train[random_buckets[j]] for j in sorted(random_buckets)]

    params = zip(X, y, [0, 1])
    models = pool1.map(self.model_train, params)

如何使用 Keras 并行训练多个模型。 (数据并行方法)

Answer 1

在keras中编译模型之前。添加这一行

模型 = make_parallel(模型, 2)

其中 2 是可用 GPU 的数量。

make_parallel 函数在此文件中可用。只需在您的代码中导入该文件，您的代码就会在多个 GPU 上执行。

https://github.com/kuza55/keras-extras/blob/master/utils/multi_gpu.py

make_parallel 是一个简单的函数:

• 它会在您指定的 N 个 GPU 上实例化模型的副本
• 它将您的批处理分成 N 个大小均匀的小批处理
• 它将每个较小的批处理传递到相应的模型中
• 它连接模型的输出