wehub-resource-sync
365 lines
12 KiB
Python
365 lines
12 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# 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.
|
|
#
|
|
|
|
|
|
import tensorflow as tf
|
|
from collections import deque
|
|
from typing import List
|
|
import os
|
|
import shutil
|
|
from tf2onnx import tf_loader, utils, convert
|
|
import copy
|
|
|
|
|
|
def ensure_and_clean_dir(dir_path, do_clean_dir=True) -> None:
|
|
"""Create a directory to save test logs
|
|
|
|
Args:
|
|
dir_path (str): directory to create / clean.
|
|
do_clean_dir (bool): boolean indicating whether to clean the directory if it already exists (remove+create).
|
|
Returns:
|
|
None
|
|
"""
|
|
if not os.path.exists(dir_path):
|
|
os.makedirs(dir_path)
|
|
elif do_clean_dir:
|
|
shutil.rmtree(dir_path)
|
|
os.makedirs(dir_path)
|
|
|
|
|
|
class Folder:
|
|
"""
|
|
Folder class that tracks all files for a single experiment.
|
|
"""
|
|
|
|
def __init__(self, folder_name) -> None:
|
|
self.base = folder_name
|
|
ensure_and_clean_dir(self.base)
|
|
self.fp32 = os.path.join(self.base, "fp32")
|
|
ensure_and_clean_dir(self.fp32)
|
|
self.fp32_saved_model = os.path.join(
|
|
self.fp32, "saved_model"
|
|
) # location of fp32 saved keras model
|
|
self.fp32_onnx_model = os.path.join(
|
|
self.fp32, "original.onnx"
|
|
) # location of fp32 onnx model
|
|
|
|
self.int8 = os.path.join(self.base, "int8")
|
|
ensure_and_clean_dir(self.int8)
|
|
self.int8_saved_model = os.path.join(
|
|
self.int8, "saved_model"
|
|
) # location of int8 saved keras model
|
|
self.int8_onnx_model = os.path.join(
|
|
self.int8, "quantized.onnx"
|
|
) # location of int8 onnx model
|
|
|
|
|
|
class CreateAssetsFolders:
|
|
"""Create empty folders to save the original and quantized TensorFlow models and their respective ONNX
|
|
models for each experiment.
|
|
|
|
The following directory structure is created: base_directory -> experiment_directory (created by `add_folder` method) -> (fp32 [saved_model, .onnx model]),
|
|
(int8 [saved_model, .onnx model]).
|
|
"""
|
|
|
|
def __init__(self, base_experiment_directory) -> None:
|
|
self.base = base_experiment_directory
|
|
if not os.path.exists(self.base):
|
|
os.mkdir(self.base)
|
|
|
|
def add_folder(self, folder_name: str) -> None:
|
|
"""
|
|
Create the experiment directory (sub-folder in the base directory passed to this class).
|
|
|
|
Args:
|
|
folder_name (str): name of folder
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
setattr(self, folder_name, Folder(os.path.join(self.base, folder_name)))
|
|
|
|
|
|
def convert_saved_model_to_onnx(
|
|
saved_model_dir: str, onnx_model_path: str, opset=13
|
|
) -> None:
|
|
"""Convert Keras saved model into ONNX format.
|
|
Works directly with CreateAssetsFolder object path.
|
|
|
|
Args:
|
|
saved_model_dir (str): Path to keras saved model.
|
|
onnx_model_path (str): Full path to ONNX model file.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
# 1. Let TensorRT optimize QDQ nodes instead of TF
|
|
from tf2onnx.optimizer import _optimizers
|
|
|
|
updated_optimizers = copy.deepcopy(_optimizers)
|
|
del updated_optimizers["q_dq_optimizer"]
|
|
del updated_optimizers["const_dequantize_optimizer"]
|
|
|
|
# 2. Extract graph definition from SavedModel
|
|
graph_def, inputs, outputs = tf_loader.from_saved_model(
|
|
model_path=saved_model_dir,
|
|
input_names=None,
|
|
output_names=None,
|
|
tag="serve",
|
|
signatures=["serving_default"],
|
|
)
|
|
|
|
# 3. Convert tf2onnx and save onnx file
|
|
model_proto, _ = convert._convert_common(
|
|
graph_def,
|
|
opset=opset,
|
|
input_names=inputs,
|
|
output_names=outputs,
|
|
output_path=onnx_model_path,
|
|
optimizers=updated_optimizers,
|
|
)
|
|
|
|
utils.save_protobuf(onnx_model_path, model_proto)
|
|
print("ONNX conversion Done!")
|
|
|
|
|
|
def convert_keras_model_to_onnx(
|
|
keras_model: tf.keras.Model, onnx_model_path: str, opset=13
|
|
) -> None:
|
|
"""Convert in-memory Keras model into ONNX format.
|
|
Works directly with CreateAssetsFolder object path.
|
|
|
|
Args:
|
|
keras_model (tf.keras.Model): Keras model.
|
|
onnx_model_path (str): Full path to ONNX model file.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
# 1. Let TensorRT optimize QDQ nodes instead of TF
|
|
from tf2onnx.optimizer import _optimizers
|
|
|
|
updated_optimizers = copy.deepcopy(_optimizers)
|
|
del updated_optimizers["q_dq_optimizer"]
|
|
del updated_optimizers["const_dequantize_optimizer"]
|
|
|
|
# 2. Convert keras model directly and save onnx file.
|
|
onnx_model_proto, _ = convert.from_keras(keras_model, opset=opset, optimizers=updated_optimizers)
|
|
utils.save_protobuf(onnx_model_path, onnx_model_proto)
|
|
|
|
|
|
class KerasModelTraveller:
|
|
"""
|
|
Utility class to travel Keras model and print out detailed layer information.
|
|
"""
|
|
|
|
def __init__(self, print_layer_config=False) -> None:
|
|
self._pc = print_layer_config
|
|
self.model_list = deque([])
|
|
# Used to filter which classes you want printed, by layer.__class__
|
|
self._filter_by_class = None
|
|
self._layer_names = []
|
|
self._print_basic_info = None
|
|
|
|
def _print_layer_info(self, layer):
|
|
assert isinstance(layer, tf.keras.layers.Layer)
|
|
if self._filter_by_class is None or layer.__class__ in self._filter_by_class:
|
|
self._layer_names.append(layer.name)
|
|
if self._print_basic_info:
|
|
print(
|
|
"layer name:{layer_name}, layer class:{layer_class}".format(
|
|
layer_name=layer.name, layer_class=layer.__class__
|
|
)
|
|
)
|
|
if self._pc:
|
|
print(layer.get_config())
|
|
if self._print_basic_info:
|
|
print("-----------------")
|
|
|
|
def _dissect(self):
|
|
if not self.model_list:
|
|
return
|
|
number_of_models = len(self.model_list)
|
|
for _ in range(number_of_models):
|
|
# Get a subclassed model
|
|
current_model = self.model_list.pop()
|
|
print("Keras Subclassed Model: {}".format(current_model.__class__.__name__))
|
|
assert isinstance(current_model, tf.keras.Model)
|
|
for l in current_model.layers:
|
|
if isinstance(l, tf.keras.Model):
|
|
# This is another subclassed model inside
|
|
# Add this model to model queue for further analysis
|
|
self.model_list.appendleft(l)
|
|
self._dissect()
|
|
else:
|
|
# This is a layer
|
|
self._print_layer_info(l)
|
|
|
|
def _travel(
|
|
self, keras_model: tf.keras.Model, filter_by_class=None, print_basic_info=False
|
|
):
|
|
"""Gets layer info by dissecting the model (need for multi-layered models)
|
|
|
|
Args:
|
|
keras_model (tf.keras.Model): Keras model
|
|
filter_by_class (str): None or array of layer.__class__ to print
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
self.filter_by_class = filter_by_class
|
|
self._print_basic_info = print_basic_info
|
|
assert isinstance(
|
|
keras_model, tf.keras.Model
|
|
), "Model passed is not Keras model"
|
|
self.model_list.appendleft(keras_model)
|
|
self._dissect()
|
|
self.filter_by_class = None
|
|
|
|
def get_layer_names(self, keras_model: tf.keras.Model, filter_by_class=None):
|
|
"""Get name of all layers in the model.
|
|
|
|
Args:
|
|
keras_model (tf.keras.Model): Keras model
|
|
filter_by_class (str): None or array of layer.__class__ to print
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
self._travel(keras_model=keras_model, filter_by_class=filter_by_class)
|
|
return self._layer_names
|
|
|
|
def get_layer_information(self, keras_model: tf.keras.Model, filter_by_class=None):
|
|
"""Print information about all layers.
|
|
|
|
Args:
|
|
keras_model (tf.keras.Model): Keras model
|
|
filter_by_class (str): None or array of layer.__class__ to print
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
self._travel(
|
|
keras_model=keras_model,
|
|
filter_by_class=filter_by_class,
|
|
print_basic_info=True,
|
|
)
|
|
|
|
|
|
def _get_layer_info(layer: tf.keras.layers.Layer) -> dict:
|
|
"""
|
|
Returns the layer's class, module, and name
|
|
"""
|
|
return {
|
|
"class": layer.__class__.__name__,
|
|
"module": layer.__class__.__module__,
|
|
"name": layer.name,
|
|
"layer": layer,
|
|
}
|
|
|
|
|
|
def _get_previous_layers_class_and_module_and_name(
|
|
layer: tf.keras.layers.Layer,
|
|
) -> List[dict]:
|
|
"""
|
|
For a given layer return a dictionary with name, module and class information of all previous layers.
|
|
"""
|
|
r = []
|
|
if isinstance(layer.input, list):
|
|
for layer_input_tensor in layer.input:
|
|
ip_tensor_parent_layer = layer_input_tensor._keras_history.layer
|
|
r.append(_get_layer_info(ip_tensor_parent_layer))
|
|
else:
|
|
ip_tensor_parent_layer = layer.input._keras_history.layer
|
|
r.append(_get_layer_info(ip_tensor_parent_layer))
|
|
return r
|
|
|
|
|
|
def find_my_predecessors(model: tf.keras.Model, current_layer_name: str) -> List[dict]:
|
|
"""
|
|
Given a layer name, find all predecessors of that layer.
|
|
|
|
Args:
|
|
model (tf.keras.Model): Keras functional model
|
|
current_layer_name (str): name of a model layer for which predecessors has to be found.
|
|
|
|
Returns:
|
|
List[dict]: List of predecessors. Each dictionary has three keys as follows,
|
|
::
|
|
{'class':<pred_layer_class>, 'module':<pred_layer_module>, 'name':<pred_layer_name>}
|
|
|
|
Raises:
|
|
AssertionError: If model is subclassed or current_layer_name is not string.
|
|
"""
|
|
supported_model_classes = {"Functional", "Sequential"}
|
|
assert isinstance(current_layer_name, str), "current layer name should be passed."
|
|
assert (
|
|
model.__class__.__name__ in supported_model_classes
|
|
), "model should be Functional or Sequential."
|
|
|
|
for layer in model.layers:
|
|
if layer.name == current_layer_name:
|
|
return _get_previous_layers_class_and_module_and_name(layer)
|
|
|
|
|
|
def find_my_successors(model: tf.keras.Model, current_layer_name: str) -> List[dict]:
|
|
"""
|
|
Given a layer name, find all successors of that layer.
|
|
|
|
Args:
|
|
model (tf.keras.Model): Keras functional model
|
|
current_layer_name (str): name of a model layer for which successors has to be found.
|
|
|
|
Returns:
|
|
List[dict]: List of predecessors. Each dictionary has three keys as follows,
|
|
::
|
|
{'class':<pred_layer_class>, 'module':<pred_layer_module>, 'name':<pred_layer_name>}
|
|
|
|
Raises:
|
|
AssertionError: If model is subclassed or current_layer_name is not string.
|
|
"""
|
|
supported_model_classes = {"Functional", "Sequential"}
|
|
assert isinstance(current_layer_name, str), "current layer name should be passed."
|
|
assert (
|
|
model.__class__.__name__ in supported_model_classes
|
|
), "model should be Functional or Sequential."
|
|
|
|
def _check_all_next_layers_with_connection_to_current(
|
|
next_layers: List[tf.keras.layers.Layer],
|
|
current_layer_name: str,
|
|
current_layer_class: str,
|
|
):
|
|
successors = []
|
|
for layer in next_layers:
|
|
p_layers = _get_previous_layers_class_and_module_and_name(layer)
|
|
for p_layer in p_layers:
|
|
if (
|
|
p_layer["class"] == current_layer_class
|
|
and p_layer["name"] == current_layer_name
|
|
):
|
|
successors.append(_get_layer_info(layer))
|
|
return successors
|
|
|
|
all_layers = model.layers
|
|
for i, layer in enumerate(all_layers):
|
|
if layer.name == current_layer_name:
|
|
next_layers = all_layers[i + 1 :]
|
|
layer_info = _get_layer_info(layer)
|
|
return _check_all_next_layers_with_connection_to_current(
|
|
next_layers, layer_info["name"], layer_info["class"]
|
|
)
|