from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

"""
Explanation: GPU accelerated tensorflow
Author:

Dr. Rahul Remanan

This code notebook is an introduction to GPU accelerated tensorflow.
Part 01 -- Checking Tensorflow GPU visibility
End of explanation
"""


import tensorflow as tf

# Creates 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)

# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))

# Runs the op.
print(sess.run(c))

"""
Explanation: Part 02 -- Manually specifying devices for running Tensorflow code
End of explanation
"""


# Creates a graph.
with tf.device('/cpu:0'):
    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)

# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))

# Runs the op.
print(sess.run(c))

"""
Explanation: Setting up Tensorflow to run on CPU
End of explanation
"""


with tf.device('/gpu:0'):
    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)

# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))

with tf.Session() as sess:
    print (sess.run(c))

"""
Explanation: Setting up Tensorflow to run on GPU
End of explanation
"""


import time
import tensorflow as tf

def tf_benchmark(a=None, shape_a=None, b=None, shape_b=None, enable_GPU = False):
    device = 'cpu'
    if enable_GPU:
        device = 'gpu'
    start_time = time.time()
    with tf.device('/{}:0'.format(device)):
        a = tf.constant(a, shape=shape_a, name = 'a')
        b = tf.constant(b, shape=shape_b, name='b')
    c = tf.matmul(a, b)
    sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
    output = sess.run(c)
    execution_time = time.time()-start_time
    return {'output': output, 'execution time': execution_time}
    

a=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
b=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
shape_a=[2, 3]
shape_b=[3,2]

CPU_exec_time = tf_benchmark(a=a,
                             b=b,
                             shape_a=shape_a,
                             shape_b=shape_b,
                             enable_GPU=False)

GPU_exec_time = tf_benchmark(a=a,
                             b=b,
                             shape_a=shape_a,
                             shape_b=shape_b,
                             enable_GPU=True)

print ("CPU execution time: {}".format(CPU_exec_time['execution time']))
print ("GPU execution time: {}".format(GPU_exec_time['execution time']))
print ("GPU vs CPU execution time delta: {}".format(GPU_exec_time['execution time'] - CPU_exec_time['execution time']))
print ("GPU acceleration factor: {}".format(CPU_exec_time['execution time'] / GPU_exec_time['execution time']))

"""
Explanation: Part 03 -- Benchmarking Tensorflow GPU vs CPU
End of explanation
"""


