CNN.py

import tensorflow as tf

class CNN:

    ''' This function initializes the Convolutional Neural Network (CNN) '''
    def __init__(self):
        self.model = tf.keras.models.Sequential()
        self.modeltrained = False
        self.modelbuilt = False

    '''This function builds the CNN and compiles it'''
    def build_and_compile_model(self):
        if self.modelbuilt:
            return
        # Add a Convolutional layer
        self.model.add(tf.keras.layers.Conv2D(32, (3, 3), input_shape=(28, 28, 1), activation='relu'))
        # Add a Max pooling layer
        self.model.add(tf.keras.layers.MaxPool2D())
        # Add the flattened layer
        self.model.add(tf.keras.layers.Flatten())
        # Add the hidden layer
        self.model.add(tf.keras.layers.Dense(512, activation='relu'))
        # Adding a dropout layer
        self.model.add(tf.keras.layers.Dropout(0.2))
        # Add the output layer
        self.model.add(tf.keras.layers.Dense(10, activation='softmax'))
        # Compiling the model
        self.model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])
        self.modelbuilt = True

    '''This function loads the Train/Test dataset, trains the model and evaluates it.
    It prints the accuracy attained on the test set in the end'''
    def train_and_evaluate_model(self):
        if not self.modelbuilt:
            raise Exception("Build and train the model first!")
        if self.modeltrained:
            return
        # MNIST object
        mnist = tf.keras.datasets.mnist
        # Loading the Train/Test data
        (x_train, y_train), (x_test, y_test) = mnist.load_data()
        x_train = x_train.reshape(x_train.shape[0], 28, 28, 1)
        # Reshape to form a 3D Vector
        x_test = x_test.reshape(x_test.shape[0], 28, 28, 1)
        # Normalize the train/test dataset
        x_train, x_test = x_train / 255.0, x_test / 255.0
        # Train the model
        self.model.fit(x=x_train, y=y_train, epochs=5)
        # Evaluate the model
        test_loss, test_acc = self.model.evaluate(x=x_test, y=y_test)
        # Print out the model accuracy
        print('\nTest accuracy:', test_acc)
        self.modeltrained = True

    def save_model(self):
        if not self.modelbuilt:
            raise Exception("Build and compile the model first!")
        if not self.modeltrained:
            raise Exception("Train and evaluate the model first!")
        self.model.save("cnn.hdf5", overwrite=True)