Image Autoencoder Training Tutorial

This notebook guides you through setting up and training a Quantized Convolutional Autoencoder using QKeras. Autoencoders learn compressed spatial representations (embeddings) of input data. By using hardware-optimized operations like QConv2D, QConv2DTranspose, and quantized_relu, this network targets ultra-low precision footprints (e.g., 4-bit configurations) suitable for efficient edge deployment.

Prerequisites

Make sure you have qkeras and tensorflow installed. If running inside Google Colab, uncomment and run the block below:

# !pip install qkeras-v3 tensorflow numpy
import numpy as np
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 QConv2D, QConv2DTranspose, QActivation, quantized_bits, print_qstats

1. Dataset Preparation and Preprocessing

We download the classic MNIST handwritten digits dataset. To prepare it for our autoencoder architecture, we normalize the raw pixel values to a $[0, 1]$ floating-point range and expand the tensor dimensions to specify a single channel ($1$), satisfying the standard (Batch, Height, Width, Channels) shape layout.

# Configuration Hyperparameters
NB_EPOCH = 10
BATCH_SIZE = 64
VERBOSE = 1
NB_CLASSES = 10
OPTIMIZER = Adam(learning_rate=0.0001, decay=0.000025)
VALIDATION_SPLIT = 0.1

# Load dataset
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# Cast types and expand to image channel dimension (NHWC layout)
x_train = x_train.astype(float)[..., np.newaxis] / 256.0
x_test = x_test.astype(float)[..., np.newaxis] / 256.0

y_train = to_categorical(y_train, NB_CLASSES)
y_test = to_categorical(y_test, NB_CLASSES)

print(f"Training configurations parsed: {x_train.shape[0]} samples ready for training.")
print(f"Testing configurations parsed: {x_test.shape[0]} samples ready for validation.")
/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(
Training configurations parsed: 60000 samples ready for training.
Testing configurations parsed: 10000 samples ready for validation.

2. Compiling the Quantized Autoencoder Network Model

The network structure defines an symmetrical bottleneck architecture:

  1. Encoder Path: Progressively scales down image features via QConv2D layers down to a highly constrained latent layer bottleneck.

  2. Decoder Path: Mirror-reconstructs structural outputs back to original pixel dimension layouts using Transposed Quantized Convolutions (QConv2DTranspose).

Every layer uses a restrictive 4-bit quantizer rule: quantized_bits(4, 0, 1) for structural tensors, and quantized_relu(4, 0) for intermediate non-linear states.

# Input Layer
x_in = Input(shape=(28, 28, 1), name="autoencoder_input")

# --- ENCODER ---
x = QConv2D(32, kernel_size=(3, 3), kernel_quantizer=quantized_bits(4, 0, 1), bias_quantizer=quantized_bits(4, 0, 1), name="qconv_enc_1")(x_in)
x = QActivation("quantized_relu(4,0)", name="qact_enc_1")(x)

x = QConv2D(16, kernel_size=(3, 3), kernel_quantizer=quantized_bits(4, 0, 1), bias_quantizer=quantized_bits(4, 0, 1), name="qconv_enc_2")(x)
x = QActivation("quantized_relu(4,0)", name="qact_enc_2")(x)

x = QConv2D(8, kernel_size=(3, 3), kernel_quantizer=quantized_bits(4, 0, 1), bias_quantizer=quantized_bits(4, 0, 1), name="qconv_enc_3")(x)
x = QActivation("quantized_relu(4,0)", name="qact_enc_3")(x)

# --- DECODER ---
x = QConv2DTranspose(8, kernel_size=(3, 3), kernel_quantizer=quantized_bits(4, 0, 1), bias_quantizer=quantized_bits(4, 0, 1), name="qconv_dec_1")(x)
x = QActivation("quantized_relu(4,0)", name="qact_dec_1")(x)

x = QConv2DTranspose(16, kernel_size=(3, 3), kernel_quantizer=quantized_bits(4, 0, 1), bias_quantizer=quantized_bits(4, 0, 1), name="qconv_dec_2")(x)
x = QActivation("quantized_relu(4,0)", name="qact_dec_2")(x)

x = QConv2DTranspose(32, kernel_size=(3, 3), kernel_quantizer=quantized_bits(4, 0, 1), bias_quantizer=quantized_bits(4, 0, 1), name="qconv_dec_3")(x)
x = QActivation("quantized_relu(4,0)", name="qact_dec_3")(x)

# Final projection block mapping back to 1 image channel
x_out = QConv2D(1, kernel_size=(3, 3), padding="same", kernel_quantizer=quantized_bits(4, 0, 1), bias_quantizer=quantized_bits(4, 0, 1), name="qconv_final")(x)
x_predictions = Activation("sigmoid", name="sigmoid_reconstruction")(x_out)

# Instantiate functional models
model = Model(inputs=[x_in], outputs=[x_predictions])
model.compile(loss="binary_crossentropy", optimizer=OPTIMIZER, metrics=["accuracy"])
model.summary()
Model: "functional"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                     Output Shape                  Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ autoencoder_input (InputLayer)  │ (None, 28, 28, 1)      │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qconv_enc_1 (QConv2D)           │ (None, 26, 26, 32)     │           320 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qact_enc_1 (QActivation)        │ (None, 26, 26, 32)     │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qconv_enc_2 (QConv2D)           │ (None, 24, 24, 16)     │         4,624 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qact_enc_2 (QActivation)        │ (None, 24, 24, 16)     │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qconv_enc_3 (QConv2D)           │ (None, 22, 22, 8)      │         1,160 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qact_enc_3 (QActivation)        │ (None, 22, 22, 8)      │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qconv_dec_1 (QConv2DTranspose)  │ (None, 24, 24, 8)      │           584 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qact_dec_1 (QActivation)        │ (None, 24, 24, 8)      │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qconv_dec_2 (QConv2DTranspose)  │ (None, 26, 26, 16)     │         1,168 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qact_dec_2 (QActivation)        │ (None, 26, 26, 16)     │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qconv_dec_3 (QConv2DTranspose)  │ (None, 28, 28, 32)     │         4,640 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qact_dec_3 (QActivation)        │ (None, 28, 28, 32)     │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qconv_final (QConv2D)           │ (None, 28, 28, 1)      │           289 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ sigmoid_reconstruction          │ (None, 28, 28, 1)      │             0 │
│ (Activation)                    │                        │               │
└─────────────────────────────────┴────────────────────────┴───────────────┘
 Total params: 12,785 (49.94 KB)
 Trainable params: 12,785 (49.94 KB)
 Non-trainable params: 0 (0.00 B)

3. Network Training Loop Execution

We train the model by passing x_train as both the input and the target goal. This setup forces the model to minimize error configurations when reconstructing its original pixel profiles.

history = model.fit(
    x_train,
    x_train,
    batch_size=BATCH_SIZE,
    epochs=NB_EPOCH,
    initial_epoch=1,
    verbose=VERBOSE,
    validation_split=VALIDATION_SPLIT
)
Epoch 2/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 63s 71ms/step - accuracy: 0.8057 - loss: 0.2527 - val_accuracy: 0.8085 - val_loss: 0.0801
Epoch 3/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 76s 90ms/step - accuracy: 0.8087 - loss: 0.0742 - val_accuracy: 0.8088 - val_loss: 0.0718
Epoch 4/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 67s 79ms/step - accuracy: 0.8088 - loss: 0.0695 - val_accuracy: 0.8088 - val_loss: 0.0688
Epoch 5/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 68s 81ms/step - accuracy: 0.8088 - loss: 0.0676 - val_accuracy: 0.8088 - val_loss: 0.0674
Epoch 6/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 66s 78ms/step - accuracy: 0.8088 - loss: 0.0663 - val_accuracy: 0.8088 - val_loss: 0.0664
Epoch 7/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 69s 81ms/step - accuracy: 0.8088 - loss: 0.0658 - val_accuracy: 0.8088 - val_loss: 0.0662
Epoch 8/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 70s 83ms/step - accuracy: 0.8088 - loss: 0.0653 - val_accuracy: 0.8088 - val_loss: 0.0653
Epoch 9/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 67s 79ms/step - accuracy: 0.8088 - loss: 0.0648 - val_accuracy: 0.8088 - val_loss: 0.0652
Epoch 10/10
844/844 ━━━━━━━━━━━━━━━━━━━━ 66s 79ms/step - accuracy: 0.8088 - loss: 0.0646 - val_accuracy: 0.8088 - val_loss: 0.0651
Training optimization process completed.

4. Inference and Parameter Profile Analysis

With training complete, we can generate compressed predictions from our hidden nodes and query weight tensor metadata blocks. Finally, we execute print_qstats to review the memory savings achieved across our custom low-precision constraints.

# Execute inference reconstructions on sample segments
num_reco = 8
samples = x_test[:num_reco]
reconstructions = model.predict(samples)

print("\n--- Layer Weight Metric Outlines ---")
for layer in model.layers:
    for w, weight in enumerate(layer.get_weights()):
        print(f"Layer Target: {layer.name:<25} | Weight Component Index: {w} | Layout Dimensions: {weight.shape}")

print("\n--- QKeras Quantized Bit Allocation Summary Profiles ---")
print_qstats(model)
1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 946ms/step

--- Layer Weight Metric Outlines ---
Layer Target: qconv_enc_1               | Weight Component Index: 0 | Layout Dimensions: (3, 3, 1, 32)
Layer Target: qconv_enc_1               | Weight Component Index: 1 | Layout Dimensions: (32,)
Layer Target: qconv_enc_2               | Weight Component Index: 0 | Layout Dimensions: (3, 3, 32, 16)
Layer Target: qconv_enc_2               | Weight Component Index: 1 | Layout Dimensions: (16,)
Layer Target: qconv_enc_3               | Weight Component Index: 0 | Layout Dimensions: (3, 3, 16, 8)
Layer Target: qconv_enc_3               | Weight Component Index: 1 | Layout Dimensions: (8,)
Layer Target: qconv_dec_1               | Weight Component Index: 0 | Layout Dimensions: (3, 3, 8, 8)
Layer Target: qconv_dec_1               | Weight Component Index: 1 | Layout Dimensions: (8,)
Layer Target: qconv_dec_2               | Weight Component Index: 0 | Layout Dimensions: (3, 3, 16, 8)
Layer Target: qconv_dec_2               | Weight Component Index: 1 | Layout Dimensions: (16,)
Layer Target: qconv_dec_3               | Weight Component Index: 0 | Layout Dimensions: (3, 3, 32, 16)
Layer Target: qconv_dec_3               | Weight Component Index: 1 | Layout Dimensions: (32,)
Layer Target: qconv_final               | Weight Component Index: 0 | Layout Dimensions: (3, 3, 32, 1)
Layer Target: qconv_final               | Weight Component Index: 1 | Layout Dimensions: (1,)

--- QKeras Quantized Bit Allocation Summary Profiles ---

Number of operations in model:
    qconv_enc_1                   : 194688 (smult_4_8)
    qconv_enc_2                   : 2654208 (smult_4_4)
    qconv_enc_3                   : 557568 (smult_4_4)
    qconv_final                   : 225792 (smult_4_4)

Number of operation types in model:
    smult_4_4                     : 3437568
    smult_4_8                     : 194688

Weight profiling:
    qconv_enc_1_weights            : 288   (4-bit unit)
    qconv_enc_1_bias               : 32    (4-bit unit)
    qconv_enc_2_weights            : 4608  (4-bit unit)
    qconv_enc_2_bias               : 16    (4-bit unit)
    qconv_enc_3_weights            : 1152  (4-bit unit)
    qconv_enc_3_bias               : 8     (4-bit unit)
    qconv_final_weights            : 288   (4-bit unit)
    qconv_final_bias               : 1     (4-bit unit)
    ----------------------------------------
    Total Bits                     : 25572

Weight sparsity:
... quantizing model
    qconv_enc_1                    : 0.1250
    qconv_enc_2                    : 0.1347
    qconv_enc_3                    : 0.1129
    qconv_dec_1                    : 0.0908
    qconv_dec_2                    : 0.1438
    qconv_dec_3                    : 0.1386
    qconv_final                    : 0.1384
    ----------------------------------------
    Total Sparsity                 : 0.1328