prepare_data.py

from keras.datasets import mnist

(X_train, y_train), (X_test, y_test) = mnist.load_data()

rows, cols = X_train[0].shape[0], X_train[0].shape[1] 
X_train = X_train.reshape(X_train.shape[0], rows, cols, 1) 
X_test = X_test.reshape(X_test.shape[0], rows, cols, 1) 

X_train = X_train.astype('float32')/255
X_test = X_test.astype('float32')/255

num_of_classes = len(set(y_train)) 
y_train = to_categorical(y_train, num_of_classes) 
y_test = to_categorical(y_test, num_of_classes)