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)
]
output_meta = _metadata_fb.TensorMetadataT()
output_meta.associatedFiles = [associated_file2]
subgraph = _metadata_fb.SubGraphMetadataT()
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 StripTfliteFile(input_tflite_file, output_tflite_file):
"""Strips all nonessential strings from the model to reduce model size.
Args:
input_tflite_file: Full path name to the input tflite file
output_tflite_file: Full path name to the stripped output tflite file.
Raises:
RuntimeError: If input_tflite_file is not found.
IOError: If input_tflite_file or output_tflite_file cannot be opened.
"""
if not os.path.exists(input_tflite_file):
raise RuntimeError('Input file not found at %r\n' % input_tflite_file)
with open(input_tflite_file, 'rb') as file_handle:
file_data = bytearray(file_handle.read())
model_obj = schema_fb.Model.GetRootAsModel(file_data, 0)
model = schema_fb.ModelT.InitFromObj(model_obj)
model.description = ''
for subgraph in model.subgraphs:
subgraph.name = ''
for tensor in subgraph.tensors:
tensor.name = ''
builder = flatbuffers.Builder(1024) # Initial size of the buffer, which
# will grow automatically if needed
model_offset = model.Pack(builder)
builder.Finish(model_offset)
model_data = builder.Output()
with open(output_tflite_file, 'wb') as out_file:
out_file.write(model_data)
def _create_empty_model_buf(self):
model = _schema_fb.ModelT()
model_builder = flatbuffers.Builder(0)
model_builder.Finish(
model.Pack(model_builder),
_metadata.MetadataPopulator.TFLITE_FILE_IDENTIFIER)
return model_builder.Output()
def _populate_metadata_buffer(self):
"""Populates the metadata buffer (in bytearray) into the model file.
Inserts metadata_buf into the metadata field of schema.Model. If the
MetadataPopulator object is created using the method,
with_model_file(model_file), the model file will be updated.
Existing metadata buffer (if applied) will be overridden by the new metadata
buffer.
"""
with open(self._model_file, "rb") as f:
model_buf = f.read()
model = _schema_fb.ModelT.InitFromObj(
_schema_fb.Model.GetRootAsModel(model_buf, 0))
buffer_field = _schema_fb.BufferT()
buffer_field.data = self._metadata_buf
is_populated = False
if not model.metadata:
model.metadata = []
else:
# Check if metadata has already been populated.
for meta in model.metadata:
if meta.name.decode("utf-8") == self.METADATA_FIELD_NAME:
is_populated = True
model.buffers[meta.buffer] = buffer_field
if not is_populated:
if not model.buffers:
model.buffers = []
model.buffers.append(buffer_field)
# Creates a new metadata field.
metadata_field = _schema_fb.MetadataT()
metadata_field.name = self.METADATA_FIELD_NAME
metadata_field.buffer = len(model.buffers) - 1
model.metadata.append(metadata_field)
# Packs model back to a flatbuffer binaray file.
b = flatbuffers.Builder(0)
b.Finish(model.Pack(b), self.TFLITE_FILE_IDENTIFIER)
model_buf = b.Output()
# Saves the updated model buffer to model file.
# Gets files that have been packed to self._model_file.
packed_files = self.get_packed_associated_file_list()
if packed_files:
# Writes the updated model buffer and associated files into a new model
# file. Then overwrites the original model file.
with tempfile.NamedTemporaryFile() as temp:
new_file = temp.name
with open(new_file, "wb") as f:
f.write(model_buf)
self._copy_archived_files(self._model_file, new_file, packed_files)
shutil.copy(new_file, self._model_file)
os.remove(new_file)
else:
with open(self._model_file, "wb") as f:
f.write(model_buf)
def load_metadata_buffer(self, metadata_buf):
"""Loads the metadata buffer (in bytearray) to be populated.
Args:
metadata_buf: metadata buffer (in bytearray) to be populated.
Raises:
ValueError: The metadata to be populated is empty.
ValueError: The metadata does not have the expected flatbuffer identifer.
ValueError: Error occurs when getting the minimum metadata parser version.
"""
if not metadata_buf:
raise ValueError("The metadata to be populated is empty.")
_assert_metadata_buffer_identifier(metadata_buf)
# Gets the minimum metadata parser version of the metadata_buf.
min_version = _pywrap_metadata_version.GetMinimumMetadataParserVersion(
bytes(metadata_buf))
# Inserts in the minimum metadata parser version into the metadata_buf.
metadata = _metadata_fb.ModelMetadataT.InitFromObj(
_metadata_fb.ModelMetadata.GetRootAsModelMetadata(metadata_buf, 0))
metadata.minParserVersion = min_version
b = flatbuffers.Builder(0)
b.Finish(metadata.Pack(b), self.METADATA_FILE_IDENTIFIER)
metadata_buf_with_version = b.Output()
self._metadata_buf = metadata_buf_with_version
def _convert_model_from_object_to_bytearray(model_object): """Converts a tflite model from a parsable object into a bytearray.""" # Initial size of the buffer, which will grow automatically if needed builder = flatbuffers.Builder(1024) model_offset = model_object.Pack(builder) builder.Finish(model_offset, file_identifier=_TFLITE_FILE_IDENTIFIER) return bytes(builder.Output())
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 write_model(model, output_tflite_file):
"""Writes the model, a python flatbuffer object, into the output tflite file.
Args:
model: tflite model
output_tflite_file: Full path name to the output tflite file.
Raises:
IOError: If output_tflite_file cannot be opened.
"""
# Initial size of the buffer, which will grow automatically if needed
builder = flatbuffers.Builder(1024)
model_offset = model.Pack(builder)
builder.Finish(model_offset, file_identifier=TFLITE_FILE_IDENTIFIER)
model_data = builder.Output()
with open(output_tflite_file, 'wb') as out_file:
out_file.write(model_data)
def _create_model_file_with_metadata_and_buf_fields(self):
metadata_field = _schema_fb.MetadataT()
metadata_field.name = "meta"
buffer_field = _schema_fb.BufferT()
model = _schema_fb.ModelT()
model.metadata = [metadata_field, metadata_field]
model.buffers = [buffer_field, buffer_field, buffer_field]
model_builder = flatbuffers.Builder(0)
model_builder.Finish(
model.Pack(model_builder),
_metadata.MetadataPopulator.TFLITE_FILE_IDENTIFIER)
mnodel_file = self.create_tempfile().full_path
with open(mnodel_file, "wb") as f:
f.write(model_builder.Output())
return mnodel_file
def _populate_metadata_with_identifier(self, model_buf, metadata_buf,
identifier):
# For testing purposes only. MetadataPopulator cannot populate metadata with
# wrong identifiers.
model = _schema_fb.ModelT.InitFromObj(
_schema_fb.Model.GetRootAsModel(model_buf, 0))
buffer_field = _schema_fb.BufferT()
buffer_field.data = metadata_buf
model.buffers = [buffer_field]
# Creates a new metadata field.
metadata_field = _schema_fb.MetadataT()
metadata_field.name = _metadata.MetadataPopulator.METADATA_FIELD_NAME
metadata_field.buffer = len(model.buffers) - 1
model.metadata = [metadata_field]
b = flatbuffers.Builder(0)
b.Finish(model.Pack(b), identifier)
return b.Output()
def testPopulateMetadataFileToModelWithMetadataAndAssociatedFiles(self):
# First, creates a dummy metadata. Populates it and the associated files
# into the model.
model_meta = _metadata_fb.ModelMetadataT()
model_meta.name = "Mobilenet_quantized"
b = flatbuffers.Builder(0)
b.Finish(
model_meta.Pack(b),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
metadata_buf = b.Output()
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, populates the metadata again.
populator2 = _metadata.MetadataPopulator.with_model_file(self._model_file)
populator2.load_metadata_file(self._metadata_file)
populator2.populate()
# Tests if the metadata is populated correctly.
self._assert_golden_metadata(self._model_file)
def build_mock_model():
"""Creates a flatbuffer containing an example model."""
builder = flatbuffers.Builder(1024)
schema_fb.BufferStart(builder)
buffer0_offset = schema_fb.BufferEnd(builder)
schema_fb.BufferStartDataVector(builder, 10)
builder.PrependUint8(9)
builder.PrependUint8(8)
builder.PrependUint8(7)
builder.PrependUint8(6)
builder.PrependUint8(5)
builder.PrependUint8(4)
builder.PrependUint8(3)
builder.PrependUint8(2)
builder.PrependUint8(1)
builder.PrependUint8(0)
buffer1_data_offset = builder.EndVector(10)
schema_fb.BufferStart(builder)
schema_fb.BufferAddData(builder, buffer1_data_offset)
buffer1_offset = schema_fb.BufferEnd(builder)
schema_fb.BufferStart(builder)
buffer2_offset = schema_fb.BufferEnd(builder)
schema_fb.ModelStartBuffersVector(builder, 3)
builder.PrependUOffsetTRelative(buffer2_offset)
builder.PrependUOffsetTRelative(buffer1_offset)
builder.PrependUOffsetTRelative(buffer0_offset)
buffers_offset = builder.EndVector(3)
string0_offset = builder.CreateString('input_tensor')
schema_fb.TensorStartShapeVector(builder, 3)
builder.PrependInt32(1)
builder.PrependInt32(2)
builder.PrependInt32(5)
shape0_offset = builder.EndVector(3)
schema_fb.TensorStart(builder)
schema_fb.TensorAddName(builder, string0_offset)
schema_fb.TensorAddShape(builder, shape0_offset)
schema_fb.TensorAddType(builder, 0)
schema_fb.TensorAddBuffer(builder, 0)
tensor0_offset = schema_fb.TensorEnd(builder)
schema_fb.QuantizationParametersStartMinVector(builder, 5)
builder.PrependFloat32(0.5)
builder.PrependFloat32(2.0)
builder.PrependFloat32(5.0)
builder.PrependFloat32(10.0)
builder.PrependFloat32(20.0)
quant1_min_offset = builder.EndVector(5)
schema_fb.QuantizationParametersStartMaxVector(builder, 5)
builder.PrependFloat32(10.0)
builder.PrependFloat32(20.0)
builder.PrependFloat32(-50.0)
builder.PrependFloat32(1.0)
builder.PrependFloat32(2.0)
quant1_max_offset = builder.EndVector(5)
schema_fb.QuantizationParametersStartScaleVector(builder, 5)
builder.PrependFloat32(3.0)
builder.PrependFloat32(4.0)
builder.PrependFloat32(5.0)
builder.PrependFloat32(6.0)
builder.PrependFloat32(7.0)
quant1_scale_offset = builder.EndVector(5)
schema_fb.QuantizationParametersStartZeroPointVector(builder, 5)
builder.PrependInt64(1)
builder.PrependInt64(2)
builder.PrependInt64(3)
builder.PrependInt64(-1)
builder.PrependInt64(-2)
quant1_zero_point_offset = builder.EndVector(5)
schema_fb.QuantizationParametersStart(builder)
schema_fb.QuantizationParametersAddMin(builder, quant1_min_offset)
schema_fb.QuantizationParametersAddMax(builder, quant1_max_offset)
schema_fb.QuantizationParametersAddScale(builder, quant1_scale_offset)
schema_fb.QuantizationParametersAddZeroPoint(builder,
quant1_zero_point_offset)
quantization1_offset = schema_fb.QuantizationParametersEnd(builder)
string1_offset = builder.CreateString('constant_tensor')
schema_fb.TensorStartShapeVector(builder, 3)
builder.PrependInt32(1)
builder.PrependInt32(2)
builder.PrependInt32(5)
shape1_offset = builder.EndVector(3)
schema_fb.TensorStart(builder)
schema_fb.TensorAddName(builder, string1_offset)
schema_fb.TensorAddShape(builder, shape1_offset)
schema_fb.TensorAddType(builder, 0)
schema_fb.TensorAddBuffer(builder, 1)
schema_fb.TensorAddQuantization(builder, quantization1_offset)
tensor1_offset = schema_fb.TensorEnd(builder)
string2_offset = builder.CreateString('output_tensor')
schema_fb.TensorStartShapeVector(builder, 3)
builder.PrependInt32(1)
builder.PrependInt32(2)
builder.PrependInt32(5)
shape2_offset = builder.EndVector(3)
schema_fb.TensorStart(builder)
schema_fb.TensorAddName(builder, string2_offset)
schema_fb.TensorAddShape(builder, shape2_offset)
schema_fb.TensorAddType(builder, 0)
schema_fb.TensorAddBuffer(builder, 2)
tensor2_offset = schema_fb.TensorEnd(builder)
schema_fb.SubGraphStartTensorsVector(builder, 3)
builder.PrependUOffsetTRelative(tensor2_offset)
builder.PrependUOffsetTRelative(tensor1_offset)
builder.PrependUOffsetTRelative(tensor0_offset)
tensors_offset = builder.EndVector(3)
schema_fb.SubGraphStartInputsVector(builder, 1)
builder.PrependInt32(0)
inputs_offset = builder.EndVector(1)
schema_fb.SubGraphStartOutputsVector(builder, 1)
builder.PrependInt32(2)
outputs_offset = builder.EndVector(1)
schema_fb.OperatorCodeStart(builder)
schema_fb.OperatorCodeAddBuiltinCode(builder,
schema_fb.BuiltinOperator.ADD)
schema_fb.OperatorCodeAddVersion(builder, 1)
code_offset = schema_fb.OperatorCodeEnd(builder)
schema_fb.ModelStartOperatorCodesVector(builder, 1)
builder.PrependUOffsetTRelative(code_offset)
codes_offset = builder.EndVector(1)
schema_fb.OperatorStartInputsVector(builder, 2)
builder.PrependInt32(0)
builder.PrependInt32(1)
op_inputs_offset = builder.EndVector(2)
schema_fb.OperatorStartOutputsVector(builder, 1)
builder.PrependInt32(2)
op_outputs_offset = builder.EndVector(1)
schema_fb.OperatorStart(builder)
schema_fb.OperatorAddOpcodeIndex(builder, 0)
schema_fb.OperatorAddInputs(builder, op_inputs_offset)
schema_fb.OperatorAddOutputs(builder, op_outputs_offset)
op_offset = schema_fb.OperatorEnd(builder)
schema_fb.SubGraphStartOperatorsVector(builder, 1)
builder.PrependUOffsetTRelative(op_offset)
ops_offset = builder.EndVector(1)
string3_offset = builder.CreateString('subgraph_name')
schema_fb.SubGraphStart(builder)
schema_fb.SubGraphAddName(builder, string3_offset)
schema_fb.SubGraphAddTensors(builder, tensors_offset)
schema_fb.SubGraphAddInputs(builder, inputs_offset)
schema_fb.SubGraphAddOutputs(builder, outputs_offset)
schema_fb.SubGraphAddOperators(builder, ops_offset)
subgraph_offset = schema_fb.SubGraphEnd(builder)
schema_fb.ModelStartSubgraphsVector(builder, 1)
builder.PrependUOffsetTRelative(subgraph_offset)
subgraphs_offset = builder.EndVector(1)
string4_offset = builder.CreateString('model_description')
schema_fb.ModelStart(builder)
schema_fb.ModelAddOperatorCodes(builder, codes_offset)
schema_fb.ModelAddSubgraphs(builder, subgraphs_offset)
schema_fb.ModelAddDescription(builder, string4_offset)
schema_fb.ModelAddBuffers(builder, buffers_offset)
model_offset = schema_fb.ModelEnd(builder)
builder.Finish(model_offset)
model = builder.Output()
return model
def _create_model_buffer_with_wrong_identifier(self):
wrong_identifier = b"widn"
model = _schema_fb.ModelT()
model_builder = flatbuffers.Builder(0)
model_builder.Finish(model.Pack(model_builder), wrong_identifier)
return model_builder.Output()
