tfmot.quantization.keras.quantize_annotate_model

Annotate a tf.keras model to be quantized.

tfmot.quantization.keras.quantize_annotate_model(
    to_annotate
)

Used in the notebooks

Used in the guide
Cluster preserving quantization aware training (CQAT) Keras example Sparsity and cluster preserving quantization aware training (PCQAT) Keras example Pruning preserving quantization aware training (PQAT) Keras example

This function does not actually quantize the model. It merely specifies that the model needs to be quantized. quantize_apply can then be used to quantize the model.

This function is intended to be used in conjunction with the quantize_annotate_layer API. Otherwise, it is simpler to use quantize_model.

Annotate a model while overriding the default behavior for a layer:

quantize_config = MyDenseQuantizeConfig()

model = quantize_annotate_model(
  keras.Sequential([
    layers.Dense(10, activation='relu', input_shape=(100,)),
    quantize_annotate_layer(
        layers.Dense(2, activation='sigmoid'),
        quantize_config=quantize_config)
  ]))

# The first Dense layer gets quantized with the default behavior,
# but the second layer uses `MyDenseQuantizeConfig` for quantization.
quantized_model = quantize_apply(model)

Note that this function removes the optimizer from the original model.

Args
to_annotate	tf.keras model which needs to be quantized.

Returns
New tf.keras model with each layer in the model wrapped with QuantizeAnnotate. The new model preserves weights from the original model.

Raises
ValueError	if the model cannot be annotated.