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 testFlatbufferToDict(self): model_data = test_utils.build_mock_model() model_dict = visualize.CreateDictFromFlatbuffer(model_data) self.assertEqual(0, model_dict['version']) self.assertEqual(1, len(model_dict['subgraphs'])) self.assertEqual(1, len(model_dict['operator_codes'])) self.assertEqual(3, len(model_dict['buffers'])) self.assertEqual(3, len(model_dict['subgraphs'][0]['tensors'])) self.assertEqual(0, model_dict['subgraphs'][0]['tensors'][0]['buffer'])
def testStripStrings(self):
# 1. SETUP
# Define the initial model
initial_model = test_utils.build_mock_model()
final_model = copy.deepcopy(initial_model)
# 2. INVOKE
# Invoke the strip_strings function
flatbuffer_utils.strip_strings(final_model)
# 3. VALIDATE
# Validate that the initial and final models are the same except strings
# Validate the description
self.assertIsNotNone(initial_model.description)
self.assertIsNone(final_model.description)
self.assertIsNotNone(initial_model.signatureDefs)
self.assertIsNone(final_model.signatureDefs)
# Validate the main subgraph's name, inputs, outputs, operators and tensors
initial_subgraph = initial_model.subgraphs[0]
final_subgraph = final_model.subgraphs[0]
self.assertIsNotNone(initial_model.subgraphs[0].name)
self.assertIsNone(final_model.subgraphs[0].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.assertIsNotNone(initial_tensors[i].name)
self.assertIsNone(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 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 testRandomizeSomeWeights(self):
# 1. SETUP
# Define the initial model
initial_model = test_utils.build_mock_model()
final_model = copy.deepcopy(initial_model)
# 2. INVOKE
# Invoke the randomize_weights function, but skip the first buffer
flatbuffer_utils.randomize_weights(
final_model, buffers_to_skip=[_SKIPPED_BUFFER_INDEX])
# 3. VALIDATE
# Validate that the initial and final models are the same, except that
# the weights in the model buffer have been modified (i.e, randomized)
# 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 enumerate(initial_subgraph.inputs):
self.assertEqual(initial_subgraph.inputs[i],
final_subgraph.inputs[i])
for i, _ in enumerate(initial_subgraph.outputs):
self.assertEqual(initial_subgraph.outputs[i],
final_subgraph.outputs[i])
for i, _ in enumerate(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 enumerate(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 that the skipped buffer is unchanged.
initial_buffer = initial_model.buffers[_SKIPPED_BUFFER_INDEX].data
final_buffer = final_model.buffers[_SKIPPED_BUFFER_INDEX].data
for j in range(initial_buffer.size):
self.assertEqual(initial_buffer.data[j], final_buffer.data[j])
def testVisualize(self):
model_data = test_utils.build_mock_model()
tmp_dir = self.get_temp_dir()
model_filename = os.path.join(tmp_dir, 'model.tflite')
with open(model_filename, 'wb') as model_file:
model_file.write(model_data)
html_filename = os.path.join(tmp_dir, 'visualization.html')
visualize.CreateHtmlFile(model_filename, html_filename)
with open(html_filename, 'r') as html_file:
html_text = html_file.read()
# It's hard to test debug output without doing a full HTML parse,
# but at least sanity check that expected identifiers are present.
self.assertRegex(
html_text, re.compile(r'%s' % model_filename, re.MULTILINE | re.DOTALL))
self.assertRegex(html_text,
re.compile(r'input_tensor', re.MULTILINE | re.DOTALL))
self.assertRegex(html_text,
re.compile(r'constant_tensor', re.MULTILINE | re.DOTALL))
self.assertRegex(html_text, re.compile(r'ADD', re.MULTILINE | re.DOTALL))
