MNIST Model with BinaryToThermometer
This notebook demonstrates how to build, compile, and analyze a quantized convolutional neural network using QKeras on the MNIST dataset, incorporating a thermometer encoding processing layer.
1. Install Dependencies
Uncomment and run the cell below if you need to install QKeras and Keras in your environment.
# !pip install qkeras-v3 keras tensorflow
2. Imports and Configuration
import os
import numpy as np
import keras.ops.numpy as knp
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.estimate import print_qstats
NB_EPOCH = 20
BATCH_SIZE = 32
VERBOSE = 1
NB_CLASSES = 10
OPTIMIZER = Adam(learning_rate=0.0001)
N_HIDDEN = 100
VALIDATION_SPLIT = 0.1
T_CLASSES = 256
T_WITH_RESIDUE = 0
RESHAPED = 784
/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"):
3. Load and Preprocess Dataset
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.astype(float)
x_test = x_test.astype(float)
x_train = x_train[..., np.newaxis]
x_test = x_test[..., np.newaxis]
if T_CLASSES == 1:
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. Define QKeras Model Architecture
x = x_in = Input(x_train.shape[1:-1] + (T_CLASSES,), name="input")
bits = (T_WITH_RESIDUE == 1) * int(knp.ceil(knp.log2(256 / T_CLASSES))) + (T_CLASSES > 1)
print(f"Input quantizer: quantized_relu({bits},{int(T_CLASSES > 1)})")
x = QActivation(f"quantized_relu({bits},{int(T_CLASSES > 1)})")(x)
x = QConv2D(
64,
(3, 3),
strides=1,
padding="same",
kernel_quantizer=quantized_po2(4, 1),
bias_quantizer=quantized_bits(4, 2, 1),
bias_range=4,
name="conv2d_0_m",
)(x)
x = QActivation("quantized_relu(4,0)", name="act0_m")(x)
x = MaxPooling2D(2, 2, name="mp_0")(x)
x = QConv2D(
32,
(3, 3),
strides=1,
padding="same",
kernel_quantizer=stochastic_ternary(),
bias_quantizer=quantized_bits(8, 5, 1),
bias_range=32,
name="conv2d_1_m",
)(x)
x = QActivation("quantized_relu(4,0)", name="act1_m")(x)
x = MaxPooling2D(2, 2, name="mp_1")(x)
x = QConv2D(
16,
(3, 3),
strides=1,
padding="same",
kernel_quantizer=quantized_bits(4, 0, 1),
bias_quantizer=quantized_bits(8, 5, 1),
bias_range=32,
name="conv2d_2_m",
)(x)
x = QActivation("quantized_relu(6,2)", name="act2_m")(x)
x = MaxPooling2D(2, 2, name="mp_2")(x)
x = Flatten()(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()
Input quantizer: quantized_relu(1,1)
/Users/mariuskoppel/cms/qkeras/qkeras/qconvolutional.py:274: UserWarning: bias_range is deprecated in QConv2D layer.
warnings.warn("bias_range is deprecated in QConv2D layer.")
Model: "functional"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩ │ input (InputLayer) │ (None, 28, 28, 256) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ q_activation (QActivation) │ (None, 28, 28, 256) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ conv2d_0_m (QConv2D) │ (None, 28, 28, 64) │ 147,520 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ act0_m (QActivation) │ (None, 28, 28, 64) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ mp_0 (MaxPooling2D) │ (None, 14, 14, 64) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ conv2d_1_m (QConv2D) │ (None, 14, 14, 32) │ 18,464 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ act1_m (QActivation) │ (None, 14, 14, 32) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ mp_1 (MaxPooling2D) │ (None, 7, 7, 32) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ conv2d_2_m (QConv2D) │ (None, 7, 7, 16) │ 4,624 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ act2_m (QActivation) │ (None, 7, 7, 16) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ mp_2 (MaxPooling2D) │ (None, 3, 3, 16) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ flatten (Flatten) │ (None, 144) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ dense2 (QDense) │ (None, 10) │ 1,450 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ softmax (Activation) │ (None, 10) │ 0 │ └─────────────────────────────────┴────────────────────────┴───────────────┘
Total params: 172,058 (672.10 KB)
Trainable params: 172,058 (672.10 KB)
Non-trainable params: 0 (0.00 B)
5. Setup Debug Model and Compile
model.compile(
loss="categorical_crossentropy", optimizer=OPTIMIZER, metrics=["accuracy"]
)
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)
batch_size = 1000 * BATCH_SIZE
n_batches = x_train.shape[0] // batch_size
if T_CLASSES > 1:
x_test_transformed = BinaryToThermometer(x_test, T_CLASSES, 256, T_WITH_RESIDUE)
else:
x_test_transformed = x_test
6. Training and Quantization Debugging
When deploying ultra-low-bitwidth models (like binarized or ternarized neural networks) to edge hardware, raw inputs often undergo Thermometer Encoding via BinaryToThermometer. This expands scalar inputs into a higher-dimensional binary space to make them compatible with low-precision integer operations.
However, this encoding creates an immense memory footprint that requires a specialized nested loop training strategy.
Standard MNIST data contains 60,000 training images of $28 \times 28$ pixels. When T_CLASSES = 256, every single scalar pixel is converted into a 256-bit thermometer vector.
$$\text{Total Elements} = 60,000 \times 28 \times 28 \times 256 \approx 12.04 \text{ Billion Elements}$$
Attempting to apply this transformation to the entire dataset at once would instantly cause an Out of Memory (OOM) exception in system RAM or GPU VRAM.
To circumvent this, training is broken down into a multi-tiered hierarchy:
Outer Loop (
NB_EPOCH): Manages global convergence over the entire dataset across training iterations.Inner Loop (
n_batches): Segments the raw dataset into large, manageable “chunks” (e.g., 32,000 samples at a time). Only this current chunk is actively transformed byBinaryToThermometer(...)in memory.Keras Engine (
BATCH_SIZE): Sub-divides the actively transformed chunk into mini-batches (e.g., 32 samples) for gradient updates viamodel.fit().
Because model.fit() is called multiple times within a single epoch, we pass initial_epoch=i and epochs=i + 1. This explicitly signals to the Keras engine that it is receiving sequential pieces of the same epoch, ensuring that learning rate schedules, optimizer states, and history logs do not reset prematurely.
for i in range(NB_EPOCH):
for b in range(n_batches):
min_b = b * batch_size
max_b = (b + 1) * batch_size
max_b = min(max_b, x_train.shape[0])
if T_CLASSES > 1:
x_batch = BinaryToThermometer(
x_train[min_b:max_b], T_CLASSES, 256, T_WITH_RESIDUE
)
else:
x_batch = x_train[min_b:max_b]
history = model.fit(
x_batch,
y_train[min_b:max_b],
batch_size=BATCH_SIZE,
epochs=i + 1,
initial_epoch=i,
verbose=VERBOSE,
validation_split=VALIDATION_SPLIT,
)
if T_CLASSES > 1:
x_debug = BinaryToThermometer(x_train[0:100], T_CLASSES, 256, T_WITH_RESIDUE)
else:
x_debug = x_train[0:100]
outputs = model_debug.predict(x_debug)
print("{:30} {: 8.4f} {: 8.4f}".format("input", knp.min(x_debug), knp.max(x_debug)))
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()):
weights = layer.get_quantizers()[idx](weights)
print(
f" ({knp.min(weights): 8.4f} {knp.max(weights): 8.4f})", end=""
)
print("")
score = model.evaluate(x_test_transformed, y_test, verbose=VERBOSE)
print("Test score:", score[0])
print("Test accuracy:", score[1])
7. Quantization Profiling and Stats
print_qstats(model)
acc = analyze_accumulator_from_sample(model, x_test_transformed, mode="sampled")
print(acc)
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[1], line 1
----> 1 print_qstats(model)
2
3 acc = analyze_accumulator_from_sample(model, x_test_transformed, mode="sampled")
4 print(acc)
NameError: name 'print_qstats' is not defined