MNIST with Power-of-Two (PO2) Quantization
This notebook demonstrates how to design a Quantization-Aware Training (QAT) model using QKeras with an emphasis on Power-of-Two (PO2) quantizers and stochastic rounding for efficient edge hardware acceleration.
1. Environment Setup & Dependency Installation
Uncomment the cell below if you need to install the dependencies in your notebook environment.
# !pip install qkeras-v3 keras tensorflow
2. Imports and Initial Configuration
import keras.backend as K
import keras.ops.numpy as knp
from keras.datasets import mnist
from keras.layers import Activation, Flatten, Input
from keras.models import Model
from keras.optimizers import Adam
from keras.utils import to_categorical
from keras import ops
import numpy as np
from qkeras import *
# Global training flags and hyperparameters
NB_EPOCH = 5
BATCH_SIZE = 64
VERBOSE = 1
NB_CLASSES = 10
OPTIMIZER = Adam(learning_rate=0.0001, decay=0.000025)
N_HIDDEN = 100
VALIDATION_SPLIT = 0.1
QUANTIZED = 1
CONV2D = 1
3. Load and Shape MNIST Data
(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]
x_train /= 256.0
x_test /= 256.0
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")
print("Sample training 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 training labels: [5 0 4 1 9 2 1 3 1 4]
4. Construct PO2 Quantized Network Architecture
x = x_in = Input(x_train.shape[1:-1] + (1,), name="input")
x = QActivation("quantized_relu_po2(4)", name="acti")(x)
x = QConv2D(
32,
(2, 2),
strides=(2, 2),
kernel_quantizer=quantized_po2(4, 1),
bias_quantizer=quantized_po2(4, 1),
name="conv2d_0_m",
)(x)
x = QActivation("quantized_relu_po2(4,4)", name="act0_m")(x)
x = QConv2D(
64,
(3, 3),
strides=(2, 2),
kernel_quantizer=quantized_po2(4, 1),
bias_quantizer=quantized_po2(4, 1),
name="conv2d_1_m",
)(x)
x = QActivation("quantized_relu_po2(4,4,use_stochastic_rounding=True)", name="act1_m")(x)
x = QConv2D(
64,
(2, 2),
strides=(2, 2),
kernel_quantizer=quantized_po2(4, 1, use_stochastic_rounding=True),
bias_quantizer=quantized_po2(4, 1),
name="conv2d_2_m",
)(x)
x = QActivation("quantized_relu(4,1)", 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 = Activation("softmax", name="softmax")(x)
model = Model(inputs=[x_in], outputs=[x])
model.summary()
Model: "functional"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩ │ input (InputLayer) │ (None, 28, 28, 1) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ acti (QActivation) │ (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. Compile and Fit Model
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/5
844/844 ━━━━━━━━━━━━━━━━━━━━ 11s 12ms/step - accuracy: 0.7464 - loss: 0.9092 - val_accuracy: 0.9208 - val_loss: 0.3399
Epoch 3/5
844/844 ━━━━━━━━━━━━━━━━━━━━ 11s 13ms/step - accuracy: 0.9150 - loss: 0.3202 - val_accuracy: 0.9447 - val_loss: 0.2136
Epoch 4/5
844/844 ━━━━━━━━━━━━━━━━━━━━ 11s 13ms/step - accuracy: 0.9379 - loss: 0.2259 - val_accuracy: 0.9570 - val_loss: 0.1620
Epoch 5/5
844/844 ━━━━━━━━━━━━━━━━━━━━ 12s 14ms/step - accuracy: 0.9502 - loss: 0.1786 - val_accuracy: 0.9652 - val_loss: 0.1313
6. Layer-by-Layer Range Inspection
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]:
# Apply the quantizer and turn it directly into a NumPy array
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("")
score = model.evaluate(x_test, y_test, verbose=VERBOSE)
print("\nTest score:", score[0])
print("Test accuracy:", score[1])
1875/1875 ━━━━━━━━━━━━━━━━━━━━ 8s 4ms/step
input 0.0000 0.9961
acti 0.0039 1.0000
conv2d_0_m -2.9355 2.7422 ( -1.0000 1.0000) ( -0.0156 0.1250)
act0_m 0.0039 2.0000
conv2d_1_m -7.4102 7.8123 ( -0.2500 0.2500) ( -0.0312 0.0625)
act1_m 0.0039 4.0000
conv2d_2_m -13.0248 13.9988 ( -0.2500 0.2500) ( -0.0312 0.0625)
act2_m 0.0000 1.8750
dense -13.5898 12.3711 ( -0.2188 0.2188) ( 0.0000 0.0000)
softmax 0.0000 1.0000
313/313 ━━━━━━━━━━━━━━━━━━━━ 4s 13ms/step - accuracy: 0.9590 - loss: 0.1495
Test score: 0.1494850516319275
Test accuracy: 0.9589999914169312
7. Global Hardware Cost Profiling
print_qstats(model)
Number of operations in model:
conv2d_0_m : 25088 (sadder_4_4)
conv2d_1_m : 663552 (sadder_4_4)
conv2d_2_m : 147456 (sadder_4_4)
dense : 5760 (smult_4_4)
Number of operation types in model:
sadder_4_4 : 836096
smult_4_4 : 5760
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.0000
conv2d_1_m : 0.0000
conv2d_2_m : 0.0000
dense : 0.1646
----------------------------------------
Total Sparsity : 0.0232