tf.quantization.experimental.quantize_saved_model

tf.quantization.experimental.quantize_saved_model(
    saved_model_path: str,
    output_directory: Optional[str] = None,
    quantization_options: Optional[tf.quantization.experimental.QuantizationOptions] = None,
    representative_dataset: Optional[repr_dataset.RepresentativeDatasetOrMapping] = None,
    *,
    overwrite_output_directory: bool = False
) -> autotrackable.AutoTrackable

# Quantizing a model trained with QAT.
quantization_options = tf.quantization.experimental.QuantizationOptions(
    signature_keys=['your_signature_key'],
)
tf.quantization.experimental.quantize_saved_model(
    '/tmp/input_model',
    '/tmp/output_model',
    quantization_options=quantization_options,
)

# When quantizing a model trained without QAT (Post-Training Quantization),
# a representative dataset is required.
representative_dataset = [{"input": tf.random.uniform(shape=(3, 3))}
                      for _ in range(256)]
tf.quantization.experimental.quantize_saved_model(
    '/tmp/input_model',
    '/tmp/output_model',
    quantization_options=quantization_options,
    representative_dataset={'your_signature_key': representative_dataset},
  )

# In addition to preset quantization methods, fine-grained control of
# quantization for each component is also supported.
_QuantizationComponentSpec = (
    tf.quantization.experimental.QuantizationComponentSpec
)
quantization_options = tf.quantization.experimental.QuantizationOptions(
    signature_keys=['your_signature_key'],
    quantization_method=tf.quantization.experimental.QuantizationMethod(
        quantization_component_specs=[
            _QuantizationComponentSpec(
                quantization_component=(
                    _QuantizationComponentSpec.COMPONENT_ACTIVATION
                ),
                tensor_type=_QuantizationComponentSpec.TENSORTYPE_INT_8,
            )
        ]
    )
)
tf.quantization.experimental.quantize_saved_model(
    '/tmp/input_model',
    '/tmp/output_model',
    quantization_options=quantization_options,
)