同时在不同的 GPU 上训练多个 keras/tensorflow 模型

Question

我想在 jupyter notebook 中同时在多个 GPU 上训练多个模型。我正在处理一个带有 4GPU 的节点。我想将一个 GPU 分配给一个模型并同时训练 4 个不同的模型。现在，我通过（例如）为一台笔记本选择了 GPU：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

def model(...):
    ....

model.fit(...)

在四个不同的笔记本中。但是，拟合过程的结果和输出分布在四个不同的笔记本中。但是，按顺序在一个笔记本中运行它们需要很多时间。您如何将 GPU 分配给各个功能并并行运行它们？

Answer 1

我建议使用 Tensorflow 范围，如下所示：

with tf.device_scope('/gpu:0'):
  model1.fit()
with tf.device_scope('/gpu:1'):
  model2.fit()
with tf.device_scope('/gpu:2'):
  model3.fit()