How to quantize a Keras model to INT8

Keras INT8 post-training quantization stores supported model weights in lower precision while keeping an inference path for normal float input tensors. It belongs after training, when a model needs a smaller .keras artifact or a lower-memory serving path without changing the model architecture.

The model.quantize("int8") API rewrites a built model in place. Current Keras quantization support covers Dense, EinsumDense, Embedding, ReversibleEmbedding, and composite layers built from those layers; this example uses Dense layers so the quantized layer names appear directly in the output.

Quantization changes the inference artifact, not the training workflow. Keep the full-precision model for retraining, measure INT8 output against validation data, and save the quantized model only after the size and accuracy trade-off is acceptable for the target backend and device.

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

   $ python -m pip install --upgrade keras jax

   Use the backend already selected for the project when the model depends on TensorFlow or PyTorch-specific layers.
   Related: How to install Keras with pip

2. Create quantize_int8.py with backend selection and sample tensors.

   quantize_int8.py

   import os
   from pathlib import Path
    
   os.environ.setdefault("KERAS_BACKEND", "jax")
    
   import keras
   import numpy as np
    
    
   keras.utils.set_random_seed(31)
   rng = np.random.default_rng(31)
    
   x_train = rng.normal(size=(64, 256)).astype("float32")
   y_train = rng.normal(size=(64, 16)).astype("float32")
   x_eval = rng.normal(size=(8, 256)).astype("float32")

   KERAS_BACKEND must be set before import keras. The script-level setting keeps the example self-contained, while a shell export or notebook kernel setting works for project code.
   Related: How to set the Keras backend

3. Add the Dense model and one training epoch.

   model = keras.Sequential(
       [
           keras.Input(shape=(256,), name="features"),
           keras.layers.Dense(512, activation="relu", name="hidden_large"),
           keras.layers.Dense(128, activation="relu", name="hidden_small"),
           keras.layers.Dense(16, name="score"),
       ]
   )
   model.compile(optimizer="adam", loss="mse")
   model.fit(x_train, y_train, epochs=1, batch_size=16, verbose=0)

   Replace the training block with the project model load call when the trained model already exists. Call build() or run one forward pass before quantization if the model was created without materialized weights.

4. Save the full-precision output and artifact before quantization.

   fp32_output = model(x_eval, training=False)
   fp32_path = Path("credit-risk-fp32.keras")
   int8_path = Path("credit-risk-int8.keras")
   for path in (fp32_path, int8_path):
       if path.exists():
           path.unlink()
    
   model.save(fp32_path)
   fp32_size = fp32_path.stat().st_size

   The full-precision artifact remains the retraining and rollback source. INT8 quantization is a post-training inference step, so do not use the quantized copy as the only saved model.

5. Quantize the built model and save the INT8 artifact.

   model.quantize("int8")
   int8_output = model(x_eval, training=False)
   model.save(int8_path)

   Pass a filters list or callable to model.quantize() when sensitive layers such as logits or small residual paths should remain in full precision.

6. Add reload and comparison checks to the script.

   reloaded = keras.saving.load_model(int8_path)
   reloaded_output = reloaded(x_eval, training=False)
    
   mse_int8 = keras.ops.mean(keras.ops.square(fp32_output - int8_output))
   mse_roundtrip = keras.ops.mean(keras.ops.square(int8_output - reloaded_output))
    
   quantized_layers = [
       layer.name
       for layer in model.layers
       if getattr(getattr(layer, "dtype_policy", None), "quantization_mode", None) == "int8"
   ]
    
   int8_size = int8_path.stat().st_size
   size_reduction = 100.0 * (1.0 - (int8_size / fp32_size))
    
   print(f"Backend: {keras.backend.backend()}")
   print("Quantized layers:", ", ".join(quantized_layers))
   print(f"FP32 file size: {fp32_size} bytes")
   print(f"INT8 file size: {int8_size} bytes")
   print(f"Size reduction: {size_reduction:.1f}%")
   print(f"FP32 vs INT8 MSE: {float(mse_int8):.8f}")
   print(f"Reloaded INT8 MSE: {float(mse_roundtrip):.8f}")
   print(f"Output shape: {tuple(np.asarray(int8_output).shape)}")

7. Run the script and confirm the quantized artifact reloads.

   $ python quantize_int8.py
   Backend: jax
   Quantized layers: hidden_large, hidden_small, score
   FP32 file size: 2417383 bytes
   INT8 file size: 1826008 bytes
   Size reduction: 24.5%
   FP32 vs INT8 MSE: 0.00007353
   Reloaded INT8 MSE: 0.00000000
   Output shape: (8, 16)

   Quantized layers confirms that every Dense layer in the sample switched to the INT8 policy. Reloaded INT8 MSE confirms that the saved .keras artifact reloads with the same quantized output on the sample batch.