Quantized RNN

This notebook shows how to build and quantize recurrent models with QKerasV3 using the standalone keras package. It updates the original QKeras RNN tutorial for the Keras 3 API and avoids older tf.keras/tf_keras idioms.

The examples use the IMDB sentiment dataset and compare:

  1. a floating-point Bidirectional(LSTM) baseline,

  2. a hand-written quantized recurrent model,

  3. automatic quantization with model_quantize, and

  4. an optional AutoQKeras search section.

QKerasV3 is installed from the package qkeras-v3, but imported as qkeras.

# In a fresh environment, uncomment one of these:
# %pip install -q "qkeras-v3" "tensorflow>=2.16" "keras>=3"
# %pip install -q git+https://github.com/fastmachinelearning/qkerasV3.git
import os
import warnings

# Pick the TensorFlow backend before importing keras when running outside this notebook.
os.environ.setdefault("KERAS_BACKEND", "tensorflow")
warnings.filterwarnings("ignore")

import keras
import numpy as np
import tensorflow as tf
from keras import layers
from keras.datasets import imdb
from keras.optimizers import Adam
from keras.utils import pad_sequences

from qkeras import QLSTM, QActivation, QBidirectional, QDense
from qkeras.utils import model_quantize

print("keras:", keras.__version__)
print("tensorflow:", tf.__version__)
print("KERAS_BACKEND:", os.environ.get("KERAS_BACKEND"))
keras: 3.13.0
tensorflow: 2.20.0
KERAS_BACKEND: tensorflow

Reproducibility and device setup

The tutorial runs on CPU, GPU, or TPU. The dataset is intentionally subsetted so that the notebook remains quick in CI and Colab. Increase TRAIN_SAMPLES, TEST_SAMPLES, and EPOCHS for a more meaningful accuracy comparison.

keras.utils.set_random_seed(812)

physical_devices = tf.config.list_physical_devices()
print("Physical devices:")
for device in physical_devices:
    print(" -", device)

try:
    resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
    tf.config.experimental_connect_to_cluster(resolver)
    tf.tpu.experimental.initialize_tpu_system(resolver)
    strategy = tf.distribute.TPUStrategy(resolver)
    print("Running with TPU strategy")
except Exception:
    strategy = tf.distribute.get_strategy()
    print("Running with default strategy")

print("Replicas:", strategy.num_replicas_in_sync)
Physical devices:
 - PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')
Running with default strategy
Replicas: 1

Load and prepare IMDB

Keras 3 exposes sequence padding as keras.utils.pad_sequences. Avoid importing keras.preprocessing.sequence, which is a legacy path in modern Keras code.

MAX_FEATURES = 10_000
MAXLEN = 100
BATCH_SIZE = 128
TRAIN_SAMPLES = 4_000
TEST_SAMPLES = 1_000
EPOCHS = 10

print("Loading IMDB data...")
(x_train, y_train), (x_test, y_test) = imdb.load_data(num_words=MAX_FEATURES)

x_train = pad_sequences(x_train[:TRAIN_SAMPLES], maxlen=MAXLEN)
y_train = np.asarray(y_train[:TRAIN_SAMPLES]).astype("float32")
x_test = pad_sequences(x_test[:TEST_SAMPLES], maxlen=MAXLEN)
y_test = np.asarray(y_test[:TEST_SAMPLES]).astype("float32")

train_dataset = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .shuffle(min(TRAIN_SAMPLES, 10_000), seed=812)
    .batch(BATCH_SIZE, drop_remainder=False)
    .prefetch(tf.data.AUTOTUNE)
)

test_dataset = (
    tf.data.Dataset.from_tensor_slices((x_test, y_test))
    .batch(BATCH_SIZE, drop_remainder=False)
    .prefetch(tf.data.AUTOTUNE)
)

print("x_train:", x_train.shape, x_train.dtype)
print("x_test: ", x_test.shape, x_test.dtype)
Loading IMDB data...
x_train: (4000, 100) int32
x_test:  (1000, 100) int32

Floating-point baseline

The important Keras 3 change is that activations are layers when you need naming or graph placement. Use layers.Activation("linear", name=...), not keras.activations(...).

UNITS = 64
EMBEDDING_DIM = 64
LOSS = "binary_crossentropy"


def create_model(
    maxlen=MAXLEN,
    max_features=MAX_FEATURES,
    embedding_dim=EMBEDDING_DIM,
    units=UNITS,
):
    inputs = keras.Input(shape=(maxlen,), dtype="int32", name="tokens")
    x = layers.Embedding(
        input_dim=max_features,
        output_dim=embedding_dim,
        name="embedding",
    )(inputs)
    x = layers.Activation("linear", name="embedding_act")(x)
    x = layers.Bidirectional(layers.LSTM(units), name="bidirectional")(x)
    outputs = layers.Dense(1, activation="sigmoid", name="dense")(x)
    return keras.Model(inputs=inputs, outputs=outputs, name="float_lstm")


keras.backend.clear_session()
with strategy.scope():
    model = create_model()
    model.compile(optimizer=Adam(learning_rate=1e-3), loss=LOSS, metrics=["accuracy"])

model.summary()
Model: "float_lstm"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                     Output Shape                  Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ tokens (InputLayer)             │ (None, 100)            │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ embedding (Embedding)           │ (None, 100, 64)        │       640,000 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ embedding_act (Activation)      │ (None, 100, 64)        │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ bidirectional (Bidirectional)   │ (None, 128)            │        66,048 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense (Dense)                   │ (None, 1)              │           129 │
└─────────────────────────────────┴────────────────────────┴───────────────┘
 Total params: 706,177 (2.69 MB)
 Trainable params: 706,177 (2.69 MB)
 Non-trainable params: 0 (0.00 B)
history = model.fit(
    train_dataset,
    epochs=EPOCHS,
    validation_data=test_dataset,
    verbose=2,
)
Epoch 1/10
32/32 - 4s - 125ms/step - accuracy: 0.5475 - loss: 0.6858 - val_accuracy: 0.7050 - val_loss: 0.6029
Epoch 2/10
32/32 - 4s - 117ms/step - accuracy: 0.7765 - loss: 0.4888 - val_accuracy: 0.7960 - val_loss: 0.4502
Epoch 3/10
32/32 - 4s - 111ms/step - accuracy: 0.8990 - loss: 0.2505 - val_accuracy: 0.8130 - val_loss: 0.4490
Epoch 4/10
32/32 - 5s - 155ms/step - accuracy: 0.9503 - loss: 0.1411 - val_accuracy: 0.7810 - val_loss: 0.5751
Epoch 5/10
32/32 - 5s - 153ms/step - accuracy: 0.9783 - loss: 0.0691 - val_accuracy: 0.7950 - val_loss: 0.6073
Epoch 6/10
32/32 - 4s - 135ms/step - accuracy: 0.9883 - loss: 0.0353 - val_accuracy: 0.7870 - val_loss: 0.7224
Epoch 7/10
32/32 - 4s - 136ms/step - accuracy: 0.9967 - loss: 0.0160 - val_accuracy: 0.7970 - val_loss: 0.7506
Epoch 8/10
32/32 - 4s - 127ms/step - accuracy: 0.9970 - loss: 0.0108 - val_accuracy: 0.7920 - val_loss: 0.9091
Epoch 9/10
32/32 - 4s - 124ms/step - accuracy: 0.9970 - loss: 0.0100 - val_accuracy: 0.7850 - val_loss: 0.8242
Epoch 10/10
32/32 - 4s - 127ms/step - accuracy: 0.9962 - loss: 0.0121 - val_accuracy: 0.7860 - val_loss: 0.9753

Hand-written quantized recurrent model

For recurrent layers, QKerasV3 keeps the QKeras style: use QLSTM for a quantized recurrent cell and QBidirectional for a quantized wrapper. The final dense layer is replaced by QDense, followed by a QActivation for the output sigmoid.

def create_qmodel(
    maxlen=MAXLEN,
    max_features=MAX_FEATURES,
    embedding_dim=EMBEDDING_DIM,
    units=UNITS,
):
    inputs = keras.Input(shape=(maxlen,), dtype="int32", name="tokens")
    x = layers.Embedding(
        input_dim=max_features,
        output_dim=embedding_dim,
        name="embedding",
    )(inputs)

    # Quantize the embedding output before the recurrent block.
    x = QActivation("quantized_bits(4,0,1)", name="embedding_act")(x)

    lstm = QLSTM(
        units,
        activation="quantized_tanh(4)",
        recurrent_activation="quantized_sigmoid(4)",
        kernel_quantizer="quantized_bits(4,0,1,alpha='auto')",
        recurrent_quantizer="quantized_bits(4,0,1,alpha='auto')",
        bias_quantizer="quantized_bits(4,0,1,alpha='auto')",
        name="qlstm",
    )
    x = QBidirectional(lstm, name="qbidirectional")(x)

    x = QDense(
        1,
        kernel_quantizer="quantized_bits(4,0,1,alpha='auto')",
        bias_quantizer="quantized_bits(4,0,1,alpha='auto')",
        name="qdense",
    )(x)
    outputs = QActivation("sigmoid", name="sigmoid")(x)
    return keras.Model(inputs=inputs, outputs=outputs, name="manual_quant_lstm")


keras.backend.clear_session()
with strategy.scope():
    qmodel = create_qmodel()
    qmodel.compile(optimizer=Adam(learning_rate=1e-3), loss=LOSS, metrics=["accuracy"])

qmodel.summary()
Model: "manual_quant_lstm"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                     Output Shape                  Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ tokens (InputLayer)             │ (None, 100)            │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ embedding (Embedding)           │ (None, 100, 64)        │       640,000 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ embedding_act (QActivation)     │ (None, 100, 64)        │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qbidirectional (QBidirectional) │ (None, 128)            │        66,048 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ qdense (QDense)                 │ (None, 1)              │           129 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ sigmoid (QActivation)           │ (None, 1)              │             0 │
└─────────────────────────────────┴────────────────────────┴───────────────┘
 Total params: 706,177 (2.69 MB)
 Trainable params: 706,177 (2.69 MB)
 Non-trainable params: 0 (0.00 B)
qhistory = qmodel.fit(
    train_dataset,
    epochs=EPOCHS,
    validation_data=test_dataset,
    verbose=2,
)
Epoch 1/10
32/32 - 14s - 435ms/step - accuracy: 0.4942 - loss: 0.6939 - val_accuracy: 0.4910 - val_loss: 0.6938
Epoch 2/10
32/32 - 8s - 262ms/step - accuracy: 0.6690 - loss: 0.5974 - val_accuracy: 0.7280 - val_loss: 0.5630
Epoch 3/10
32/32 - 11s - 352ms/step - accuracy: 0.7598 - loss: 0.4981 - val_accuracy: 0.7320 - val_loss: 0.5471
Epoch 4/10
32/32 - 11s - 341ms/step - accuracy: 0.7725 - loss: 0.4805 - val_accuracy: 0.7110 - val_loss: 0.5682
Epoch 5/10
32/32 - 9s - 287ms/step - accuracy: 0.7605 - loss: 0.4904 - val_accuracy: 0.7300 - val_loss: 0.5592
Epoch 6/10
32/32 - 10s - 324ms/step - accuracy: 0.7778 - loss: 0.4749 - val_accuracy: 0.7250 - val_loss: 0.5525
Epoch 7/10
32/32 - 9s - 276ms/step - accuracy: 0.7760 - loss: 0.4669 - val_accuracy: 0.7120 - val_loss: 0.5780
Epoch 8/10
32/32 - 9s - 282ms/step - accuracy: 0.7828 - loss: 0.4678 - val_accuracy: 0.7210 - val_loss: 0.5635
Epoch 9/10
32/32 - 9s - 276ms/step - accuracy: 0.8077 - loss: 0.4377 - val_accuracy: 0.7170 - val_loss: 0.5641
Epoch 10/10
32/32 - 10s - 320ms/step - accuracy: 0.8075 - loss: 0.4267 - val_accuracy: 0.7070 - val_loss: 0.5850

Automatic quantization with model_quantize

model_quantize can clone a floating-point model and replace supported layers using a quantizer configuration. For wrappers such as Bidirectional, quantization support may depend on the exact QKerasV3 version. The cell below is therefore written defensively and prints the error if the installed version does not support a given conversion path yet.

BITS = 4
quantizer_config = {
    "embedding_act": "quantized_bits(4,0,1)",
    "bidirectional": {
        "activation": f"quantized_tanh({BITS})",
        "recurrent_activation": f"quantized_sigmoid({BITS})",
        "kernel_quantizer": f"quantized_bits({BITS},0,1,alpha='auto')",
        "recurrent_quantizer": f"quantized_bits({BITS},0,1,alpha='auto')",
        "bias_quantizer": f"quantized_bits({BITS},0,1,alpha='auto')",
    },
    "dense": {
        "kernel_quantizer": f"quantized_bits({BITS},0,1,alpha='auto')",
        "bias_quantizer": f"quantized_bits({BITS},0,1,alpha='auto')",
    },
}

keras.backend.clear_session()
with strategy.scope():
    base_model = create_model()
    auto_qmodel = model_quantize(
        base_model,
        quantizer_config,
        activation_bits=BITS,
        custom_objects={},
        transfer_weights=True,
    )
    auto_qmodel.compile(
        optimizer=Adam(learning_rate=1e-3),
        loss=LOSS,
        metrics=["accuracy"],
    )

auto_qmodel.summary()
Model: "float_lstm"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                     Output Shape                  Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ tokens (InputLayer)             │ (None, 100)            │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ embedding (Embedding)           │ (None, 100, 64)        │       640,000 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ embedding_act (QActivation)     │ (None, 100, 64)        │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ bidirectional (QBidirectional)  │ (None, 128)            │        66,048 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense (QDense)                  │ (None, 1)              │           129 │
└─────────────────────────────────┴────────────────────────┴───────────────┘
 Total params: 706,177 (2.69 MB)
 Trainable params: 706,177 (2.69 MB)
 Non-trainable params: 0 (0.00 B)
auto_qhistory = auto_qmodel.fit(
    train_dataset,
    epochs=EPOCHS,
    validation_data=test_dataset,
    verbose=2,
)
Epoch 1/10
32/32 - 10s - 311ms/step - accuracy: 0.4955 - loss: 0.6926 - val_accuracy: 0.5410 - val_loss: 0.6990
Epoch 2/10
32/32 - 7s - 228ms/step - accuracy: 0.5807 - loss: 1.4403 - val_accuracy: 0.6470 - val_loss: 2.6811
Epoch 3/10
32/32 - 8s - 241ms/step - accuracy: 0.6933 - loss: 1.6885 - val_accuracy: 0.6890 - val_loss: 0.7481
Epoch 4/10
32/32 - 10s - 324ms/step - accuracy: 0.7685 - loss: 0.5999 - val_accuracy: 0.7190 - val_loss: 0.6991
Epoch 5/10
32/32 - 9s - 277ms/step - accuracy: 0.7903 - loss: 0.5725 - val_accuracy: 0.7070 - val_loss: 0.7455
Epoch 6/10
32/32 - 7s - 216ms/step - accuracy: 0.8073 - loss: 0.5155 - val_accuracy: 0.7150 - val_loss: 0.8558
Epoch 7/10
32/32 - 6s - 183ms/step - accuracy: 0.8332 - loss: 0.5208 - val_accuracy: 0.7260 - val_loss: 0.8274
Epoch 8/10
32/32 - 6s - 182ms/step - accuracy: 0.8490 - loss: 0.4677 - val_accuracy: 0.7320 - val_loss: 0.9162
Epoch 9/10
32/32 - 6s - 202ms/step - accuracy: 0.8585 - loss: 0.4391 - val_accuracy: 0.7410 - val_loss: 0.9431
Epoch 10/10
32/32 - 6s - 192ms/step - accuracy: 0.9057 - loss: 0.3311 - val_accuracy: 0.7590 - val_loss: 0.8895

Save and reload with the Keras 3 .keras format

Use the native .keras format for Keras 3 models. Avoid legacy .h5 unless you specifically need HDF5 compatibility.

model_path = "manual_quant_lstm.keras"
qmodel.save(model_path)

reloaded = keras.models.load_model(model_path)
print("Reloaded model:", reloaded.name)
print("Reloaded layers:", [layer.__class__.__name__ for layer in reloaded.layers])
Reloaded model: manual_quant_lstm
Reloaded layers: ['InputLayer', 'Embedding', 'QActivation', 'QBidirectional', 'QDense', 'QActivation']

Notes for QKerasV3 migration

  • Install QKerasV3 as qkeras-v3, import it as qkeras.

  • Prefer standalone keras imports over tensorflow.keras imports.

  • Set KERAS_BACKEND before importing keras when running scripts.

  • Use layers.Activation(...) or layer-level activation=...; do not call keras.activations(...) as if it were a layer.

  • Use keras.utils.pad_sequences instead of legacy preprocessing imports.

  • Save full models with .keras for Keras 3 compatibility.

  • Keep quantizer strings simple and syntactically valid, e.g. quantized_bits(4,0,1,alpha='auto').