feat: add tool for visualizing compressed (and uncompressed) models #3099
Conversation
Add a development tool that prints compressed and uncompressed .tflite models to stdout in a human-readable, searchable, structured, text format. Helpful annotations (indexes of lists, names of operators, etc.) derived from the model are added as virtual fields with names beginning with an _underscore. Add a unit test which simply ensures the viewer does not crash when run on several models found in the source tree. BUG=see description
| from tensorflow.lite.python import schema_py_generated as tflite_schema | ||
|
|
||
| USAGE = textwrap.dedent(f"""\ | ||
| Usage: {os.path.basename(sys.argv[0])} $(realpath <MODEL>) |
There was a problem hiding this comment.
Is that $(realpath <MODEL>) somehow expanded? I'm thinking the output should just be view <model path>. On the other hand, if tensorflow.lite.micro.compression isn't in the PyPi, then this script can only be run using Bazel, in which case the person running the Bazel command would probably know what $(realpath) is?
There was a problem hiding this comment.
Fixed. See 19f7f21.
| # Convert the model into a Python dictionary, expressing the hierarchial nature | ||
| # of the model, and pretty print the dictionary. Please extend as needed for | ||
| # your use case. | ||
|
|
There was a problem hiding this comment.
I would add a bazel run command line example here.
There was a problem hiding this comment.
Fixed. See 757e95e.
| } for t in sorted(lut_tensors, key=lambda x: x.tensor)] | ||
|
|
||
|
|
||
| def is_compressed_buffer(buffer_index, unpacked_metadata): |
There was a problem hiding this comment.
Method not used. I believe this was part of the code I added to show compressed bitstrings as lists of indices.
There was a problem hiding this comment.
Fixed. See 2d56eef.
| return False, None, None | ||
|
|
||
|
|
||
| def unpack_indices(buffer, lut_data): |
There was a problem hiding this comment.
Method not used. I believe this was part of the code I added to show compressed bitstrings as lists of indices.
There was a problem hiding this comment.
Fixed. See 2d56eef.
| return indices | ||
|
|
||
|
|
||
| def unpack_compression_metadata(buffer): |
There was a problem hiding this comment.
Fixed. See 2d56eef.
| return buffers | ||
|
|
||
|
|
||
| def get_compression_metadata_buffer(model): |
There was a problem hiding this comment.
Fixed. See 2d56eef.
Add functionality to decode and display LUT indices for compressed tensor buffers in the model viewer. When viewing compressed models, buffers containing compressed data now show a "_lut_indices" field with the decoded indices displayed in a readable multi-line format. Also improve the numpy array pretty printer to display all values without ellipsis truncation, making large arrays more readable.
Make the usage message context-aware by detecting if the script is running under Bazel. When run through Bazel, show the full bazel command and include an example using $(realpath ...) to convert relative paths to absolute paths, which is required due to Bazel's sandboxing. When run standalone, show a simpler usage message.
Remove several functions which were part of an earlier implementation.
Good idea! Implemented in 4919c04, without a flag. Nothing has a flag yet. It's expected you'll pipe this output a pager or write it to a file and search it. |
Add a development tool that prints compressed and uncompressed
.tflite models to stdout in a human-readable, searchable,
structured, text format. Helpful annotations (indexes of lists,
names of operators, etc.) derived from the model are added as
virtual fields with names beginning with an _underscore.
Add a unit test which simply ensures the viewer does not crash
when run on several models found in the source tree.
BUG=see description