tf.quantization.experimental.TfRecordRepresentativeDatasetSaver

Representative dataset saver in TFRecord format.

tf.quantization.experimental.TfRecordRepresentativeDatasetSaver(
    path_map: Mapping[str, os.PathLike[str]],
    expected_input_key_map: Optional[Mapping[str, Collection[str]]] = None
)

Saves representative datasets for quantization calibration in TFRecord format. The samples are serialized as RepresentativeDataSample.

The save method return a signature key to RepresentativeDatasetFile map, which can be used for QuantizationOptions.

Example usage:

# Creating the representative dataset.
representative_dataset = [{"input": tf.random.uniform(shape=(3, 3))}
                      for _ in range(256)]

# Saving to a TFRecord file.
dataset_file_map = (
  tf.quantization.experimental.TfRecordRepresentativeDatasetSaver(
        path_map={'serving_default': '/tmp/representative_dataset_path'}
    ).save({'serving_default': representative_dataset})
)

# Using in QuantizationOptions.
quantization_options = tf.quantization.experimental.QuantizationOptions(
    signature_keys=['serving_default'],
    representative_datasets=dataset_file_map,
)
tf.quantization.experimental.quantize_saved_model(
    '/tmp/input_model',
    '/tmp/output_model',
    quantization_options=quantization_options,
)

path_map Signature def key -> path mapping. Each path is a TFRecord file to which a RepresentativeDataset is saved. The signature def keys should be a subset of the SignatureDef keys of the representative_dataset argument of the save() call.
expected_input_key_map Signature def key -> expected input keys. If set, validate that the sample has same set of input keys before saving.

KeyError If path_map and expected_input_key_map have different keys.

Methods

save

View source

save(
    representative_dataset: RepresentativeDatasetMapping
) -> Mapping[str, _RepresentativeDatasetFile]

Saves the representative dataset.

Args
representative_dataset Signature def key -> representative dataset mapping. Each dataset is saved in a separate TFRecord file whose path matches the signature def key of path_map.

Raises
ValueError When the signature def key in representative_dataset is not present in the path_map.

Returns
A map from signature key to the RepresentativeDatasetFile instance contains the path to the saved file.