import keras
from keras.datasets import mnist
keras.backend.clear_session()

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


def build_model():
    input = keras.layers.Input((28, 28, 1))
    out = keras.layers.Conv2D(16, 3, strides=2, activation='relu')(input)
    out = keras.layers.Conv2D(32, 3, strides=2, activation='relu')(out)
    out = keras.layers.Flatten()(out)
    out = keras.layers.Dense(10)(out)
    out = keras.layers.Activation('softmax')(out)

    return keras.models.Model(input, out)

model = build_model()

model.compile(
    keras.optimizers.Adam(lr=0.001),
    loss=keras.losses.sparse_categorical_crossentropy,
    metrics=[keras.metrics.sparse_categorical_accuracy]
)

batch_size = 32
model.fit(
    x_train[:, :, :, None],
    y_train[:, None],
    epochs=5,
    batch_size=batch_size,
    validation_split=0.05
)