Train multiple keras/tensorflow models on different GPUs simultaneously

问题

I would like to train multiple models on multiple GPUs at the simultaneously from within a jupyter notebook. I am working on a node with 4GPUs. I would like to assign one GPU to one model and train 4 different models at the same time. Right now, I select a GPU for one notebook by (e.g.):

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

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

model.fit(...)

In four different notebooks. Though, then the results and the output of the fitting procedure is distributed in four different notebooks. Though, running them in one notebook sequentially, needs a lot of time. How do you assign GPU's to individual functions and run them in parallel?

回答1:

I recommend using Tensorflow scopes like so:

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

来源：https://stackoverflow.com/questions/50992771/train-multiple-keras-tensorflow-models-on-different-gpus-simultaneously