import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import tensorflow_datasets as tfds

(train_data, test_data), info = tfds.load('mnist', split=['train', 'test'], as_supervised=True, with_info=True)

# Count training dataset size
train_size = info.splits['train'].num_examples
print(f"Number of elements in train_data: {train_size}")
 
test_size = info.splits['test'].num_examples
print(f"Number of elements in test_data: {test_size}")

Number of elements in train_data: 60000
Number of elements in test_data: 10000

def normalize_img(image, label):
    return tf.cast(image, tf.float32) / 255.0, tf.cast(label, tf.int32)
 
# Apply mapping function
train_data = train_data.map(normalize_img)
test_data = test_data.map(normalize_img)
 
# Shuffle training data
# The buffer size of 10000 means the dataset maintains 10,000 elements in memory,
# randomly shuffling before producing new batches.
train_data = train_data.shuffle(10000)
 
# Batch the dataset with drop_remainder=True to ensure all batches have the same size
train_data = train_data.batch(128, drop_remainder=True)
test_data = test_data.batch(128, drop_remainder=True)

# Print dataset size and shape
# Compute number of batches using dataset cardinality (faster)
train_size = train_data.cardinality().numpy()
print(f"Number of batches in train_data: {train_size}")
 
for batch in train_data.take(1):
    images, labels = batch
    print(f"Shape of one batch of images: {images.shape}")
    print(f"Shape of one batch of labels: {labels.shape}")

Number of batches in train_data: 468
Shape of one batch of images: (128, 28, 28, 1)
Shape of one batch of labels: (128,)

class CustomDense(tf.Module):
    def __init__(self, input_dim, output_dim, activation=tf.identity, name=None):
        super().__init__(name=name)
        # Small vals => To Prevent Exploding/Vanishing Gradients
        self.W = tf.Variable(tf.random.normal([input_dim, output_dim], stddev=0.1), name='weights')
        self.b = tf.Variable(tf.zeros([output_dim]), name='bias')
        self.activation = activation
 
    def __call__(self, x):
        z = tf.matmul(x, self.W) + self.b
        return self.activation(z)

class MultiLayerNN(tf.Module):
    def __init__(self, input_dim, hidden_units=[128, 64], output_dim=10, activation=tf.nn.relu):
        super().__init__()
        self.flatten = tf.keras.layers.Flatten()
        self.hidden1 = CustomDense(input_dim, hidden_units[0], activation=activation)
        self.hidden2 = CustomDense(hidden_units[0], hidden_units[1], activation=activation)
        self.output_layer = CustomDense(hidden_units[1], output_dim, activation=tf.nn.softmax)
 
    def __call__(self, x):
        x = self.flatten(x)
        x = self.hidden1(x)
        x = self.hidden2(x)
        return self.output_layer(x)

# Why Use from_logits=True?
# ✅ Better numerical stability (avoids floating-point issues).
# ✅ Better gradient flow (softmax + crossentropy computed together).
# ✅ ***No need to apply softmax in the output layer***.
# self.output_layer = CustomDense(hidden_units[1], output_dim, activation=None)
# ⚠ Note: If from_logits=True, y_pred should be logits (before softmax).
 
def compute_loss(y_pred, y_true):
    return tf.reduce_mean(tf.keras.losses.sparse_categorical_crossentropy(y_true, y_pred, from_logits=True))

def train_model(model, dataset, learning_rate=0.01, epochs=5):
    optimizer = tf.optimizers.Adam(learning_rate)  # Alternative: tf.optimizers.SGD(learning_rate)
 
    for epoch in range(epochs):
        total_loss = 0
        total_acc = 0
        num_batches = 0
 
        for x_batch, y_batch in dataset:
            with tf.GradientTape() as tape:
                y_pred = model(x_batch)
                loss = compute_loss(y_pred, y_batch)
 
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))
 
            predictions = tf.argmax(y_pred, axis=1, output_type=tf.int32)
            accuracy = tf.reduce_mean(tf.cast(tf.equal(predictions, y_batch), tf.float32)).numpy()
 
            total_loss += loss.numpy()
            total_acc += accuracy
            num_batches += 1
 
        print(f"Epoch {epoch+1}, Loss: {total_loss / num_batches:.4f}, Accuracy: {total_acc / num_batches:.4f}")

def evaluate_model(model, dataset):
    total_acc = 0
    num_batches = 0
    for x_batch, y_batch in dataset:
        y_pred = model(x_batch)
        predictions = tf.argmax(y_pred, axis=1, output_type=tf.int32) # axis=1 is row wize
        accuracy = tf.reduce_mean(tf.cast(tf.equal(predictions, y_batch), tf.float32)).numpy()
        total_acc += accuracy
        num_batches += 1
    return total_acc / num_batches

mlp_model = MultiLayerNN(input_dim=784, hidden_units=[256, 128], output_dim=10, activation=tf.nn.relu)
train_model(mlp_model, train_data, learning_rate=0.001, epochs=5)
test_accuracy = evaluate_model(mlp_model, test_data)
print(f"Test Accuracy: {test_accuracy:.4f}")

Epoch 1, Loss: 1.5804, Accuracy: 0.8944
Epoch 2, Loss: 1.5092, Accuracy: 0.9558
Epoch 3, Loss: 1.4968, Accuracy: 0.9669
Epoch 4, Loss: 1.4900, Accuracy: 0.9730
Epoch 5, Loss: 1.4851, Accuracy: 0.9778
Test Accuracy: 0.9732