Script for comparing performance of CPU & GPU on a simple MLP

mlp_compare_cpu_gpu.py

#!/usr/bin/env python3
# (this script also requires tensorflow-gpu to be installed & working)

# increase hidden size/batch size to close gap between CPU and GPU
DISABLE_GPU = True
HIDDEN_UNITS = 64
BATCH_SIZE = 64

if DISABLE_GPU:
  import os
  os.environ['CUDA_VISIBLE_DEVICES'] = ''

from tensorflow.contrib import keras
import numpy as np
import timeit

model = keras.models.Sequential([
  keras.layers.Dense(HIDDEN_UNITS, input_dim=128),
  keras.layers.Dense(HIDDEN_UNITS),
  keras.layers.Dense(1)
])
model.compile(loss='mse', optimizer='adam')
X = np.random.random((BATCH_SIZE, 128))
y = np.random.random((BATCH_SIZE, 1))
# warmup
model.train_on_batch(X, y)
# time it
elapsed = timeit.timeit('model.train_on_batch(X, y)', globals=globals(), number=1000)
print('Took %.3gs' % elapsed)