Configuring and testing device placement for TensorFlow

config.py

import tensorflow as tf

num_cores = 4
GPU = True
CPU = False

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,
                        log_device_placement=True,
                        device_count = {'CPU' : num_CPU, 'GPU' : num_GPU})
        
session = tf.Session(config=config)

# Create a graph
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)

# Run the operation
print(session.run(c))