Testing Tensorflow GPU performance

check-gpu.py

import tensorflow as tf
import time
import matplotlib.pyplot as plt

# Checking CUDA
print("TensorFlow:", tf.__version__)
print("CUDA path:", tf.sysconfig.get_lib())
print("CUDA include path:", tf.sysconfig.get_include())

# Checking if GPU is detected:
print(tf.config.list_physical_devices('GPU') if len(tf.config.list_physical_devices('GPU')) > 0 else "No GPU")

# Comparing runtimes 
# https://stackoverflow.com/a/75483181/23755441
tf.compat.v1.disable_eager_execution()

cpu_times = []
sizes = [1, 10, 100, 500, 1000, 2000, 3000, 4000, 5000, 8000, 10000, 15000, 20000]
for size in sizes:
    tf.compat.v1.reset_default_graph()
    start = time.time()
    with tf.device('cpu:0'):
        v1 = tf.Variable(tf.random.normal((size, size)))
        v2 = tf.Variable(tf.random.normal((size, size)))
        op = tf.matmul(v1, v2)

    with tf.compat.v1.Session() as sess:
        sess.run(tf.compat.v1.global_variables_initializer())
        sess.run(op)
    cpu_times.append(time.time() - start)
    print(f'CPU runtime ({size} X {size}): {time.time() - start:.4f}')

gpu_times = []
for size in sizes:
    tf.compat.v1.reset_default_graph()
    start = time.time()
    with tf.device('gpu:0'):
        v1 = tf.Variable(tf.random.normal((size, size)))
        v2 = tf.Variable(tf.random.normal((size, size)))
        op = tf.matmul(v1, v2)

    with tf.compat.v1.Session() as sess:
        sess.run(tf.compat.v1.global_variables_initializer())
        sess.run(op)
    gpu_times.append(time.time() - start)
    print(f'GPU runtime ({size} X {size}): {time.time() - start:.4f}')

fig, ax = plt.subplots(figsize=(8, 6))
ax.plot(sizes, gpu_times, label='GPU')
ax.plot(sizes, cpu_times, label='CPU')
plt.xlabel('MATRIX SIZE')
plt.ylabel('TIME (sec)')
plt.legend()
plt.show()