def _create_metadata_file(self):
associated_file1 = _metadata_fb.AssociatedFileT()
associated_file1.name = b"file1"
associated_file2 = _metadata_fb.AssociatedFileT()
associated_file2.name = b"file2"
self.expected_recorded_files = [
six.ensure_str(associated_file1.name),
six.ensure_str(associated_file2.name)
]
input_meta = _metadata_fb.TensorMetadataT()
output_meta = _metadata_fb.TensorMetadataT()
output_meta.associatedFiles = [associated_file2]
subgraph = _metadata_fb.SubGraphMetadataT()
# Create a model with two inputs and one output.
subgraph.inputTensorMetadata = [input_meta, input_meta]
subgraph.outputTensorMetadata = [output_meta]
model_meta = _metadata_fb.ModelMetadataT()
model_meta.name = "Mobilenet_quantized"
model_meta.associatedFiles = [associated_file1]
model_meta.subgraphMetadata = [subgraph]
b = flatbuffers.Builder(0)
b.Finish(
model_meta.Pack(b),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
metadata_file = self.create_tempfile().full_path
with open(metadata_file, "wb") as f:
f.write(b.Output())
return metadata_file
def testPopulateMetadataFileToModelWithMetadataAndAssociatedFiles(self):
# First, creates a dummy metadata different from self._metadata_file. It
# needs to have the same input/output tensor numbers as self._model_file.
# Populates it and the associated files into the model.
input_meta = _metadata_fb.TensorMetadataT()
output_meta = _metadata_fb.TensorMetadataT()
subgraph = _metadata_fb.SubGraphMetadataT()
# Create a model with two inputs and one output.
subgraph.inputTensorMetadata = [input_meta, input_meta]
subgraph.outputTensorMetadata = [output_meta]
model_meta = _metadata_fb.ModelMetadataT()
model_meta.subgraphMetadata = [subgraph]
b = flatbuffers.Builder(0)
b.Finish(
model_meta.Pack(b),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
metadata_buf = b.Output()
# Populate the metadata.
populator1 = _metadata.MetadataPopulator.with_model_file(self._model_file)
populator1.load_metadata_buffer(metadata_buf)
populator1.load_associated_files([self._file1, self._file2])
populator1.populate()
# Then, populate the metadata again.
populator2 = _metadata.MetadataPopulator.with_model_file(self._model_file)
populator2.load_metadata_file(self._metadata_file)
populator2.populate()
# Test if the metadata is populated correctly.
self._assert_golden_metadata(self._model_file)
def _create_metadata_buffer_with_wrong_identifier(self): # Creates a metadata with wrong identifier wrong_identifier = b"widn" metadata = _metadata_fb.ModelMetadataT() metadata_builder = flatbuffers.Builder(0) metadata_builder.Finish(metadata.Pack(metadata_builder), wrong_identifier) return metadata_builder.Output()
def _create_model_meta_with_subgraph_meta(self, subgraph_meta):
model_meta = _metadata_fb.ModelMetadataT()
model_meta.subgraphMetadata = [subgraph_meta]
b = flatbuffers.Builder(0)
b.Finish(model_meta.Pack(b),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
return b.Output()
def _create_dummy_metadata(self):
# Create dummy input metadata
input_metadata = _metadata_fb.TensorMetadataT()
input_metadata.name = _INPUT_NAME
# Create dummy output metadata
output_metadata = _metadata_fb.TensorMetadataT()
output_metadata.name = _OUTPUT_NAME
# Create dummy model_metadata
model_metadata = _metadata_fb.ModelMetadataT()
model_metadata.name = _MODEL_NAME
return model_metadata, input_metadata, output_metadata
def create_from_metadata(
cls,
model_buffer: bytearray,
model_metadata: Optional[_metadata_fb.ModelMetadataT] = None,
input_metadata: Optional[List[
_metadata_fb.TensorMetadataT]] = None,
output_metadata: Optional[List[
_metadata_fb.TensorMetadataT]] = None,
associated_files: Optional[List[str]] = None):
"""Creates MetadataWriter based on the metadata Flatbuffers Python Objects.
Args:
model_buffer: valid buffer of the model file.
model_metadata: general model metadata [1]. The subgraph_metadata will be
refreshed with input_metadata and output_metadata.
input_metadata: a list of metadata of the input tensors [2].
output_metadata: a list of metadata of the output tensors [3].
associated_files: path to the associated files to be populated.
[1]:
https://github.com/tensorflow/tflite-support/blob/b80289c4cd1224d0e1836c7654e82f070f9eefaa/tensorflow_lite_support/metadata/metadata_schema.fbs#L640-L681
[2]:
https://github.com/tensorflow/tflite-support/blob/b80289c4cd1224d0e1836c7654e82f070f9eefaa/tensorflow_lite_support/metadata/metadata_schema.fbs#L590
[3]:
https://github.com/tensorflow/tflite-support/blob/b80289c4cd1224d0e1836c7654e82f070f9eefaa/tensorflow_lite_support/metadata/metadata_schema.fbs#L599
Returns:
A MetadataWriter Object.
"""
# Create empty tensor metadata when input_metadata/output_metadata are None
# to bypass MetadataPopulator verification.
if not input_metadata:
model = _schema_fb.Model.GetRootAsModel(model_buffer, 0)
num_input_tensors = model.Subgraphs(0).InputsLength()
input_metadata = [_metadata_fb.TensorMetadataT()
] * num_input_tensors
if not output_metadata:
model = _schema_fb.Model.GetRootAsModel(model_buffer, 0)
num_output_tensors = model.Subgraphs(0).OutputsLength()
output_metadata = [_metadata_fb.TensorMetadataT()
] * num_output_tensors
subgraph_metadata = _metadata_fb.SubGraphMetadataT()
subgraph_metadata.inputTensorMetadata = input_metadata
subgraph_metadata.outputTensorMetadata = output_metadata
if model_metadata is None:
model_metadata = _metadata_fb.ModelMetadataT()
model_metadata.subgraphMetadata = [subgraph_metadata]
b = flatbuffers.Builder(0)
b.Finish(model_metadata.Pack(b),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
return cls(model_buffer, b.Output(), associated_files)
def _create_dummy_model_metadata(
tensor_metadata: _metadata_fb.TensorMetadataT) -> bytes:
# Create a dummy model using the tensor metadata.
subgraph_metadata = _metadata_fb.SubGraphMetadataT()
subgraph_metadata.inputTensorMetadata = [tensor_metadata]
model_metadata = _metadata_fb.ModelMetadataT()
model_metadata.subgraphMetadata = [subgraph_metadata]
# Create the Flatbuffers object and convert it to the json format.
builder = flatbuffers.Builder(0)
builder.Finish(model_metadata.Pack(builder),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
return bytes(builder.Output())
def create_metadata(self) -> _metadata_fb.ModelMetadataT:
"""Creates the model metadata based on the general model information.
Returns:
A Flatbuffers Python object of the model metadata.
"""
model_metadata = _metadata_fb.ModelMetadataT()
model_metadata.name = self.name
model_metadata.version = self.version
model_metadata.description = self.description
model_metadata.author = self.author
model_metadata.license = self.licenses
return model_metadata
def testLoadMetadataBufferWithNoSubgraphMetadataThrowsException(self):
# Create a dummy metadata without Subgraph.
model_meta = _metadata_fb.ModelMetadataT()
builder = flatbuffers.Builder(0)
builder.Finish(
model_meta.Pack(builder),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
meta_buf = builder.Output()
populator = _metadata.MetadataPopulator.with_model_buffer(self._model_buf)
with self.assertRaises(ValueError) as error:
populator.load_metadata_buffer(meta_buf)
self.assertEqual(
"The number of SubgraphMetadata should be exactly one, but got 0.",
str(error.exception))
def test_create_metadata_should_succeed(self):
file_md = metadata_info.AssociatedFileMd(
file_path="label.txt",
description="The label file.",
file_type=_metadata_fb.AssociatedFileType.TENSOR_AXIS_LABELS,
locale="en")
file_metadata = file_md.create_metadata()
# Create the Flatbuffers object and convert it to the json format.
model_metadata = _metadata_fb.ModelMetadataT()
model_metadata.associatedFiles = [file_metadata]
builder = flatbuffers.Builder(0)
builder.Finish(model_metadata.Pack(builder),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
metadata_json = _metadata.convert_to_json(bytes(builder.Output()))
expected_json = test_utils.load_file(self._EXPECTED_META_JSON, "r")
self.assertEqual(metadata_json, expected_json)
def test_create_score_calibration_file_md_should_succeed(self):
score_calibration_md = metadata_info.ScoreCalibrationMd(
_metadata_fb.ScoreTransformationType.LOG, self._DEFAULT_VALUE,
self._SCORE_CALIBRATION_FILE)
score_calibration_file_md = (
score_calibration_md.create_score_calibration_file_md())
file_metadata = score_calibration_file_md.create_metadata()
# Create the Flatbuffers object and convert it to the json format.
model_metadata = _metadata_fb.ModelMetadataT()
model_metadata.associatedFiles = [file_metadata]
builder = flatbuffers.Builder(0)
builder.Finish(model_metadata.Pack(builder),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
metadata_json = _metadata.convert_to_json(bytes(builder.Output()))
expected_json = test_utils.load_file(self._EXPECTED_MODEL_META_JSON,
"r")
self.assertEqual(metadata_json, expected_json)
def testLoadMetadataBufferWithWrongOutputMetaNumberThrowsException(self):
# Create a dummy metadata with no output tensor metadata, while the expected
# number is 1.
input_meta = _metadata_fb.TensorMetadataT()
subgprah_meta = _metadata_fb.SubGraphMetadataT()
subgprah_meta.inputTensorMetadata = [input_meta, input_meta]
model_meta = _metadata_fb.ModelMetadataT()
model_meta.subgraphMetadata = [subgprah_meta]
builder = flatbuffers.Builder(0)
builder.Finish(
model_meta.Pack(builder),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
meta_buf = builder.Output()
populator = _metadata.MetadataPopulator.with_model_buffer(self._model_buf)
with self.assertRaises(ValueError) as error:
populator.load_metadata_buffer(meta_buf)
self.assertEqual(
("The number of output tensors (1) should match the number of "
"output tensor metadata (0)"), str(error.exception))
