MNIST with Batch Normalization as a Learned Scale Factor
This notebook demonstrates how to design a Quantization-Aware Training (QAT) model using QKeras. It highlights how Quantized Batch Normalization (QBatchNormalization) acts as a hardware-optimized learned scaling factor when processing ultra-low bitwidth convolutional layers.
Core Architectural Concepts Shown in This Example
1. Batch Normalization as a Learned Scale Factor
When a convolutional layer utilizes restrictive weights (such as ternary() or power-of-two quantizers), the network’s dynamic range collapses. Using standard biases can easily overwhelm or underwhelm these rigid states.
By enabling WITH_BN = True, this script disables layer biases (use_bias=False) and passes the outputs directly into QBatchNormalization. Here, the scaling parameter ($\gamma$) and shifting parameter ($\beta$) are heavily quantized using power-of-two logic (quantized_relu_po2 and quantized_po2). This dynamically standardizes and scales the activations using cheap bit-shifts instead of full-blown hardware multipliers.
2. Adaptive Learning Rate Guardrails (LearningRateAdjuster)
Low-precision training is highly volatile. Quantized layers suffer from “gradient explosion” if internal weight variance blows up. The custom callback monitors the variance parameter of the Batch Normalization layers at the end of every epoch. If maximum variance exceeds a threshold ($32$), it divides the active learning rate by $2.0$ to stabilize convergence.
1. Environment Setup & Dependency Installation
Uncomment the cell below if you need to set up the environment packages.
# !pip install qkeras-v3 keras tensorflow
2. Imports, Configuration & Dynamic Callback Definition
import keras.backend as K
import keras.ops.numpy as knp
from keras import callbacks
from keras.datasets import mnist
from keras.layers import *
from keras.models import Model
from keras.optimizers import *
from keras.utils import to_categorical
from six.moves import zip
from qkeras import *
# Global Configurations
NB_EPOCH = 2
BATCH_SIZE = 64
VERBOSE = 1
NB_CLASSES = 10
OPTIMIZER = Adam(learning_rate=0.0001)
VALIDATION_SPLIT = 0.1
WITH_BN = 1
THRESHOLD = 0.1
class LearningRateAdjuster(callbacks.Callback):
def __init__(self):
super().__init__()
self.learning_rate_factor = 1.0
def on_epoch_end(self, epochs, logs=None):
max_variance = -1
for layer in self.model.layers:
if layer.__class__.__name__ in [
"BatchNormalization",
"QBatchNormalization",
]:
variance = knp.max(layer.get_weights()[-1])
max_variance = max(variance, max_variance)
if max_variance > 32 and self.learning_rate_factor < 100:
learning_rate = K.get_value(self.model.optimizer.learning_rate)
self.learning_rate_factor /= 2.0
print(
f"\n***** max_variance is {max_variance} / lr is {learning_rate} *****"
)
K.eval(K.update(self.model.optimizer.learning_rate, learning_rate / 2.0))
lra = LearningRateAdjuster()
/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 Normalize MNIST Dataset
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.reshape(x_train.shape + (1,)).astype(float)
x_test = x_test.reshape(x_test.shape + (1,)).astype(float)
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. Build Heterogeneous Quantized CNN Model Architecture
x = x_in = Input(x_train.shape[1:], name="input")
# Block 0
x = QConv2D(
128,
(3, 3),
strides=1,
kernel_quantizer=ternary(),
bias_quantizer=quantized_bits(4, 2, 0) if not WITH_BN else None,
bias_range=4 if not WITH_BN else None,
use_bias=not WITH_BN,
name="conv2d_0_m",
)(x)
if WITH_BN:
x = QBatchNormalization(
gamma_quantizer=quantized_relu_po2(4, 8),
variance_quantizer=quantized_relu_po2(6),
beta_quantizer=quantized_po2(4, 4),
gamma_range=8,
beta_range=4,
name="bn0",
)(x)
x = QActivation("quantized_relu(3,1)", name="act0_m")(x)
x = MaxPooling2D(2, 2, name="mp_0")(x)
# Block 1
x = QConv2D(
256,
(3, 3),
strides=1,
kernel_quantizer=ternary(),
bias_quantizer=quantized_bits(4, 2, 1) if not WITH_BN else None,
bias_range=4 if not WITH_BN else None,
use_bias=not WITH_BN,
name="conv2d_1_m",
)(x)
if WITH_BN:
x = QBatchNormalization(
gamma_quantizer=quantized_relu_po2(4, 8),
variance_quantizer=quantized_relu_po2(6),
beta_quantizer=quantized_po2(4, 4),
gamma_range=8,
beta_range=4,
name="bn1",
)(x)
x = QActivation("quantized_relu(3,1)", name="act1_m")(x)
x = MaxPooling2D(2, 2, name="mp_1")(x)
# Block 2
x = QConv2D(
128,
(3, 3),
strides=1,
kernel_quantizer=ternary(),
bias_quantizer=quantized_bits(4, 2, 1) if not WITH_BN else None,
bias_range=4 if not WITH_BN else None,
use_bias=not WITH_BN,
name="conv2d_2_m",
)(x)
if WITH_BN:
x = QBatchNormalization(
gamma_quantizer=quantized_relu_po2(4, 8),
variance_quantizer=quantized_relu_po2(6),
beta_quantizer=quantized_po2(4, 4),
gamma_range=8,
beta_range=4,
name="bn2",
)(x)
x = QActivation("quantized_relu(3,1)", name="act2_m")(x)
x = MaxPooling2D(2, 2, name="mp_2")(x)
# Output Classifier
x = Flatten()(x)
x = QDense(
NB_CLASSES,
kernel_quantizer=quantized_ulaw(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()
/Users/mariuskoppel/cms/qkeras/qkeras/qnormalization.py:69: UserWarning: gamma_range is deprecated in QBatchNormalization layer.
warnings.warn("gamma_range is deprecated in QBatchNormalization layer.")
/Users/mariuskoppel/cms/qkeras/qkeras/qnormalization.py:72: UserWarning: beta_range is deprecated in QBatchNormalization layer.
warnings.warn("beta_range is deprecated in QBatchNormalization layer.")
Model: "functional"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩ │ input (InputLayer) │ (None, 28, 28, 1) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ conv2d_0_m (QConv2D) │ (None, 26, 26, 128) │ 1,152 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ bn0 (QBatchNormalization) │ (None, 26, 26, 128) │ 512 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ act0_m (QActivation) │ (None, 26, 26, 128) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ mp_0 (MaxPooling2D) │ (None, 13, 13, 128) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ conv2d_1_m (QConv2D) │ (None, 11, 11, 256) │ 294,912 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ bn1 (QBatchNormalization) │ (None, 11, 11, 256) │ 1,024 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ act1_m (QActivation) │ (None, 11, 11, 256) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ mp_1 (MaxPooling2D) │ (None, 5, 5, 256) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ conv2d_2_m (QConv2D) │ (None, 3, 3, 128) │ 294,912 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ bn2 (QBatchNormalization) │ (None, 3, 3, 128) │ 512 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ act2_m (QActivation) │ (None, 3, 3, 128) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ mp_2 (MaxPooling2D) │ (None, 1, 1, 128) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ flatten (Flatten) │ (None, 128) │ 0 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ dense (QDense) │ (None, 10) │ 1,290 │ ├─────────────────────────────────┼────────────────────────┼───────────────┤ │ softmax (Activation) │ (None, 10) │ 0 │ └─────────────────────────────────┴────────────────────────┴───────────────┘
Total params: 594,314 (2.27 MB)
Trainable params: 593,290 (2.26 MB)
Non-trainable params: 1,024 (4.00 KB)
5. Compile and Execute Training
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,
callbacks=[], # Add [lra] here to activate adaptive learning rates
)
Epoch 2/2
844/844 ━━━━━━━━━━━━━━━━━━━━ 66s 76ms/step - accuracy: 0.8841 - loss: 0.3780 - val_accuracy: 0.9643 - val_loss: 0.1246
Der Kernel ist beim Ausführen von Code in der aktuellen Zelle oder einer vorherigen Zelle abgestürzt.
Bitte überprüfen Sie den Code in der/den Zelle(n), um eine mögliche Fehlerursache zu identifizieren.
Klicken Sie <a href='https://aka.ms/vscodeJupyterKernelCrash'>hier</a>, um weitere Informationen zu erhalten.
Weitere Informationen finden Sie unter Jupyter <a href='command:jupyter.viewOutput'>Protokoll</a>.
6. Activation Boundary and Weight Bound Diagnostics
outputs = []
output_names = []
for layer in model.layers:
if layer.__class__.__name__ in [
"QActivation",
"QBatchNormalization",
"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]:
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=False)
print("\nTest loss:", score[0])
print("Test accuracy:", score[1])
1654/1875 ━━━━━━━━━━━━━━━━━━━━ 7s 33ms/step
7. Quantitative Profiling Operations Stats
print_qstats(model)