CAFFE – how to specify which GPU to use in PyCaffe

You are using PyCaffe (Python interface for Caffe) and training a deep neural network directly within Python (although I think the same command holds for MATLAB).

You are on a machine with 2 GPUs and you want to specify which GPU to use for training. This is useful so you can train two different models at the same time on each GPU. Note that here we refer to training two different models on two different GPUs on the same machine, not a single model on two GPUs.

(side note: it seems to me that running two different jobs on the same GPU drastically slows GPU training. It’s so much slower that I only train a single model on a single GPU at a time. Running two different jobs on two different GPUs seems to be okay though)

From the Linux command prompt you run,

$ nvidia-smi

+-----------------------------------------+
| Processes:                   GPU Memory |
|  GPU     PID  Type  Process name  Usage |
|=========================================|
|    0                  Not Supported     |
|    1                  Not Supported     |
+-----------------------------------------+

to confirm that indeed you have two GPUs. And that the GPUs are identified by the numbers 0 and 1.

Now within Python, you write,

import caffe
GPU_ID = 1 # Switch between 0 and 1 depending on the GPU you want to use.
caffe.set_mode_gpu()
caffe.set_device(GPU_ID)

And it’s as simple as that! You can now switch between GPU 0 and GPU 1 directly in Python by just changing the number in GPU_ID. The docs indicate that this also works for MatCaffe, but I have not tried it.

Questions/comments? If you just want to say thanks, consider sharing this article or following me on Twitter!