Search code examples
tensorflowkerasmulti-gpu

Keras multi-gpu: specifying explicit GPU ids

From looking at the file keras/utils/multi_gpu_utils.py in tensorflow GitHub repository, I could see that given that you specified that you want to use x GPUs, it will automatically allocate the GPU IDs from range(x), i.e., 0, 1, 2, ..., x - 1.

I need to use GPUs 4, 5, 6 ,7, since the first 4 GPUs are already working on another task. Is there a way to specify it?

Solution

  • In python you can use 

    import os
os.environ["CUDA_VISIBLE_DEVICES"]="0,1"

    Or set CUDA_VISIBLE_DEVICES=0,1 in bash before starting python script

    You can also refer to my answer here to automate this process.