How to quantize a Keras model to int4

INT4 quantization reduces the precision of trained Keras model weights for inference-time storage and memory bandwidth. It is useful after training when a smaller model artifact matters more than continued training in full precision.

Keras applies post-training quantization after the model has built its weights. In the default INT4 path, supported layers store weights in 4-bit values while activations are dynamically quantized for inference, so the model should be checked with the same input shape and task metric used by the serving path.

A small JAX-backed Dense classifier keeps the local check focused on the quantization state, saved .keras artifact, and reloaded prediction shape. Quantize only after training is complete, then compare the quantized model against a representative validation set before shipping it.

1. Install Keras and the JAX backend in the active Python environment.

   $ python -m pip install --upgrade keras jax

   Use the backend already chosen for the project when it supports the model and quantization path. JAX keeps this local CPU check small.
   Related: How to install Keras with pip

2. Select the JAX backend before importing Keras.

   $ export KERAS_BACKEND=jax

   Keras reads KERAS_BACKEND during import. Set it in the shell, notebook kernel, or process environment before any import keras statement.
   Related: How to set the Keras backend

3. Create quantize_int4.py with backend selection, imports, a sample batch, and a built model.

   quantize_int4.py

   import os
   from pathlib import Path
    
   os.environ.setdefault("KERAS_BACKEND", "jax")
    
   import keras
   import numpy as np
    
    
   keras.utils.set_random_seed(17)
    
   model = keras.Sequential(
       [
           keras.Input(shape=(6,), name="features"),
           keras.layers.Dense(12, activation="relu", name="feature_projection"),
           keras.layers.Dense(3, activation="softmax", name="class_score"),
       ]
   )
    
   sample = np.array(
       [
           [0.2, 0.4, 0.1, 0.9, 0.3, 0.7],
           [0.8, 0.1, 0.6, 0.2, 0.5, 0.4],
       ],
       dtype="float32",
   )

   Replace the temporary classifier and sample tensor with the trained model and validation input shape used by the target inference path.

4. Add the baseline prediction and INT4 quantization call.

   baseline = model.predict(sample, verbose=0)
   model.quantize(
       "int4",
       filters=lambda layer: isinstance(layer, keras.layers.Dense),
   )
   quantized = model.predict(sample, verbose=0)

   The filter targets the supported Dense compute layers in the classifier and avoids quantizing the input placeholder. Expand or remove the filter only after checking support and accuracy for the layer types in the trained model.

5. Save, reload, and print the quantization evidence.

   output_path = Path("int4-classifier.keras")
   if output_path.exists():
       output_path.unlink()
    
   model.save(output_path)
   reloaded = keras.models.load_model(output_path)
   reloaded_output = reloaded.predict(sample, verbose=0)
    
   print(f"backend: {keras.backend.backend()}")
   print(f"quantized file: {output_path}")
   print(f"file size: {output_path.stat().st_size} bytes")
   print(f"output shape: {quantized.shape}")
   print(f"reloaded output shape: {reloaded_output.shape}")
   print(f"max prediction delta: {np.max(np.abs(baseline - quantized)):.6f}")
   print(f"reloaded matches quantized: {np.allclose(quantized, reloaded_output, atol=1e-5)}")
    
   for layer in model.layers:
       if hasattr(layer, "quantization_mode") and layer.quantization_mode:
           print(f"{layer.name} quantization: {layer.quantization_mode}")

   max prediction delta is a quick smoke check, not a replacement for task metrics. Use accuracy, loss, perplexity, or another deployment metric on representative data before accepting an INT4 model.

6. Run the script and confirm that the quantized model reloads with the same output shape.

   $ python quantize_int4.py
   backend: jax
   quantized file: int4-classifier.keras
   file size: 17358 bytes
   output shape: (2, 3)
   reloaded output shape: (2, 3)
   max prediction delta: 0.008058
   reloaded matches quantized: True
   feature_projection quantization: int4
   class_score quantization: int4

   The True reload check proves that the saved .keras artifact can be loaded and used for inference. The int4 layer lines confirm that the supported compute layers were quantized.