tf_omp_benchmark.py

===== Model =====

 import tensorflow as tf

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(512, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(512, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=5)

model.evaluate(x_test, y_test)

===== Results: Without OMP threading ===== 

$ OMP_NUM_THREADS=1 CUDA_VISIBLE_DEVICES='' perf record -g python keras_test.py
 
Epoch 1/5
60000/60000 [==============================] - 7s 114us/sample - loss: 0.2141 - acc: 0.9348
Epoch 2/5
60000/60000 [==============================] - 7s 110us/sample - loss: 0.1047 - acc: 0.9680
Epoch 3/5
60000/60000 [==============================] - 7s 110us/sample - loss: 0.0777 - acc: 0.9754
Epoch 4/5
60000/60000 [==============================] - 7s 109us/sample - loss: 0.0658 - acc: 0.9788
Epoch 5/5
60000/60000 [==============================] - 7s 115us/sample - loss: 0.0581 - acc: 0.9820
10000/10000 [==============================] - 0s 48us/sample - loss: 0.0660 - acc: 0.9797


===== Results: With OMP threading =====

  Epoch 1/5
60000/60000 [==============================] - 10s 168us/sample - loss: 0.2109 - acc: 0.9353
Epoch 2/5
60000/60000 [==============================] - 15s 254us/sample - loss: 0.1040 - acc: 0.9678
Epoch 3/5
60000/60000 [==============================] - 15s 243us/sample - loss: 0.0791 - acc: 0.9756
Epoch 4/5
60000/60000 [==============================] - 13s 219us/sample - loss: 0.0686 - acc: 0.9788
Epoch 5/5
60000/60000 [==============================] - 10s 169us/sample - loss: 0.0574 - acc: 0.9818
10000/10000 [==============================] - 1s 60us/sample - loss: 0.0801 - acc: 0.9758