106 lines
3.8 KiB
Python
106 lines
3.8 KiB
Python
|
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
|
|||
|
|
#
|
|||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|||
|
|
# you may not use this file except in compliance with the License.
|
|||
|
|
# You may obtain a copy of the License at
|
|||
|
|
#
|
|||
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|||
|
|
#
|
|||
|
|
# Unless required by applicable law or agreed to in writing, software
|
|||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|||
|
|
# See the License for the specific language governing permissions and
|
|||
|
|
# limitations under the License.
|
|||
|
|
# ==============================================================================
|
|||
|
|
"""This tool analyzes a TensorFlow Lite graph."""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
# pylint: disable=g-import-not-at-top
|
|||
|
|
if not os.path.splitext(__file__)[0].endswith(
|
|||
|
|
os.path.join("tflite_runtime", "analyzer")):
|
|||
|
|
# This file is part of tensorflow package.
|
|||
|
|
from tensorflow.lite.python import wrap_toco
|
|||
|
|
from tensorflow.lite.python.analyzer_wrapper import _pywrap_analyzer_wrapper as _analyzer_wrapper
|
|||
|
|
from tensorflow.python.util.tf_export import tf_export as _tf_export
|
|||
|
|
else:
|
|||
|
|
# This file is part of tflite_runtime package.
|
|||
|
|
from tflite_runtime import _pywrap_analyzer_wrapper as _analyzer_wrapper
|
|||
|
|
|
|||
|
|
def _tf_export(*x, **kwargs):
|
|||
|
|
del x, kwargs
|
|||
|
|
return lambda x: x
|
|||
|
|
|
|||
|
|
|
|||
|
|
@_tf_export("lite.experimental.Analyzer")
|
|||
|
|
class ModelAnalyzer():
|
|||
|
|
"""Provides a collection of TFLite model analyzer tools.
|
|||
|
|
|
|||
|
|
Example:
|
|||
|
|
|
|||
|
|
```python
|
|||
|
|
model = tf.keras.applications.MobileNetV3Large()
|
|||
|
|
fb_model = tf.lite.TFLiteConverterV2.from_keras_model(model).convert()
|
|||
|
|
tf.lite.experimental.Analyzer.analyze(model_content=fb_model)
|
|||
|
|
# === TFLite ModelAnalyzer ===
|
|||
|
|
#
|
|||
|
|
# Your TFLite model has ‘1’ subgraph(s). In the subgraph description below,
|
|||
|
|
# T# represents the Tensor numbers. For example, in Subgraph#0, the MUL op
|
|||
|
|
# takes tensor #0 and tensor #19 as input and produces tensor #136 as output.
|
|||
|
|
#
|
|||
|
|
# Subgraph#0 main(T#0) -> [T#263]
|
|||
|
|
# Op#0 MUL(T#0, T#19) -> [T#136]
|
|||
|
|
# Op#1 ADD(T#136, T#18) -> [T#137]
|
|||
|
|
# Op#2 CONV_2D(T#137, T#44, T#93) -> [T#138]
|
|||
|
|
# Op#3 HARD_SWISH(T#138) -> [T#139]
|
|||
|
|
# Op#4 DEPTHWISE_CONV_2D(T#139, T#94, T#24) -> [T#140]
|
|||
|
|
# ...
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
WARNING: Experimental interface, subject to change.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def analyze(model_path=None,
|
|||
|
|
model_content=None,
|
|||
|
|
gpu_compatibility=False,
|
|||
|
|
**kwargs):
|
|||
|
|
"""Analyzes the given tflite_model with dumping model structure.
|
|||
|
|
|
|||
|
|
This tool provides a way to understand users' TFLite flatbuffer model by
|
|||
|
|
dumping internal graph structure. It also provides additional features
|
|||
|
|
like checking GPU delegate compatibility.
|
|||
|
|
|
|||
|
|
WARNING: Experimental interface, subject to change.
|
|||
|
|
The output format is not guaranteed to stay stable, so don't
|
|||
|
|
write scripts to this.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
model_path: TFLite flatbuffer model path.
|
|||
|
|
model_content: TFLite flatbuffer model object.
|
|||
|
|
gpu_compatibility: Whether to check GPU delegate compatibility.
|
|||
|
|
**kwargs: Experimental keyword arguments to analyze API.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Print analyzed report via console output.
|
|||
|
|
"""
|
|||
|
|
if not model_path and not model_content:
|
|||
|
|
raise ValueError("neither `model_path` nor `model_content` is provided")
|
|||
|
|
if model_path:
|
|||
|
|
print(f"=== {model_path} ===\n")
|
|||
|
|
tflite_model = model_path
|
|||
|
|
input_is_filepath = True
|
|||
|
|
else:
|
|||
|
|
print("=== TFLite ModelAnalyzer ===\n")
|
|||
|
|
tflite_model = model_content
|
|||
|
|
input_is_filepath = False
|
|||
|
|
|
|||
|
|
if kwargs.get("experimental_use_mlir", False):
|
|||
|
|
print(
|
|||
|
|
wrap_toco.wrapped_flat_buffer_file_to_mlir(tflite_model,
|
|||
|
|
input_is_filepath))
|
|||
|
|
else:
|
|||
|
|
print(
|
|||
|
|
_analyzer_wrapper.ModelAnalyzer(tflite_model, input_is_filepath,
|
|||
|
|
gpu_compatibility))
|