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()