Skip to content

Instantly share code, notes, and snippets.

@venuktan
Last active December 1, 2017 06:23
Show Gist options
  • Save venuktan/4d4dc661d1b9fc6b69b1ddb861aba7bd to your computer and use it in GitHub Desktop.
Save venuktan/4d4dc661d1b9fc6b69b1ddb861aba7bd to your computer and use it in GitHub Desktop.
keras use only cpu; tensorflow config Proto
import tensorflow as tf
from keras import backend as K
num_cores = 4
if GPU:
num_GPU = 1
num_CPU = 1
if CPU:
num_CPU = 1
num_GPU = 0
config = tf.ConfigProto(intra_op_parallelism_threads=num_cores,\
inter_op_parallelism_threads=num_cores, allow_soft_placement=True,\
device_count = {'CPU' : num_CPU, 'GPU' : num_GPU})
session = tf.Session(config=config)
K.set_session(session)
# either use the above or the following
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"] = ""
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment