experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Update Strip strings to also clear signatures in the flatbuffer.

Browse Source 
We clear the whole signature list as removing strings makes the signature useless.

PiperOrigin-RevId: 346594577
Change-Id: I1b4e2b19f69bf260fefc74ff02e92d874c8b8e76
This commit is contained in:
Karim Nosir 2020-12-09 11:19:56 -08:00 committed by TensorFlower Gardener
parent 76a17bad58
commit 698571054b
3 changed files with 40 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/lite/tools/flatbuffer_utils.py
View File

@ -126,6 +126,8 @@ def strip_strings(model):
subgraph.name = None
for tensor in subgraph.tensors:
tensor.name = None
# We clear all signature_def structure, since without names it is useless.
model.signatureDefs = None
def randomize_weights(model, random_seed=0):

3
tensorflow/lite/tools/flatbuffer_utils_test.py
View File

@ -89,6 +89,9 @@ class StripStringsTest(test_util.TensorFlowTestCase):
# 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]

35
tensorflow/lite/tools/test_utils.py
View File

@ -196,6 +196,40 @@ def build_mock_flatbuffer_model():
builder.PrependUOffsetTRelative(subgraph_offset)
subgraphs_offset = builder.EndVector(1)
signature_method = builder.CreateString('my_method')
signature_key = builder.CreateString('my_key')
input_tensor_string = builder.CreateString('input_tensor')
output_tensor_string = builder.CreateString('output_tensor')
# Signature Inputs
schema_fb.TensorMapStart(builder)
schema_fb.TensorMapAddName(builder, input_tensor_string)
schema_fb.TensorMapAddTensorIndex(builder, 1)
input_tensor = schema_fb.TensorMapEnd(builder)
# Signature Outputs
schema_fb.TensorMapStart(builder)
schema_fb.TensorMapAddName(builder, output_tensor_string)
schema_fb.TensorMapAddTensorIndex(builder, 2)
output_tensor = schema_fb.TensorMapEnd(builder)
schema_fb.SignatureDefStartInputsVector(builder, 1)
builder.PrependUOffsetTRelative(input_tensor)
signature_inputs_offset = builder.EndVector(1)
schema_fb.SignatureDefStartOutputsVector(builder, 1)
builder.PrependUOffsetTRelative(output_tensor)
signature_outputs_offset = builder.EndVector(1)
schema_fb.SignatureDefStart(builder)
schema_fb.SignatureDefAddKey(builder, signature_key)
schema_fb.SignatureDefAddMethodName(builder, signature_method)
schema_fb.SignatureDefAddInputs(builder, signature_inputs_offset)
schema_fb.SignatureDefAddOutputs(builder, signature_outputs_offset)
signature_offset = schema_fb.SignatureDefEnd(builder)
schema_fb.ModelStartSignatureDefsVector(builder, 1)
builder.PrependUOffsetTRelative(signature_offset)
signature_defs_offset = builder.EndVector(1)
string4_offset = builder.CreateString('model_description')
schema_fb.ModelStart(builder)
schema_fb.ModelAddVersion(builder, TFLITE_SCHEMA_VERSION)
@ -203,6 +237,7 @@ def build_mock_flatbuffer_model():
schema_fb.ModelAddSubgraphs(builder, subgraphs_offset)
schema_fb.ModelAddDescription(builder, string4_offset)
schema_fb.ModelAddBuffers(builder, buffers_offset)
schema_fb.ModelAddSignatureDefs(builder, signature_defs_offset)
model_offset = schema_fb.ModelEnd(builder)
builder.Finish(model_offset)
model = builder.Output()