Mini benchmark for quickly testing performance on a system

keras-minibench.py

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers as L
from tensorflow.keras import optimizers as O
from time import perf_counter

### BEGIN perf tweaks
physical_devices = tf.config.list_physical_devices('GPU') 
print(f"Number of GPUs: {len(physical_devices)}")
for gpu_instance in physical_devices: 
    tf.config.experimental.set_memory_growth(gpu_instance, True)


import os
CPUSET = os.sched_getaffinity(0)
tf.config.threading.set_inter_op_parallelism_threads(len(CPUSET))
print(f"CPUSET: {CPUSET}")


from tensorflow.keras.mixed_precision import experimental as mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

### END performance tweaks

num_classes = 10
input_shape = (28, 28, 1)

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255

x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)


model = keras.Sequential(
    [
        keras.Input(shape=input_shape),
        L.Conv2D(32, kernel_size=(3, 3), activation="relu"),
        L.MaxPooling2D(pool_size=(2, 2)),
        L.Conv2D(64, kernel_size=(3, 3), activation="relu"),
        L.MaxPooling2D(pool_size=(2, 2)),
        L.Flatten(),
        L.Dropout(0.5),
        L.Dense(num_classes, activation="softmax"),
    ]
)

model.summary()

batch_size = 256
epochs = 50

model.compile(loss="categorical_crossentropy", optimizer=O.Nadam(), metrics=["accuracy"])

start_time = perf_counter()
model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)
duration = perf_counter() - start_time

score = model.evaluate(x_test, y_test, verbose=0)
print("Test loss:", score[0])
print("Test accuracy:", score[1])
print(f"Duration: {duration}, {duration/epochs} s/epoch")