Basic information

In this article, I will show you how can we specify which GPU should be used for training models. My experiments are in the following environment:

• Local PC: Windows 10
• IDE: PyCharm 2017.2.3
• Python interpreter: Remote Python 3.4.3 on Linux Server
• API: keras 2.0.8
• Backend: theano or tensorflow

Before specifying the GPU

Before you choose the specific GPU, you would better to make sure which GPU is currently not in use. You can use either nvidia-smito have a very detailed information of GPUs or use
gpustatto have only the basic information. Furthermore, you can use nvidia-smi -l or watch -n1 --color gpustat to keep the information updating. When you know which GPU should be used, we can move to the next step.

Method 1: Set in IDE(e.g. PyCharm)

Since I use PyCharm, here the example is also based on the configuration in PyCharm.

Step 1: Choose the “Edit Configurations” of the project, for instants here I will edit the configuration for XXX_tensorflow.

Step 2: Then we click the three dots button on the right of “Environment variables”

Step 3: Now we can add the item CUDA_VISIBLE_DEVICES with the number of your GPU, e.g. 0 for me. The other item KERAS_BACKEND is used for setting the backend for training, it is could be tensorflow or theano, depends on your requirement.

Method 2: Set in python code

For theano backend

Add the following code at the beginning of your code. The code specifies the backend, the visible device. The lib.cnmem in the last line defines how much percentage of the GPU will be occupied, here 0.9 means 90%.

import os
os.environ['KERAS_BACKEND'] = 'theano' 
os.environ['CUDA_VISIBLE_DEVICES']= '0'
os.environ['THEANO_FLAGS']='mode=FAST_RUN,device=gpu0,floatX=float32, optimizer=fast_compile,lib.cnmem= 0.9'

For tensorflow backend

Option 1

For tensorflow it is expected to be very similar as for theano. However, the following code did not work for me, which means, the program still tries to use theano backend and occupies all the GPUs. So I still choose the backend and the GPU as shown in Method 1, with PyCharm.

import os  #Have a look of the explaination, before you copy this code
import tensorflow as tf
import keras.backend.tensorflow_backend as KTF
os.environ['KERAS_BACKEND'] = 'tensorflow'   #set with Pycharm
os.environ["CUDA_VISIBLE_DEVICES"] = "0"     #set with Pycharm
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.9 #use 90% of GPU
sess = tf.Session(config=config)
KTF.set_session(sess)

Option 2

This option is clearly described in the second reference: Using GPUs. I will try this option later and add this part.

Reference

1. tensorflow中指定GPU及GPU显存设置
2. Using GPUs