def testXxdOutputToBytes(self):
# 1. SETUP
# Define the initial model
initial_model = test_utils.build_mock_model()
initial_bytes = flatbuffer_utils.convert_object_to_bytearray(
initial_model)
# Define temporary files
tmp_dir = self.get_temp_dir()
model_filename = os.path.join(tmp_dir, 'model.tflite')
# 2. Write model to temporary file (will be used as input for xxd)
flatbuffer_utils.write_model(initial_model, model_filename)
# 3. DUMP WITH xxd
input_cc_file = os.path.join(tmp_dir, 'model.cc')
command = 'xxd -i {} > {}'.format(model_filename, input_cc_file)
subprocess.call(command, shell=True)
# 4. VALIDATE
final_bytes = flatbuffer_utils.xxd_output_to_bytes(input_cc_file)
# Validate that the initial and final bytearray are the same
self.assertEqual(initial_bytes, final_bytes)
def generate_single_layer_model(self, model, subgraph, op, opcode_idx):
generated_model = schema_fb.ModelT()
generated_model.buffers = []
generated_model.version = 3
# Create subgraph.
generated_subgraph = schema_fb.SubGraphT()
generated_subgraph.inputs = self.inputs
generated_subgraph.outputs = [len(op.inputs)]
generated_subgraph.tensors = []
for input_idx, tensor_idx in enumerate(op.inputs):
tensor = copy.deepcopy(subgraph.tensors[tensor_idx])
tensor.buffer = len(generated_model.buffers)
buffer = copy.deepcopy(
model.buffers[subgraph.tensors[tensor_idx].buffer])
if input_idx in self.inputs:
buffer.data = None
bytes_per_element = BytesFromFlatbufferType(tensor.type)
if buffer.data is not None and len(tensor.shape) > 1:
for i in range(len(buffer.data)):
buffer.data[i] = buffer.data[i] * np.random.uniform(
low=0.5, high=1.0, size=1)
all_equal = True
for i, elem in enumerate(buffer.data):
all_equal = all_equal and elem == model.buffers[
subgraph.tensors[tensor_idx].buffer].data[i]
assert not all_equal
generated_model.buffers.append(buffer)
generated_subgraph.tensors.append(tensor)
for tensor_idx in op.outputs:
tensor = copy.deepcopy(subgraph.tensors[tensor_idx])
tensor.buffer = len(generated_model.buffers)
buffer = copy.deepcopy(
model.buffers[subgraph.tensors[tensor_idx].buffer])
generated_model.buffers.append(buffer)
generated_subgraph.tensors.append(tensor)
# Create op.
generated_op = copy.deepcopy(op)
generated_op.inputs = [i for i in range(len(op.inputs))]
generated_op.outputs = [len(op.inputs)]
generated_op.opcodeIndex = 0
generated_subgraph.operators = [generated_op]
generated_model.subgraphs = [generated_subgraph]
generated_model.operatorCodes = [model.operatorCodes[opcode_idx]]
model_name = self.output_dir + '/' + self.output_dir.split(
'/')[-1] + str(self.op_idx) + '.tflite'
self.op_idx += 1
flatbuffer_utils.write_model(generated_model, model_name)
return model_name
def testWriteReadModel(self):
# 1. SETUP
# Define the initial model
initial_model = test_utils.build_mock_model()
# Define temporary files
tmp_dir = self.get_temp_dir()
model_filename = os.path.join(tmp_dir, 'model.tflite')
# 2. INVOKE
# Invoke the write_model and read_model functions
flatbuffer_utils.write_model(initial_model, model_filename)
final_model = flatbuffer_utils.read_model(model_filename)
# 3. VALIDATE
# Validate that the initial and final models are the same
# Validate the description
self.assertEqual(initial_model.description, final_model.description)
# Validate the main subgraph's name, inputs, outputs, operators and tensors
initial_subgraph = initial_model.subgraphs[0]
final_subgraph = final_model.subgraphs[0]
self.assertEqual(initial_subgraph.name, final_subgraph.name)
for i in range(len(initial_subgraph.inputs)):
self.assertEqual(initial_subgraph.inputs[i],
final_subgraph.inputs[i])
for i in range(len(initial_subgraph.outputs)):
self.assertEqual(initial_subgraph.outputs[i],
final_subgraph.outputs[i])
for i in range(len(initial_subgraph.operators)):
self.assertEqual(initial_subgraph.operators[i].opcodeIndex,
final_subgraph.operators[i].opcodeIndex)
initial_tensors = initial_subgraph.tensors
final_tensors = final_subgraph.tensors
for i in range(len(initial_tensors)):
self.assertEqual(initial_tensors[i].name, final_tensors[i].name)
self.assertEqual(initial_tensors[i].type, final_tensors[i].type)
self.assertEqual(initial_tensors[i].buffer,
final_tensors[i].buffer)
for j in range(len(initial_tensors[i].shape)):
self.assertEqual(initial_tensors[i].shape[j],
final_tensors[i].shape[j])
# Validate the first valid buffer (index 0 is always None)
initial_buffer = initial_model.buffers[1].data
final_buffer = final_model.buffers[1].data
for i in range(initial_buffer.size):
self.assertEqual(initial_buffer.data[i], final_buffer.data[i])
def main(_):
model = flatbuffer_utils.read_model(FLAGS.input_tflite_file)
flatbuffer_utils.strip_strings(model)
flatbuffer_utils.write_model(model, FLAGS.output_tflite_file)
def main(_):
model = flatbuffer_utils.read_model(FLAGS.input_tflite_file)
flatbuffer_utils.randomize_weights(model, FLAGS.random_seed,
FLAGS.buffers_to_skip)
flatbuffer_utils.write_model(model, FLAGS.output_tflite_file)