View source on GitHub |
Representative dataset saver in TFRecord format.
tf.quantization.experimental.TfRecordRepresentativeDatasetSaver(
path_map: Mapping[str, os.PathLike[str]]
)
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,
)
Methods
save
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. |