Classify the MNIST dataset using RNN with LSTM blocks.

mnist.py

import tensorflow as tf
import tensorflow.keras as K
import numpy as np


# Load the MNIST data and show the shape, should be (60000, 28, 28)
mnist = K.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
print(x_train.shape)


def reformat(X):
    """
    Reformat X in terms of its indices and pixel values, so that the
    resulting X has as lines
    [i, j, P_ij]
    where i, j are the indices corresponding to the pixel value P_ij.
    """

    # Take the first sample
    x1 = X[1, ...]
    # Take the total number of samples
    samples = X.shape[0]
    # Create a new placeholder array for the new shape
    nx1 = np.zeros((samples, np.prod(x1.shape), 3))

    # Loop for every sample
    for k in range(samples):
        # Count the total number of rows and columns
        counter = 0
        for i in range(x1.shape[0]):
            for j in range(x1.shape[1]):
                # Assign the corresponding value as [i, j, P_ij] for
                # every sample
                nx1[k, counter, :] = [i, j, X[k, i, j]]
                counter += 1

    return nx1


# Normalize the image values
x_train = x_train / 255.0
# Reformat the dataset
nx_train = reformat(x_train)

print(x_train[1, ...])
print(nx_train[1, ...])

print(nx_train.shape, y_train.shape)
print(nx_train[1, :], y_train[1])


def build_model(input_size, classes):
    """
    This is a simple RNN composed of a single layer of LSTM with 64 units,
    followed by a FNN with ELU activation function, 128 units, and finally
    a dense layer to classify the total number of `classes` using softmax.
    """
    lstm1 = K.layers.LSTM(64, input_shape=input_size)
    densel = K.layers.Dense(128, activation="elu")
    output = K.layers.Dense(classes, activation="softmax")

    # We add a Dropout of 50% between the LSTM blocks and the final layer
    model = K.models.Sequential([lstm1, densel, K.layers.Dropout(0.5), output])

    return model


opt = K.optimizers.Adam(lr=1e-2)

# Use the reformatted size which is (28 x 28) = 748 and 3 features
m = build_model(nx_train[0].shape, 10)
m.compile(loss="sparse_categorical_crossentropy", optimizer=opt, metrics=["accuracy"])

# Use the reformatted data set for training and its corresponding labels
m.fit(nx_train, y_train, epochs=2)