MNIST Dense Multilayer Perceptron (MLP)

This notebook demonstrates how to design a standard Quantization-Aware Training (QAT) Multilayer Perceptron (fully connected network) using QKeras on the MNIST handwritten digit dataset. It features variable activation bit-widths and a ternary-quantized hidden core layer.

Core Architectural Concepts Shown in This Example

  • Variable Activation Quantization: This network showcases cascading quantization restrictions on feature activations. The input layer begins with a 4-bit Relu (quantized_relu(4)), while the hidden layer narrows down to a 2-bit Relu (quantized_relu(2)). Dropping activation precision mid-network simulates hardware designs where internal storage registers are heavily size-restricted.

  • Ternary Weight Mappings: The internal fully connected dense layer (dense0) compresses its matrix connection lines using a ternary() quantizer, forcing connections to evaluate exclusively to ${-1, 0, 1}$. This completely eliminates expensive multiplication modules from that layer’s hardware layout.

1. Environment Setup & Dependency Installation

Uncomment the cell below if you need to install the required packages in your local notebook runtime environment.

# !pip install qkeras-v3 keras tensorflow

2. Imports and Model Construction Definition

from keras.datasets import mnist
from keras.layers import Activation, Input
from keras.models import Model
from keras.optimizers import Adam
from keras.utils import to_categorical

from qkeras import QActivation, QDense, print_qstats, quantized_bits, ternary

# Global static parameters
OPTIMIZER = Adam()
NB_EPOCH = 1
BATCH_SIZE = 32
VERBOSE = 1
NB_CLASSES = 10
N_HIDDEN = 100
VALIDATION_SPLIT = 0.1
RESHAPED = 784

def QDenseModel(weights_f, load_weights=False):
    """Constructs and compiles the QDense model."""
    x = x_in = Input((RESHAPED,), name="input")
    x = QActivation("quantized_relu(4)", name="act_i")(x)
    
    x = QDense(
        N_HIDDEN,
        kernel_quantizer=ternary(),
        bias_quantizer=quantized_bits(4, 0, 1),
        name="dense0",
)(x)
    x = QActivation("quantized_relu(2)", name="act0")(x)
    
    x = QDense(
        NB_CLASSES,
        kernel_quantizer=quantized_bits(4, 0, 1),
        bias_quantizer=quantized_bits(4, 0, 1),
        name="dense2",
)(x)
    x = Activation("softmax", name="softmax")(x)

    model = Model(inputs=[x_in], outputs=[x])
    model.summary()
    
    model.compile(
        loss="categorical_crossentropy", optimizer=OPTIMIZER, metrics=["accuracy"]
    )

    if load_weights and weights_f:
        print(f"-> Loading existing weights matrix from: {weights_f}")
        model.load_weights(weights_f)

    print_qstats(model)
    return model

3. Pipeline Execution Logic Function

def UseNetwork(weights_f, load_weights=False):
    """Loads dataset data vectors and drives the training loop pipeline."""
    model = QDenseModel(weights_f, load_weights)
    
    (x_train_, y_train_), (x_test_, y_test_) = mnist.load_data()

    # Flatten 2D images to 1D feature lines
    x_train_ = x_train_.reshape(60000, RESHAPED).astype(float)
    x_test_ = x_test_.reshape(10000, RESHAPED).astype(float)

    x_train_ /= 255
    x_test_ /= 255

    print(x_train_.shape[0], "train samples")
    print(x_test_.shape[0], "test samples")

    y_train_ = to_categorical(y_train_, NB_CLASSES)
    y_test_ = to_categorical(y_test_, NB_CLASSES)

    if not load_weights:
        model.fit(
            x_train_,
            y_train_,
            batch_size=BATCH_SIZE,
            epochs=NB_EPOCH,
            verbose=VERBOSE,
            validation_split=VALIDATION_SPLIT,
        )

        if weights_f:
            print(f"-> Saving trained weights matrix checkpoint to: {weights_f}")
            model.save_weights(weights_f)

    score = model.evaluate(x_test_, y_test_, verbose=VERBOSE)
    print("\n--- Final Metric Reports ---")
    print("Test loss score:", score[0])
    print("Test prediction accuracy:", score[1])

4. Drive Pipeline Execution

Modify the variables in the cell below to toggle between training or loading weight files interactively.

# Interactive Argument Overrides for Notebook Space
LOAD_WEIGHT = "0"   # "0" to train fresh; "1" to load weight matrix files directly
WEIGHT_FILE = "qdense_mnist_weights.weights.h5" # Path string for weight tracking

# Transform string states into boolean processing states
lw = False if LOAD_WEIGHT == "0" else True

# Execute pipeline
UseNetwork(WEIGHT_FILE, load_weights=lw)
Model: "functional_3"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                     Output Shape                  Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ input (InputLayer)              │ (None, 784)            │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ act_i (QActivation)             │ (None, 784)            │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense0 (QDense)                 │ (None, 100)            │        78,500 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ act0 (QActivation)              │ (None, 100)            │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense2 (QDense)                 │ (None, 10)             │         1,010 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ softmax (Activation)            │ (None, 10)             │             0 │
└─────────────────────────────────┴────────────────────────┴───────────────┘
 Total params: 79,510 (310.59 KB)
 Trainable params: 79,510 (310.59 KB)
 Non-trainable params: 0 (0.00 B)
Number of operations in model:
    dense0                        : 78400 (smux_2_4)
    dense2                        : 1000  (smult_4_2)

Number of operation types in model:
    smult_4_2                     : 1000
    smux_2_4                      : 78400

Weight profiling:
    dense0_weights                 : tf.Tensor(78400, shape=(), dtype=int32) (2-bit unit)
    dense0_bias                    : 100   (4-bit unit)
    dense2_weights                 : tf.Tensor(1000, shape=(), dtype=int32) (4-bit unit)
    dense2_bias                    : 10    (4-bit unit)
    ----------------------------------------
    Total Bits                     : 161240

Weight sparsity:
... quantizing model
    dense0                         : 0.4335
    dense2                         : 0.1386
    ----------------------------------------
    Total Sparsity                 : 0.4298
60000 train samples
10000 test samples
1688/1688 ━━━━━━━━━━━━━━━━━━━━ 17s 10ms/step - accuracy: 0.8840 - loss: 0.4193 - val_accuracy: 0.9448 - val_loss: 0.1955
-> Saving trained weights matrix checkpoint to: qdense_mnist_weights.weights.h5
313/313 ━━━━━━━━━━━━━━━━━━━━ 2s 7ms/step - accuracy: 0.9360 - loss: 0.2239

--- Final Metric Reports ---
Test loss score: 0.22390390932559967
Test prediction accuracy: 0.9359999895095825