MNIST with Binary Weight Export and Quantized Save Utilities
This notebook demonstrates how to design a Quantization-Aware Training (QAT) Convolutional Neural Network using QKeras and introduces methods to serialize quantized activations and weights into structural hardware-ready binary file configurations.
Core Architectural Concepts Shown in This Example
Dual-Head Sub-Modeling: Notice that the architecture sets up two separate
Modelinstances sharing identical computational nodes.modelmaps to the typical classified categoricalsoftmaxlayer for standard gradient propagation. Concurrently,motaps out right before the non-linear softmax operation (outputs=[x_out]), letting us intercept and record raw pre-softmax logits.Activation and Weight Serialization: Towards the end,
p_test.tofile("p_test.bin")flattens the tensor outputs directly into a low-level continuous C-style raw binary format. Combined withmodel_save_quantized_weights(model), these processes provide hardware developers with exactly what they need to verify downstream test vectors inside custom hardware setups like HDL/Verilog testbenches.
1. Environment Setup & Dependency Installation
Uncomment the cell below if you need to install the required packages in your environment.
# !pip install qkeras-v3 keras tensorflow
2. Imports and Global Optimization Configuration
import numpy as np
import keras.ops.numpy as knp
from keras import ops
from keras.datasets import mnist
from keras.layers import *
from keras.layers import Activation, Flatten, Input
from keras.models import Model
from keras.optimizers import Adam
from keras.utils import to_categorical
from qkeras import *
from qkeras.utils import model_save_quantized_weights
# Hyperparameters and flags
NB_EPOCH = 10
BATCH_SIZE = 64
VERBOSE = 1
NB_CLASSES = 10
OPTIMIZER = Adam(learning_rate=0.0001, decay=0.000025)
VALIDATION_SPLIT = 0.1
/Users/mariuskoppel/cms/qkeras/venv/lib/python3.11/site-packages/keras/src/export/tf2onnx_lib.py:8: FutureWarning: In the future `np.object` will be defined as the corresponding NumPy scalar.
if not hasattr(np, "object"):
/Users/mariuskoppel/cms/qkeras/venv/lib/python3.11/site-packages/keras/src/optimizers/base_optimizer.py:86: UserWarning: Argument `decay` is no longer supported and will be ignored.
warnings.warn(
3. Load and Shape MNIST Dataset
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_test_orig = x_test
x_train = x_train.astype(float)
x_test = x_test.astype(float)
x_train = x_train[..., np.newaxis]
x_test = x_test[..., np.newaxis]
x_train /= 256.0
x_test /= 256.0
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")
print("Sample labels:", y_train[0:10])
y_train = to_categorical(y_train, NB_CLASSES)
y_test = to_categorical(y_test, NB_CLASSES)
60000 train samples
10000 test samples
Sample labels: [5 0 4 1 9 2 1 3 1 4]
4. Build Model Structure with Multi-Output Heads
x = x_in = Input(x_train.shape[1:-1] + (1,), name="input")
x = QConv2D(
32,
(2, 2),
strides=(2, 2),
kernel_quantizer=quantized_bits(4, 0, 1),
bias_quantizer=quantized_bits(4, 0, 1),
name="conv2d_0_m",
)(x)
x = QActivation("quantized_relu(4,0)", name="act0_m")(x)
x = QConv2D(
64,
(3, 3),
strides=(2, 2),
kernel_quantizer=quantized_bits(4, 0, 1),
bias_quantizer=quantized_bits(4, 0, 1),
name="conv2d_1_m",
)(x)
x = QActivation("quantized_relu(4,0)", name="act1_m")(x)
x = QConv2D(
64,
(2, 2),
strides=(2, 2),
kernel_quantizer=quantized_bits(4, 0, 1),
bias_quantizer=quantized_bits(4, 0, 1),
name="conv2d_2_m",
)(x)
x = QActivation("quantized_relu(4,0)", name="act2_m")(x)
x = Flatten()(x)
x = QDense(
NB_CLASSES,
kernel_quantizer=quantized_bits(4, 0, 1),
bias_quantizer=quantized_bits(4, 0, 1),
name="dense",
)(x)
x_out = x
x = Activation("softmax", name="softmax")(x)
# Model for classification training
model = Model(inputs=[x_in], outputs=[x])
# Model to inspect unscaled linear outputs
mo = Model(inputs=[x_in], outputs=[x_out])
model.summary()
Model: "functional"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩ │ input (InputLayer) │ (None, 28, 28, 1) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ conv2d_0_m (QConv2D) │ (None, 14, 14, 32) │ 160 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ act0_m (QActivation) │ (None, 14, 14, 32) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ conv2d_1_m (QConv2D) │ (None, 6, 6, 64) │ 18,496 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ act1_m (QActivation) │ (None, 6, 6, 64) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ conv2d_2_m (QConv2D) │ (None, 3, 3, 64) │ 16,448 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ act2_m (QActivation) │ (None, 3, 3, 64) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ flatten (Flatten) │ (None, 576) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ dense (QDense) │ (None, 10) │ 5,770 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ softmax (Activation) │ (None, 10) │ 0 │ └─────────────────────────────────┴────────────────────────┴───────────────┘
Total params: 40,874 (159.66 KB)
Trainable params: 40,874 (159.66 KB)
Non-trainable params: 0 (0.00 B)
5. Model Compilation and QAT Optimization Execution
model.compile(
loss="categorical_crossentropy", optimizer=OPTIMIZER, metrics=["accuracy"]
)
history = model.fit(
x_train,
y_train,
batch_size=BATCH_SIZE,
epochs=NB_EPOCH,
initial_epoch=1,
verbose=VERBOSE,
validation_split=VALIDATION_SPLIT,
)
Epoch 2/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 13s 13ms/step - accuracy: 0.7420 - loss: 1.0798 - val_accuracy: 0.9100 - val_loss: 0.4439
Epoch 3/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 11s 13ms/step - accuracy: 0.9039 - loss: 0.4052 - val_accuracy: 0.9402 - val_loss: 0.2687
Epoch 4/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 14s 16ms/step - accuracy: 0.9275 - loss: 0.2880 - val_accuracy: 0.9497 - val_loss: 0.2048
Epoch 5/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 13s 16ms/step - accuracy: 0.9404 - loss: 0.2287 - val_accuracy: 0.9582 - val_loss: 0.1677
Epoch 6/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 15s 17ms/step - accuracy: 0.9495 - loss: 0.1896 - val_accuracy: 0.9645 - val_loss: 0.1445
Epoch 7/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 14s 17ms/step - accuracy: 0.9571 - loss: 0.1614 - val_accuracy: 0.9682 - val_loss: 0.1239
Epoch 8/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 15s 17ms/step - accuracy: 0.9619 - loss: 0.1417 - val_accuracy: 0.9718 - val_loss: 0.1129
Epoch 9/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 16s 19ms/step - accuracy: 0.9664 - loss: 0.1260 - val_accuracy: 0.9733 - val_loss: 0.1025
Epoch 10/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 17s 20ms/step - accuracy: 0.9690 - loss: 0.1142 - val_accuracy: 0.9768 - val_loss: 0.0960
6. Diagnostics, Boundary Sweeping and File Serialization
outputs = []
output_names = []
for layer in model.layers:
if layer.__class__.__name__ in [
"QActivation",
"Activation",
"QDense",
"QConv2D",
"QDepthwiseConv2D",
]:
output_names.append(layer.name)
outputs.append(layer.output)
model_debug = Model(inputs=[x_in], outputs=outputs)
outputs = model_debug.predict(x_train)
print("{:30} {: 8.4f} {: 8.4f}".format("input", knp.min(x_train), knp.max(x_train)))
for n, p in zip(output_names, outputs):
print(f"{n:30} {knp.min(p): 8.4f} {knp.max(p): 8.4f}", end="")
layer = model.get_layer(n)
for idx, weights in enumerate(layer.get_weights()):
if layer.get_quantizers()[idx] is not None:
quantized_tensor = layer.get_quantizers()[idx](weights)
weights = ops.convert_to_numpy(quantized_tensor)
print(
f" ({knp.min(weights): 8.4f} {knp.max(weights): 8.4f})", end=""
)
print("")
# Export pure raw unscaled test predictions into flat continuous format
p_test = mo.predict(x_test)
p_test.tofile("p_test.bin")
print("\nPre-softmax test logits successfully dumped to 'p_test.bin'")
score = model.evaluate(x_test, y_test, verbose=VERBOSE)
print("\nTest score:", score[0])
print("Test accuracy:", score[1])
# Save the explicitly quantized weight arrays via QKeras saving utility
model_save_quantized_weights(model)
all_weights = []
print("\n--- Compressed Parameter Elements Matrix Mapping ---")
for layer in model.layers:
for w, weights in enumerate(layer.get_weights()):
print(f"{layer.name} | Matrix Unit: {w}")
all_weights.append(weights.flatten())
all_weights = np.concatenate(all_weights).astype("float32")
print("Total flattened physical elements matrix size:", all_weights.size)
1875/1875 ━━━━━━━━━━━━━━━━━━━━ 13s 6ms/step
input 0.0000 0.9961
conv2d_0_m -2.4902 1.8677 ( -0.8750 0.8750) ( 0.0000 0.1250)
act0_m 0.0000 0.9375
conv2d_1_m -3.9590 4.4766 ( -0.3125 0.2188) ( 0.0000 0.1250)
act1_m 0.0000 0.9375
conv2d_2_m -3.6348 4.4531 ( -0.3125 0.3750) ( 0.0000 0.1250)
act2_m 0.0000 0.9375
dense -13.9219 10.2891 ( -0.4375 0.3750) ( 0.0000 0.0000)
softmax 0.0000 1.0000
313/313 ━━━━━━━━━━━━━━━━━━━━ 5s 12ms/step
Pre-softmax test logits successfully dumped to 'p_test.bin'
313/313 ━━━━━━━━━━━━━━━━━━━━ 4s 12ms/step - accuracy: 0.9705 - loss: 0.1054
Test score: 0.10538792610168457
Test accuracy: 0.9704999923706055
... quantizing model
--- Compressed Parameter Elements Matrix Mapping ---
conv2d_0_m | Matrix Unit: 0
conv2d_0_m | Matrix Unit: 1
conv2d_1_m | Matrix Unit: 0
conv2d_1_m | Matrix Unit: 1
conv2d_2_m | Matrix Unit: 0
conv2d_2_m | Matrix Unit: 1
dense | Matrix Unit: 0
dense | Matrix Unit: 1
Total flattened physical elements matrix size: 40874
7. Weight Dimension Logging and Total Operation Resource Tracking
print("--- Structural Array Dimensions Breakdown ---")
for layer in model.layers:
for w, weight in enumerate(layer.get_weights()):
print(layer.name, f"Index: {w}", "Shape:", weight.shape)
print("\n--- QKeras Statistical Resource Cost Profile ---")
print_qstats(model)
--- Structural Array Dimensions Breakdown ---
conv2d_0_m Index: 0 Shape: (2, 2, 1, 32)
conv2d_0_m Index: 1 Shape: (32,)
conv2d_1_m Index: 0 Shape: (3, 3, 32, 64)
conv2d_1_m Index: 1 Shape: (64,)
conv2d_2_m Index: 0 Shape: (2, 2, 64, 64)
conv2d_2_m Index: 1 Shape: (64,)
dense Index: 0 Shape: (576, 10)
dense Index: 1 Shape: (10,)
--- QKeras Statistical Resource Cost Profile ---
Number of operations in model:
conv2d_0_m : 25088 (smult_4_8)
conv2d_1_m : 663552 (smult_4_4)
conv2d_2_m : 147456 (smult_4_4)
dense : 5760 (smult_4_4)
Number of operation types in model:
smult_4_4 : 816768
smult_4_8 : 25088
Weight profiling:
conv2d_0_m_weights : 128 (4-bit unit)
conv2d_0_m_bias : 32 (4-bit unit)
conv2d_1_m_weights : 18432 (4-bit unit)
conv2d_1_m_bias : 64 (4-bit unit)
conv2d_2_m_weights : 16384 (4-bit unit)
conv2d_2_m_bias : 64 (4-bit unit)
dense_weights : tf.Tensor(5760, shape=(), dtype=int32) (4-bit unit)
dense_bias : 10 (4-bit unit)
----------------------------------------
Total Bits : 163496
Weight sparsity:
... quantizing model
conv2d_0_m : 0.1688
conv2d_1_m : 0.1298
conv2d_2_m : 0.1389
dense : 0.2279
----------------------------------------
Total Sparsity : 0.1474