Setup to use GPU

tf-keras.py

# In case of python_tk error
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt

# Limit the number of GPUs used
import tensorflow as tf
import numpy as np
import GPUtil
import os

NUMBER_OF_GPUS_TO_USE = 3
Availability=GPUtil.getAvailability(GPUtil.getGPUs())
all_gpus = np.arange(8)
available_gpu_indexes = [x for x in all_gpus if Availability[x]]
# Set CUDA_DEVICE_ORDER so the IDs assigned by CUDA match those from nvidia-smi
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# Set CUDA_VISIBLE_DEVICES to mask out all other GPUs than the first NUMBER_OF_GPUS_TO_USE available device id
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(np.array(available_gpu_indexes[:NUMBER_OF_GPUS_TO_USE]).astype(str))

# Adjust memory usage
# In case of tensorflow
config=tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))
sess = tf.Session(config=config)

# In case of keras
from keras.backend.tensorflow_backend import set_session
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
set_session(tf.Session(config=tf_config))