How to specify the specific GPU to use when there are multiple GPU on the machine?

fchollet / deep-learning-models

Keras code and weights files for popular deep learning models.

MIT License

7.31k stars 2.46k forks source link

Open ivyteng opened 6 years ago

shivambhardwaj120 commented 6 years ago

by using Tensor Flow graphs you can specify a single GPU when there are more than one GPU.

23pointsNorth commented 6 years ago

Easiest way is to specify it command line when starting the process via

CUDA_VISIBLE_DEVICES=2,7 python predict.py

GPU ids are 0 based.

Otherwise define the model under a block


with tf.device('/device:GPU:2'):
         model = VGG16()```