Skip to content

Instantly share code, notes, and snippets.

@petewarden
Created January 29, 2020 20:35
Show Gist options
  • Save petewarden/180ea281bbe09a6a7cbc8b2558097fc7 to your computer and use it in GitHub Desktop.
Save petewarden/180ea281bbe09a6a7cbc8b2558097fc7 to your computer and use it in GitHub Desktop.
Example of using the new Flatbuffers Python API to modify a TensorFlow Lite file
def StripTfliteFile(tflite_input, tflite_output):
"""Given 'tflite_input`, write stripped version, 'tflite_output'."""
if not os.path.exists(tflite_input):
raise RuntimeError("Invalid filename %r" % tflite_input)
with open(tflite_input, "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)
model_offset = model.Pack(builder)
builder.Finish(model_offset)
model_data = builder.Output()
with open(tflite_output, "wb") as out_file:
out_file.write(model_data)
@bojeckkk
Copy link

Hmmm, Netron says "RangeError: Offset is outside the bounds of the DataView ". Stripped model doesn't load in TFLite aswell: "ERROR: The model is not a valid Flatbuffer buffer"

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment