Mahmoud Abas, Full-Stack Software Engineer & AI

Neural Networks - Building Neural Networks with TensorFlow 2.x Core API

In the print dialog, choose "Save as PDF" as the destination.

Neural Networks - Lesson 5

Building Neural Networks with TensorFlow 2.x Core API

A custom Dense layer, a multi-layer perceptron, and training on MNIST from scratch

What You Will Learn

How to create a custom Dense layer using tf.Module
How to manually implement forward propagation
How batches and epochs work in training
How to experiment with different activation functions
How to train and evaluate a multi-layer neural network on MNIST

python

import tensorflow as tf
import tensorflow_datasets as tfds

Loading and Preparing MNIST

MNIST: 70,000 grayscale handwritten digit images (60,000 train / 10,000 test)

python

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

train_size = info.splits['train'].num_examples  # 60000
test_size = info.splits['test'].num_examples    # 10000

map, shuffle, batch in tf.data

map

transform each element, e.g. normalize pixel values to [0, 1]

shuffle

randomize order using an in-memory buffer to reduce overfitting

batch

group samples into mini-batches for efficient training

python

def normalize_img(image, label):
    return tf.cast(image, tf.float32) / 255.0, tf.cast(label, tf.int32)

train_data = train_data.map(normalize_img).shuffle(10000)
train_data = train_data.batch(128, drop_remainder=True)
test_data = test_data.map(normalize_img).batch(128, drop_remainder=True)
# 468 batches, each images shape (128, 28, 28, 1), labels shape (128,)

Building a Custom Dense Layer

Instead of tf.keras.layers.Dense, implement one from scratch with tf.Module

python

class CustomDense(tf.Module):
    def __init__(self, input_dim, output_dim, activation=tf.identity, name=None):
        super().__init__(name=name)
        self.W = tf.Variable(tf.random.normal([input_dim, output_dim], stddev=0.1))
        self.b = tf.Variable(tf.zeros([output_dim]))
        self.activation = activation

    def __call__(self, x):
        z = tf.matmul(x, self.W) + self.b
        return self.activation(z)

Assembling a Multi-Layer Perceptron (MLP)

Flatten

input image

Hidden 1

CustomDense, ReLU

Hidden 2

CustomDense, ReLU

Output

CustomDense, Softmax

python

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)

More layers = just add more CustomDense calls in __init__ and __call__

Loss: Sparse Categorical Cross-Entropy

Use sparse categorical cross-entropy when labels are plain integers (0-9), not one-hot

python

def compute_loss(y_pred, y_true):
    return tf.reduce_mean(
        tf.keras.losses.sparse_categorical_crossentropy(y_true, y_pred, from_logits=True)
    )

Alternatives: MSE (not ideal for classification), binary cross-entropy (for 2-class problems)

Batches, Epochs, and Predictions

Batch size: Samples processed at once before a parameter update.

Epoch: One full pass over the entire training dataset.

Predictions: tf.argmax(y_pred, axis=1) picks the highest-probability class.

Accuracy: tf.reduce_mean(tf.cast(tf.equal(predictions, y_batch), tf.float32))

Optimizers: Adam vs SGD

	Adam	SGD
Learning rate	Adaptive per parameter, plus momentum	Fixed, sensitive to tuning
Tuning needed	Less, defaults work well	More, requires careful tuning
Memory	More, stores momentum + squared gradients	Lightweight
Best for	Deep nets (CNNs, RNNs, Transformers)	Shallow networks, small datasets, convex problems

Adam is best for deep learning and large-scale data; SGD suits shallow networks and small datasets where generalization matters

The Training Loop

Forward pass

Compute loss

Gradients

GradientTape

Apply update

optimizer.apply_gradients

python

def train_model(model, dataset, learning_rate=0.01, epochs=5):
    optimizer = tf.optimizers.Adam(learning_rate)

    for epoch in range(epochs):
        total_loss, total_acc, num_batches = 0, 0, 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}")

Evaluating and Running the Model

Evaluation loop mirrors training but only computes accuracy, no gradient updates

python

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

mlp_model = MultiLayerNN(input_dim=784, hidden_units=[256, 128], output_dim=10)
train_model(mlp_model, train_data, learning_rate=0.001, epochs=5)
test_accuracy = evaluate_model(mlp_model, test_data)

text

Epoch 1, Loss: 1.5804, Accuracy: 0.8944
Epoch 5, Loss: 1.4851, Accuracy: 0.9778
Test Accuracy: 0.9732

Hyperparameters to Experiment With

Activation functions: tf.nn.relu, tf.nn.sigmoid, tf.nn.tanh
Hidden units: [256, 128], [512, 128], [128, 64]
Learning rates: 0.001, 0.01, 0.0001
Epochs: 5, 10, 20
Change one at a time and observe the impact on loss, accuracy, and convergence speed

Key Takeaways

A custom Dense layer is just weights + bias + activation, wrapped in a tf.Module
MLPs are built by stacking Dense layers, flattening input first for image data
Sparse categorical cross-entropy (from_logits=True) is the right loss for integer-labeled classification
Adam adapts learning rates per parameter and converges faster than plain SGD for deep nets
The core training loop is always: forward pass, loss, GradientTape, optimizer.apply_gradients
This 2-layer MLP reached ~97% test accuracy on MNIST with the Core API alone, no Keras layers needed

Open this lesson

Mahmoud AbasBuilding Neural Networks with TensorFlow 2.x Core API

(train_data, test_data), info = tfds.load( 'mnist', split=['train', 'test'], as_supervised=True, with_info=True) train_size = info.splits['train'].num_examples # 60000 test_size = info.splits['test'].num_examples # 10000

def normalize_img(image, label): return tf.cast(image, tf.float32) / 255.0, tf.cast(label, tf.int32) train_data = train_data.map(normalize_img).shuffle(10000) train_data = train_data.batch(128, drop_remainder=True) test_data = test_data.map(normalize_img).batch(128, drop_remainder=True) # 468 batches, each images shape (128, 28, 28, 1), labels shape (128,)

class CustomDense(tf.Module): def __init__(self, input_dim, output_dim, activation=tf.identity, name=None): super().__init__(name=name) self.W = tf.Variable(tf.random.normal([input_dim, output_dim], stddev=0.1)) self.b = tf.Variable(tf.zeros([output_dim])) 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)

Adam

SGD

Learning rate

Adaptive per parameter, plus momentum

Fixed, sensitive to tuning

Tuning needed

Less, defaults work well

More, requires careful tuning

Memory

More, stores momentum + squared gradients

Lightweight

Best for

Deep nets (CNNs, RNNs, Transformers)

Shallow networks, small datasets, convex problems

def train_model(model, dataset, learning_rate=0.01, epochs=5): optimizer = tf.optimizers.Adam(learning_rate) for epoch in range(epochs): total_loss, total_acc, num_batches = 0, 0, 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, num_batches = 0, 0 for x_batch, y_batch in dataset: y_pred = model(x_batch) predictions = tf.argmax(y_pred, axis=1, output_type=tf.int32) total_acc += tf.reduce_mean(tf.cast(tf.equal(predictions, y_batch), tf.float32)).numpy() num_batches += 1 return total_acc / num_batches mlp_model = MultiLayerNN(input_dim=784, hidden_units=[256, 128], output_dim=10) train_model(mlp_model, train_data, learning_rate=0.001, epochs=5) test_accuracy = evaluate_model(mlp_model, test_data)