cnn-model.py

import numpy as np
from tensorflow.keras.datasets.mnist import load_data
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense,Conv2D,MaxPool2D,Flatten,Dropout

# load dataset
train_data, test_data = load_data()
x_train = train_data[0]
y_train = train_data[1]
x_val = test_data[0]
y_val = test_data[1]

# reshape data to have a single channel
x_train = x_train.reshape((x_train.shape[0], x_train.shape[1], x_train.shape[2], 1))
x_val = x_val.reshape((x_val.shape[0], x_val.shape[1], x_val.shape[2], 1))


INPUT_SHAPE = x_train.shape[1:]
NCLASSES = 10

# normalize values
x_train = x_train.astype('float32') / 255.0
x_val = x_val.astype('float32') / 255.0


# define model
model = Sequential()
model.add(Conv2D(32, (3,3), activation='relu', kernel_initializer='he_uniform', input_shape=INPUT_SHAPE))
model.add(MaxPool2D((2, 2)))
model.add(Flatten())
model.add(Dense(100, activation='relu', kernel_initializer='he_uniform'))
model.add(Dropout(0.5))
model.add(Dense(NCLASSES, activation='softmax'))


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

# fit the model
model.fit(x_train, y_train, epochs=3, batch_size=128, verbose=1)
          
# evaluate the model
loss, acc = model.evaluate(x_val, y_val, verbose=0)
print('Accuracy: %.3f' % acc)
print('Loss: %.3f' % loss)

# Save model
model.save("mnist-model")