chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Waiting to run
Docker Image CI / build-ubuntu2004 (push) Waiting to run
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,37 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
G_NUM_BITS: int = 8
|
||||
G_NARROW_RANGE: bool = True
|
||||
G_SYMMETRIC: bool = True
|
||||
|
||||
from tensorflow_quantization.quantize import QuantizationSpec
|
||||
from tensorflow_quantization.quantize import quantize_model
|
||||
|
||||
from tensorflow_quantization.quantize_wrapper_base import BaseQuantizeWrapper
|
||||
|
||||
from tensorflow_quantization.quantize_wrappers import WeightedBaseQuantizeWrapper
|
||||
from tensorflow_quantization.quantize_wrappers import NonWeightedBaseQuantizeWrapper
|
||||
|
||||
from tensorflow_quantization.custom_qdq_case_base import CustomQDQInsertionCase
|
||||
|
||||
from tensorflow_quantization.utils import CreateAssetsFolders
|
||||
from tensorflow_quantization.utils import convert_saved_model_to_onnx
|
||||
from tensorflow_quantization.utils import convert_keras_model_to_onnx
|
||||
|
||||
from .version import __version__
|
||||
@@ -0,0 +1,48 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
from abc import ABC
|
||||
|
||||
|
||||
class CustomQDQInsertionCase(ABC):
|
||||
"""
|
||||
This class helps user to programatically decide toolkit behavior to quantize specific layers.
|
||||
Based on the output of this class 'case' function, toolkit deviates from its standard behavior.
|
||||
"""
|
||||
|
||||
def info(self) -> str:
|
||||
return ""
|
||||
|
||||
def case(
|
||||
self, keras_model: "tf.keras.Model", qspec: "QuantizationSpec"
|
||||
) -> "QuantizationSpec":
|
||||
"""
|
||||
This function is called internally by the framework.
|
||||
Given keras model is passed as an argument and object of QuantizationSpec class
|
||||
is expcted in return.
|
||||
Returned QuantzaionSpec class object should contain information about the layers that needs
|
||||
to be treated specially/differently from default framework behavior.
|
||||
|
||||
Args:
|
||||
keras_model (tf.keras.Model): Keras functional or sequentail model
|
||||
qspec (QuantizationSpec): User passed QuantizationSpec object. It is important to note that
|
||||
new special qdq might or might not use quantizations specs user has provided.
|
||||
Returns:
|
||||
A new QuantizationSpec object.
|
||||
"""
|
||||
raise NotImplementedError("case method must be overridden by user")
|
||||
@@ -0,0 +1,363 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
|
||||
from tensorflow_quantization import CustomQDQInsertionCase
|
||||
from tensorflow_quantization import QuantizationSpec
|
||||
from tensorflow_quantization import utils
|
||||
import tensorflow as tf
|
||||
from typing import List
|
||||
|
||||
|
||||
def is_parent_type(parent_class: str, class_type="Conv") -> bool:
|
||||
"""
|
||||
Checks if 'parent_class' is of type 'type'.
|
||||
Examples of types: Conv, BatchNorm, Dropout, Activation.
|
||||
"""
|
||||
return class_type in parent_class
|
||||
|
||||
|
||||
def is_parent_pattern(parent_info: dict, pattern: List = ["BatchNorm", "Conv"]) -> bool:
|
||||
""" Checks if parent heritage follows a specific 'pattern'.
|
||||
Args:
|
||||
parent_info (dict): dictionary with parent's information.
|
||||
pattern (List): list containing a layer's parental heritage ([parent, grandparent, great-grandparent, ...]).
|
||||
Returns:
|
||||
bool: indicating whether a layer's parent heritage follows the given pattern.
|
||||
"""
|
||||
grandparent_info = parent_info
|
||||
for i, p in enumerate(pattern):
|
||||
if i > 0:
|
||||
grandparent_info = utils._get_previous_layers_class_and_module_and_name(
|
||||
grandparent_info["layer"]
|
||||
)[0]
|
||||
if not is_parent_type(grandparent_info["class"], class_type=p):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def check_is_quantizable_by_layer_name(
|
||||
qspec: QuantizationSpec, current_layer_name: str
|
||||
) -> bool:
|
||||
"""
|
||||
Checks if 'current_layer_name' is in 'qspec'. It returns True if 'current_layer_name' is NOT in 'qspec' and
|
||||
False if it is. This means that the user's request will get prioritized over our automatic methods.
|
||||
|
||||
Args:
|
||||
qspec (QuantizationSpec): quantization specification.
|
||||
|
||||
Returns:
|
||||
is_quantizable_by_layer_name (bool): boolean indicating whether 'current_layer_name' is quantizable by our
|
||||
method (is NOT in 'qspec'), or not (is in 'qspec', so that configuration should be followed).
|
||||
"""
|
||||
def _is_layer_in_user_passed_qspec(layer_name):
|
||||
for l in qspec.layers:
|
||||
if l.name == layer_name:
|
||||
return True
|
||||
return False
|
||||
|
||||
is_quantizable_by_layer_name = qspec is None or (
|
||||
qspec is not None and not _is_layer_in_user_passed_qspec(current_layer_name)
|
||||
)
|
||||
return is_quantizable_by_layer_name
|
||||
|
||||
|
||||
###################################################################
|
||||
################# General Custom QDQ Cases ########################
|
||||
###################################################################
|
||||
|
||||
|
||||
class BNQDQCase(CustomQDQInsertionCase):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def info(self) -> str:
|
||||
return "Avoids inserting QDQ before BatchNorm in cases where BN is connected to a Conv layer (since that BN " \
|
||||
"will be fused with previous Conv layer). This case happens in ResNet-v2, where the following pattern " \
|
||||
"exists: BN-ReLU-Conv blocks (pre-activation function). In that scenario, BN is sometimes connected to " \
|
||||
"`Add` layer, which doesn't fuse with BN."
|
||||
|
||||
def case(
|
||||
self, keras_model: tf.keras.Model, qspec: QuantizationSpec
|
||||
) -> QuantizationSpec:
|
||||
def _check_if_quantizable_bn(layer):
|
||||
layer_parent = layer.input._keras_history.layer
|
||||
parent_class_name = layer_parent.__class__.__name__
|
||||
if not is_parent_type(parent_class_name, class_type="Conv"):
|
||||
if check_is_quantizable_by_layer_name(qspec, layer.name):
|
||||
return True
|
||||
return False
|
||||
|
||||
bn_qspec = QuantizationSpec()
|
||||
for layer in keras_model.layers:
|
||||
if isinstance(layer, tf.keras.layers.BatchNormalization):
|
||||
"""
|
||||
Returns quantizable BatchNorm layers: All BN layers that are not connected to a Conv layer.
|
||||
In other words, don't add QDQ to BN layers in a Conv-BN sequence (and of course, if it shouldn't be
|
||||
ignored due to the user's preference)."
|
||||
"""
|
||||
if _check_if_quantizable_bn(layer):
|
||||
bn_qspec.add(
|
||||
name=layer.name, quantize_input=True, quantize_weight=False
|
||||
)
|
||||
return bn_qspec
|
||||
|
||||
|
||||
class ResidualConnectionQDQCase(CustomQDQInsertionCase):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def info(self) -> str:
|
||||
info_str = "Goal: To return all quantizable residual inputs. " \
|
||||
"Rules: Residual connection is represented by the Add layer. The recommendation from the TRT team " \
|
||||
" is to add QDQ to all of its inputs except when: " \
|
||||
" - the input is Bias. Note that TF sees MatMul+BiasAdd as a Dense layer, so no need to check " \
|
||||
" if the input is Bias. " \
|
||||
" - in the case of one of the inputs being a simple residual branch and the other Conv or " \
|
||||
" Conv+BN, add QDQ nodes to just the residual branch. This is needed to trigger an INT8 " \
|
||||
" kernel fusion with Add. " \
|
||||
" - in the case of more than one input being Conv or Conv+BN, add QDQ to all inputs except 1. " \
|
||||
" The last 2 cases are needed to trigger an INT8 kernel fusion with Add. " \
|
||||
" [ResNet-v1]: Note that the connection between Conv2D and Add layer is not direct: " \
|
||||
" Conv2D -> BatchNormalization -> Add " \
|
||||
" To get to the layer, we need to access `input._keras_history.layer` " \
|
||||
" This is the same for EfficientNet-B0. " \
|
||||
" [ResNet-v2]: Connection is direct ReLU -> Conv2D -> Add " \
|
||||
" [EfficientNet-B0]: Contains two special patterns: " \
|
||||
" 1. Conv -> BatchNorm -> Activation -> Add " \
|
||||
" 2. Conv -> BatchNorm -> Activation -> Dropout -> Add"
|
||||
return info_str
|
||||
|
||||
def case(
|
||||
self, keras_model: tf.keras.Model, qspec: QuantizationSpec
|
||||
) -> QuantizationSpec:
|
||||
|
||||
res_qspec = QuantizationSpec()
|
||||
for layer in keras_model.layers:
|
||||
if isinstance(layer, tf.keras.layers.Add) and check_is_quantizable_by_layer_name(qspec, layer.name):
|
||||
"""
|
||||
Returns quantizable inputs to Add layers: all inputs except 1 with 'pattern'.
|
||||
Patterns checked for: Conv, Conv-BN, Conv-BN-Activation, Conv-BN-Activation-Dropout.
|
||||
"""
|
||||
layer_parents = utils.find_my_predecessors(keras_model, layer.name)
|
||||
|
||||
# Collect the non-quantizable input (1 branch with Conv pattern)
|
||||
input_indices_convs = []
|
||||
for i, l_parent_info in enumerate(layer_parents):
|
||||
l_parent_class = l_parent_info["class"]
|
||||
l_parent_layer = l_parent_info["layer"]
|
||||
# Check that the input is a Conv pattern
|
||||
if (
|
||||
is_parent_type(l_parent_class, class_type="Conv")
|
||||
or is_parent_pattern(l_parent_info, pattern=["BatchNorm", "Conv"])
|
||||
or is_parent_pattern(l_parent_info, pattern=["Activation", "BatchNorm", "Conv"])
|
||||
or is_parent_pattern(l_parent_info, pattern=["Dropout", "Activation", "BatchNorm", "Conv"])
|
||||
):
|
||||
# Check that it's not a residual branch (input does not have more than 1 outbound node)
|
||||
if hasattr(l_parent_layer, 'outbound_nodes'):
|
||||
num_outbound_nodes = len(getattr(l_parent_layer, 'outbound_nodes'))
|
||||
if num_outbound_nodes == 1:
|
||||
# Branch without QDQ branch is chosen
|
||||
input_indices_convs.append(i)
|
||||
break
|
||||
|
||||
# Default behavior: add QDQ in all inputs except 1 with Conv/BN
|
||||
input_indices = list(range(0, len(layer_parents)))
|
||||
if len(input_indices_convs) > 0:
|
||||
# Don't quantize one of the Conv pattern branches.
|
||||
index_to_delete = input_indices_convs[-1]
|
||||
del input_indices[index_to_delete]
|
||||
if len(input_indices) > 0:
|
||||
res_qspec.add(
|
||||
layer.name,
|
||||
quantize_input=True,
|
||||
quantize_weight=False,
|
||||
quantization_index=input_indices,
|
||||
)
|
||||
return res_qspec
|
||||
|
||||
|
||||
class MaxPoolQDQCase(CustomQDQInsertionCase):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def info(self) -> str:
|
||||
return "Enables quantization of MaxPool layers. This is needed in cases where MaxPool is added to a residual " \
|
||||
"connection and where the other branches are already quantized (needed to trigger a horizontal fusion " \
|
||||
"in the residual connection. This case happens in ResNet-v2."
|
||||
|
||||
def case(
|
||||
self, keras_model: tf.keras.Model, qspec: QuantizationSpec
|
||||
) -> QuantizationSpec:
|
||||
mp_qspec = QuantizationSpec()
|
||||
for layer in keras_model.layers:
|
||||
if isinstance(layer, tf.keras.layers.MaxPooling2D):
|
||||
"""
|
||||
Returns quantizable MaxPooling2D layers.
|
||||
"""
|
||||
if check_is_quantizable_by_layer_name(qspec, layer.name):
|
||||
mp_qspec.add(
|
||||
name=layer.name,
|
||||
quantize_input=True,
|
||||
quantize_weight=False
|
||||
)
|
||||
return mp_qspec
|
||||
|
||||
|
||||
###################################################################
|
||||
############ Network Specific QDQ Cases ###########################
|
||||
###################################################################
|
||||
|
||||
|
||||
class ResNetV1QDQCase(CustomQDQInsertionCase):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def info(self) -> str:
|
||||
return (
|
||||
"Returns all quantizable nodes in ResNet-v1: "
|
||||
" 1. Residual connections."
|
||||
)
|
||||
|
||||
def case(
|
||||
self, keras_model: tf.keras.Model, qspec: QuantizationSpec
|
||||
) -> QuantizationSpec:
|
||||
special_qspec = QuantizationSpec()
|
||||
|
||||
# Use Residual connection QDQ
|
||||
residual_cqdq = ResidualConnectionQDQCase()
|
||||
residual_cqdq_qspec = residual_cqdq.case(keras_model, qspec)
|
||||
special_qspec.layers.extend(residual_cqdq_qspec.layers)
|
||||
|
||||
return special_qspec
|
||||
|
||||
|
||||
class ResNetV2QDQCase(CustomQDQInsertionCase):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def info(self) -> str:
|
||||
return (
|
||||
"Returns all quantizable nodes in ResNet-v2: "
|
||||
" 1. Residual connections, "
|
||||
" 2. BatchNorm not connected to Conv, "
|
||||
" 3. MaxPool layers."
|
||||
)
|
||||
|
||||
def case(
|
||||
self, keras_model: tf.keras.Model, qspec: QuantizationSpec
|
||||
) -> QuantizationSpec:
|
||||
special_qspec = QuantizationSpec()
|
||||
|
||||
# Use Residual connection QDQ
|
||||
residual_cqdq = ResidualConnectionQDQCase()
|
||||
residual_cqdq_qspec = residual_cqdq.case(keras_model, qspec)
|
||||
special_qspec.layers.extend(residual_cqdq_qspec.layers)
|
||||
|
||||
# Use BN QDQ Case
|
||||
bn_cqdq = BNQDQCase()
|
||||
bn_cqdq_qspec = bn_cqdq.case(keras_model, qspec)
|
||||
special_qspec.layers.extend(bn_cqdq_qspec.layers)
|
||||
|
||||
# Use MaxPool QDQ Case (necessary for ResNet-v2)
|
||||
mp_cqdq = MaxPoolQDQCase()
|
||||
mp_cqdq_qspec = mp_cqdq.case(keras_model, qspec)
|
||||
special_qspec.layers.extend(mp_cqdq_qspec.layers)
|
||||
|
||||
return special_qspec
|
||||
|
||||
|
||||
class EfficientNetQDQCase(CustomQDQInsertionCase):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def info(self) -> str:
|
||||
return (
|
||||
"Returns all quantizable nodes in EfficientNet:"
|
||||
" 1. Residual connections,"
|
||||
" 2. Quantize inputs (0, 1) of Multiply layers in SE (Squeeze-Excite) block."
|
||||
)
|
||||
|
||||
def case(
|
||||
self, keras_model: tf.keras.Model, qspec: QuantizationSpec
|
||||
) -> QuantizationSpec:
|
||||
special_qspec = QuantizationSpec()
|
||||
|
||||
# Use Residual connection QDQ
|
||||
residual_cqdq = ResidualConnectionQDQCase()
|
||||
residual_cqdq_qspec = residual_cqdq.case(keras_model, qspec)
|
||||
special_qspec.layers.extend(residual_cqdq_qspec.layers)
|
||||
|
||||
# Implement EfficientNet specific case to trigger horizontal fusion in Mul residual branch.
|
||||
# Gives preference to the user-specified `qspec`.
|
||||
for layer in keras_model.layers:
|
||||
if (
|
||||
isinstance(layer, tf.keras.layers.Multiply)
|
||||
and check_is_quantizable_by_layer_name(qspec, layer.name)
|
||||
):
|
||||
special_qspec.add(
|
||||
layer.name,
|
||||
quantize_input=True,
|
||||
quantize_weight=False,
|
||||
quantization_index=[0, 1],
|
||||
)
|
||||
return special_qspec
|
||||
|
||||
|
||||
class MobileNetQDQCase(CustomQDQInsertionCase):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def info(self) -> str:
|
||||
return (
|
||||
"Returns all quantizable nodes in MobileNet: "
|
||||
" 1. Residual connections."
|
||||
)
|
||||
|
||||
def case(
|
||||
self, keras_model: tf.keras.Model, qspec: QuantizationSpec
|
||||
) -> QuantizationSpec:
|
||||
special_qspec = QuantizationSpec()
|
||||
|
||||
# Use Residual connection QDQ
|
||||
residual_cqdq = ResidualConnectionQDQCase()
|
||||
residual_cqdq_qspec = residual_cqdq.case(keras_model, qspec)
|
||||
special_qspec.layers.extend(residual_cqdq_qspec.layers)
|
||||
|
||||
return special_qspec
|
||||
|
||||
|
||||
class InceptionQDQCase(CustomQDQInsertionCase):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def info(self) -> str:
|
||||
return (
|
||||
"Returns all quantizable nodes in Inception-v3: "
|
||||
" 1. MaxPool layers to trigger horizontal fusion in the output of Concat."
|
||||
)
|
||||
|
||||
def case(
|
||||
self, keras_model: tf.keras.Model, qspec: QuantizationSpec
|
||||
) -> QuantizationSpec:
|
||||
special_qspec = QuantizationSpec()
|
||||
|
||||
# Use MaxPool QDQ Case
|
||||
mp_cqdq = MaxPoolQDQCase()
|
||||
mp_cqdq_qspec = mp_cqdq.case(keras_model, qspec)
|
||||
special_qspec.layers.extend(mp_cqdq_qspec.layers)
|
||||
|
||||
return special_qspec
|
||||
@@ -0,0 +1,64 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""
|
||||
This module holds the quantization config class object that is accessed globally by library modules.
|
||||
DIRECT USE OF THIS MODULE BY USER IS PROHIBITED.
|
||||
"""
|
||||
|
||||
# List that holds quantization config class object, Length is always one!
|
||||
# Object is added automatically on class creation
|
||||
G_CONFIG_OBJECT = []
|
||||
|
||||
|
||||
def add_config_object(config_object: "BaseConfig") -> None:
|
||||
"""
|
||||
Add instance of quantize config class to the global list.
|
||||
Args:
|
||||
config_object : Instance of one of four quantize config class
|
||||
"""
|
||||
assert (
|
||||
len(G_CONFIG_OBJECT) == 0
|
||||
), "Looks like previous quatize object is alive. Did you call clear() on the object?"
|
||||
G_CONFIG_OBJECT.append(config_object)
|
||||
|
||||
|
||||
def remove_config_object() -> None:
|
||||
"""
|
||||
Remove instance of quantize config class from the global list.
|
||||
"""
|
||||
if G_CONFIG_OBJECT:
|
||||
G_CONFIG_OBJECT.clear()
|
||||
|
||||
|
||||
def get_config_object() -> "BaseConfig":
|
||||
"""
|
||||
Return quantize config class object
|
||||
"""
|
||||
assert (
|
||||
len(G_CONFIG_OBJECT) == 1
|
||||
), "Have you created quantize config object before calling `quantize_model`?"
|
||||
if G_CONFIG_OBJECT:
|
||||
return G_CONFIG_OBJECT[0]
|
||||
|
||||
|
||||
def is_config_object_created() -> bool:
|
||||
"""
|
||||
Sanity check function for whether quantize config class object is created.
|
||||
"""
|
||||
return len(G_CONFIG_OBJECT) == 1
|
||||
@@ -0,0 +1,445 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
# Copyright 2019 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.
|
||||
|
||||
from typing import List, Union
|
||||
import tensorflow as tf
|
||||
import tensorflow_quantization.quantize_wrappers as quantize_wrappers
|
||||
import tensorflow_quantization.global_config as cfg
|
||||
import tensorflow_quantization.quantize_config as quantize_config
|
||||
from dataclasses import dataclass
|
||||
from tensorflow_quantization.quantize_wrappers import DISABLED_LAYER_QUANTIZATION_DEFAULT
|
||||
|
||||
|
||||
@dataclass
|
||||
class LayerConfig:
|
||||
"""
|
||||
Internal dataclass for a single layer config.
|
||||
Args:
|
||||
name (str): Name of the layer. As seen from utilities such as `model.summary()`
|
||||
is_keras_class (bool) : Set this to True if layer_name passed represents a layer class from Keras.
|
||||
Default is False.
|
||||
quantize_input (bool): Set this to True if input to the layers should be quantized. Default is True
|
||||
since default behavior is following Nvidia quantization recipe.
|
||||
quantize_weight (bool): Set this to True if weights to the layers should be quantized. Default is True
|
||||
since default behavior is following Nvidia quantization recipe. For weightless layers, value is
|
||||
ignored.
|
||||
quantization_index (List): Indices on inputs to which quantization is applied for the layers with
|
||||
multiple inputs. E.g Add, Concatenate
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
|
||||
name: str = None
|
||||
is_keras_class: bool = False
|
||||
quantize_input: bool = True
|
||||
quantize_weight: bool = True
|
||||
quantization_index: list = None
|
||||
|
||||
|
||||
class QuantizationSpec:
|
||||
"""
|
||||
Helper class holding config objects for all layers to quantize.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.layers = []
|
||||
|
||||
def __str__(self) -> str:
|
||||
for l in self.layers:
|
||||
print(l)
|
||||
return ""
|
||||
|
||||
def add(
|
||||
self,
|
||||
name: Union[str, List],
|
||||
is_keras_class: Union[bool, List] = False,
|
||||
quantize_input: Union[bool, List] = True,
|
||||
quantize_weight: Union[bool, List] = True,
|
||||
quantization_index: Union[List, List[List]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Takes user parameters and adds LayerConfig object to a list for each add call.
|
||||
|
||||
Args:
|
||||
name (Union[str, List]): Name of the layer. As seen from utilities such as `model.summary()`
|
||||
is_keras_class (Union[bool, List]): List or a single value. Set this to True if layer_name passed represents a layer class from Keras.
|
||||
Default is False.
|
||||
quantize_input (Union[bool, List]): List or a single value. Set this to True if input to the layers should be quantized. Default is True
|
||||
since default behavior is following Nvidia quantization recipe.
|
||||
quantize_weight (Union[bool, List]): List or a single value. Set this to True if weights to the layers should be quantized. Default is True
|
||||
since default behavior is following Nvidia quantization recipe. For weightless layers, value is
|
||||
ignored.
|
||||
quantization_index (Union[List, List[List]]): List or List of List. List with indices on inputs to which quantization is applied for the layers with
|
||||
multiple inputs. E.g Add, Concatenate
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
if not isinstance(name, list):
|
||||
self.layers.append(
|
||||
LayerConfig(
|
||||
name=name,
|
||||
is_keras_class=is_keras_class,
|
||||
quantize_input=quantize_input,
|
||||
quantize_weight=quantize_weight,
|
||||
quantization_index=quantization_index,
|
||||
)
|
||||
)
|
||||
else:
|
||||
# layer names is passed as a list
|
||||
if isinstance(is_keras_class, list):
|
||||
assert len(name) == len(
|
||||
is_keras_class
|
||||
), "[E] `is_keras_class` is a list but length is not same as layer `name` list"
|
||||
if isinstance(quantize_input, list):
|
||||
assert len(name) == len(
|
||||
quantize_input
|
||||
), "[E] `quantize_input` is a list but length is not same as layer `name` list"
|
||||
if isinstance(quantize_weight, list):
|
||||
assert len(name) == len(
|
||||
quantize_weight
|
||||
), "[E] `quantize_weight` is a list but length is not same as layer `name` list"
|
||||
if isinstance(quantization_index, list):
|
||||
assert len(name) == len(
|
||||
quantization_index
|
||||
), "[E] `quantization_index` is list but length is not same as layer `name` list"
|
||||
for i, e in enumerate(name):
|
||||
cl_name = e
|
||||
cl_is_keras_class = (
|
||||
is_keras_class[i]
|
||||
if isinstance(is_keras_class, list)
|
||||
else is_keras_class
|
||||
)
|
||||
cl_quantize_input = (
|
||||
quantize_input[i]
|
||||
if isinstance(quantize_input, list)
|
||||
else quantize_input
|
||||
)
|
||||
cl_quantize_weight = (
|
||||
quantize_weight[i]
|
||||
if isinstance(quantize_weight, list)
|
||||
else quantize_weight
|
||||
)
|
||||
cl_quantization_index = (
|
||||
quantization_index[i]
|
||||
if isinstance(quantization_index, list)
|
||||
else quantization_index
|
||||
)
|
||||
self.layers.append(
|
||||
LayerConfig(
|
||||
name=cl_name,
|
||||
is_keras_class=cl_is_keras_class,
|
||||
quantize_input=cl_quantize_input,
|
||||
quantize_weight=cl_quantize_weight,
|
||||
quantization_index=cl_quantization_index,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _skip_layer(layer: tf.keras.layers.Layer) -> bool:
|
||||
"""
|
||||
Decide whether quantization wrapping should be skipped for the given layer.
|
||||
The decision is made based on an internal quantize config object parameters.
|
||||
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): Keras model layer
|
||||
Returns:
|
||||
bool: True if given layer should not be quantized else False
|
||||
"""
|
||||
config_object = cfg.get_config_object()
|
||||
|
||||
# Check if any layer with Disabled Quantization by default are in the 'config_object.layer_classes_to_quantize'.
|
||||
# If so, that layer will be enabled for quantization. Otherwise, skip (return True).
|
||||
layer_class_name = layer.__class__.__name__
|
||||
if layer_class_name in DISABLED_LAYER_QUANTIZATION_DEFAULT:
|
||||
if layer_class_name not in config_object.layer_classes_to_quantize:
|
||||
if layer.name in config_object.get_layer_config():
|
||||
# User can enable a single layer even if the default behavior of a Class is to not quantize.
|
||||
# The decision of whether to quantize this layer or not will be left for later checks, such as when
|
||||
# quantize_input and quantize_weight = False.
|
||||
pass
|
||||
else:
|
||||
# Default behavior: skip layer
|
||||
return True
|
||||
|
||||
# 1. When quantize_input = False, quantize_weight = False and quantization_index=None, don't even wrap the layer.
|
||||
if layer.name in config_object.get_layer_config():
|
||||
current_layer_config = config_object.get_layer_config()[layer.name]
|
||||
if (
|
||||
current_layer_config["qbool_list"][0] == False # quantize_input
|
||||
and current_layer_config["qbool_list"][1] == False # quantize_weight
|
||||
and "qindex_list" not in current_layer_config
|
||||
):
|
||||
print(
|
||||
"[I] Layer `{layer_name}` is not quantized. There is nothing to quantize since "
|
||||
"quantize_input = False, quantize_weight = False and quantization_index=None".format(
|
||||
layer_name=layer.name
|
||||
)
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
# 2. Called when quantization_mode is `partial`
|
||||
if config_object.config_class_id == 2:
|
||||
# A. Skip current `layer class` if current layer class is not in user provided QuantizationSpec class
|
||||
# object. However, when current layer name is passed by user to quantize, don't skip the layer.
|
||||
if (
|
||||
len(config_object.layer_classes_to_quantize) != 0
|
||||
and layer.__class__.__name__ not in config_object.layer_classes_to_quantize
|
||||
):
|
||||
if layer.name in config_object.get_layer_config():
|
||||
return False
|
||||
else:
|
||||
print(
|
||||
"[I] Layer class `{layer_class_name}` is not quantized. Partial quantization is enabled "
|
||||
"and layer class is not in user provided QuantizationSpec class object".format(
|
||||
layer_class_name=layer.__class__.__name__
|
||||
)
|
||||
)
|
||||
return True
|
||||
# B. Skip current layer if `layer.name` is not in user provided QuantizationSpec class object.
|
||||
# However, if current layer class is passed by user to quantize, don't skip the layer.
|
||||
elif layer.name not in config_object.get_layer_config():
|
||||
if layer.__class__.__name__ in config_object.layer_classes_to_quantize:
|
||||
return False
|
||||
else:
|
||||
print(
|
||||
"[I] Layer `{layer_name}` is not quantized. Partial quantization is enabled and layer name is not "
|
||||
"in user provided QuantizationSpec class object".format(
|
||||
layer_name=layer.name
|
||||
)
|
||||
)
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _quantize_model_layer_clone_function(
|
||||
layer: tf.keras.layers.Layer,
|
||||
) -> "BaseQuantizeWrapper":
|
||||
"""
|
||||
Wrap or leave given layer based on quantize config object parameters.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): Keras model layer
|
||||
Returns:
|
||||
BaseQuantizeWrapper: layer wrapped in BaseQuantizeWrapper class.
|
||||
"""
|
||||
layer_wrapper = layer
|
||||
if _skip_layer(layer):
|
||||
# Skip the layers not specified by the user.
|
||||
pass
|
||||
else:
|
||||
child_wrappers_dict = quantize_wrappers.BaseQuantizeWrapper.CHILD_WRAPPERS
|
||||
possible_wrapper_name_for_this_layer = (
|
||||
layer.__class__.__name__ + "QuantizeWrapper"
|
||||
)
|
||||
if possible_wrapper_name_for_this_layer in child_wrappers_dict:
|
||||
wrapper_function = child_wrappers_dict[possible_wrapper_name_for_this_layer]
|
||||
layer_wrapper = wrapper_function(layer)
|
||||
return layer_wrapper
|
||||
|
||||
|
||||
def _execute_quantize_model(
|
||||
model: tf.keras.Model, class_id: int, qspec: QuantizationSpec = None
|
||||
) -> tf.keras.Model:
|
||||
"""
|
||||
clone the model and apply quantization to specific layers based on quantize config object parameters.
|
||||
Args:
|
||||
model (tf.keras.Model): Keras functional or sequential model.
|
||||
* Currently Subclassed models are not supported
|
||||
class_id (int): internal quantization class ID
|
||||
qspec (QuantizationSpec): object of QuantizationSpec class. If few layers or layer classes are to be treated
|
||||
differently, LayerConfig class objects for that layer/layer class are created internally and
|
||||
added to QuantizationSpec class.
|
||||
Returns:
|
||||
tf.keras.Model: Quantized model with QDQ nodes added.
|
||||
"""
|
||||
config_id_class_name_map = {
|
||||
0: "FullNetworkQuantization",
|
||||
1: "FullNetworkSpecialQuantization",
|
||||
2: "PartialNetworkQuantization",
|
||||
}
|
||||
|
||||
# 1. Create quantize config object
|
||||
q_config_object = getattr(quantize_config, config_id_class_name_map[class_id])()
|
||||
|
||||
# 2. Update object attributes
|
||||
if qspec:
|
||||
q_config_object.add_quantization_spec_object(qspec, model.layers)
|
||||
|
||||
assert (
|
||||
cfg.is_config_object_created()
|
||||
), "[E] Have you created the quantization config object before calling `quantize_model`?"
|
||||
|
||||
# 3. Ensure that the original model is kept untouched.
|
||||
# This step is needed as `clone_model` with our custom `clone_function` wraps layers in a destructive manner.
|
||||
# TODO: delete later if a better solution is found, most likely inside our custom `clone_function`.
|
||||
cloned_model = tf.keras.models.clone_model(model)
|
||||
cloned_model.set_weights(model.get_weights())
|
||||
|
||||
# 4. Wrap quantizable layers
|
||||
quant_model = tf.keras.models.clone_model(
|
||||
cloned_model, input_tensors=None, clone_function=_quantize_model_layer_clone_function
|
||||
)
|
||||
|
||||
# 5. Clean global space afterwards
|
||||
q_config_object.clean()
|
||||
|
||||
return quant_model
|
||||
|
||||
|
||||
def _recognize_config_class_id(
|
||||
quantization_mode: str = "full", qspec: QuantizationSpec = None
|
||||
) -> int:
|
||||
"""
|
||||
Interpret internal quantize config class based on parameters passed by user to
|
||||
`quantize_model` function.
|
||||
Args:
|
||||
quantization_mode (str): Either 'full' or 'partial' quantization mode
|
||||
qspec (QuantizationSpec): object of QuantizationSpec class. If few layers or layer classes are to be treated
|
||||
differently, LayerConfig class objects for that layer/layer class are created internally and
|
||||
added to QuantizationSpec class.
|
||||
Returns:
|
||||
int: ID for quantization category class used internally.
|
||||
Raises:
|
||||
Exception: if no class can be interpreted for given parameter combination
|
||||
"""
|
||||
if quantization_mode == "full" and qspec is None:
|
||||
return 0
|
||||
elif quantization_mode == "full" and qspec is not None:
|
||||
return 1
|
||||
elif quantization_mode == "partial" and qspec is not None:
|
||||
return 2
|
||||
else:
|
||||
raise Exception(
|
||||
"Could not recognize config class ID."
|
||||
" Are parameters passed to `quantize_model` function correct?"
|
||||
)
|
||||
|
||||
|
||||
def _validate_config(
|
||||
quantization_mode: str = "full", qspec: QuantizationSpec = None
|
||||
) -> None:
|
||||
"""
|
||||
Validate if parameters passed to `quantize_model` makes sense.
|
||||
Args:
|
||||
quantization_mode (str): quantization mode can be either 'full' or 'partial'
|
||||
qspec (QuantizationSpec): object of QuantizationSpec class. If few layers or layer classes are to be treated
|
||||
differently, LayerConfig class objects for that layer/layer class are created internally and
|
||||
added to QuantizationSpec class.
|
||||
Returns:
|
||||
None
|
||||
Raises:
|
||||
AssertionError: when configuration is not valid.
|
||||
"""
|
||||
|
||||
def _verify_support_for_all_layer_classes(qspec: QuantizationSpec):
|
||||
for layer in qspec.layers:
|
||||
if layer.is_keras_class:
|
||||
# Layer class name is provided.
|
||||
child_wrappers_dict = (
|
||||
quantize_wrappers.BaseQuantizeWrapper.CHILD_WRAPPERS
|
||||
)
|
||||
possible_wrapper_name_for_this_layer = layer.name + "QuantizeWrapper"
|
||||
assert possible_wrapper_name_for_this_layer in child_wrappers_dict, (
|
||||
"[E] layer class `{layer_name}` is not supported yet! Either there is no native wrapper or user "
|
||||
"provided wrapper registration failed.".format(
|
||||
layer_name=layer.name
|
||||
)
|
||||
)
|
||||
|
||||
if qspec:
|
||||
_verify_support_for_all_layer_classes(qspec)
|
||||
|
||||
if quantization_mode == "partial":
|
||||
assert (
|
||||
qspec is not None
|
||||
), "[E] `QuantizationSpec` class object must be passed when `quantization_mode=partial`."
|
||||
|
||||
|
||||
def quantize_model(
|
||||
model,
|
||||
quantization_mode: str = "full",
|
||||
quantization_spec: QuantizationSpec = None,
|
||||
custom_qdq_cases: List["CustomQDQInsertionCase"] = None,
|
||||
) -> tf.keras.Model:
|
||||
"""
|
||||
Insert Q/DQ nodes in Keras model and return a copy. Weights are preserved unlike native keras clone.
|
||||
|
||||
Args:
|
||||
model(tf.keras.Model): Keras Functional or Sequential model.subclassed models are not yet supported.
|
||||
quantization_mode(str): quantization mode can be either 'full' or 'partial'
|
||||
quantization_spec(QuantizationSpec) : object of QuantizationSpec class. If few layers or layer classes are to
|
||||
be treated differently, LayerConfig class objects for that layer/layer class are created internally and
|
||||
added to QuantizationSpec class.
|
||||
custom_qdq_cases(List[CustomQDQInsertionCase]) : `Case` method on every object in this list is called by passing
|
||||
model and user passed quantization_spec as arguments. Each member of this list is an object of a class
|
||||
inherited from CustomQDQInsertionCase class.
|
||||
|
||||
Raises:
|
||||
AssertionError: When passed model is subclassed.
|
||||
AssertionError: When CustomQDQInsertionCase does not return QuantizationSpec object.
|
||||
AssertionError: When quantization mode is `partial` but QuantizationSpec object is not passed.
|
||||
AssertionError: When quantization wrapper is not found for desired layer class.
|
||||
ExceptionError: When internal quantization class ID can't be detected. This happens when passed parameters
|
||||
do not make sense.
|
||||
Returns:
|
||||
tf.keras.Model: Quantized model with QDQ nodes inserted according to NVIDIA quantization recipe.
|
||||
"""
|
||||
supported_model_classes = {"Functional", "Sequential"}
|
||||
assert (
|
||||
model.__class__.__name__ in supported_model_classes
|
||||
), "[E] Currently only `Functional` or `Sequential` model quantization is supported."
|
||||
|
||||
# Update quantization_spec object based on output of special QDQ cases.
|
||||
custom_quantization_spec = QuantizationSpec()
|
||||
if custom_qdq_cases:
|
||||
for custom_qdq_case in custom_qdq_cases:
|
||||
qspec_case_object = custom_qdq_case.case(model, quantization_spec)
|
||||
if qspec_case_object:
|
||||
assert isinstance(
|
||||
qspec_case_object, QuantizationSpec
|
||||
), "[E] {} \
|
||||
does not return an object of QuantizationSpec.".format(
|
||||
qspec_case_object.__class__.__name__
|
||||
)
|
||||
custom_quantization_spec.layers.extend(qspec_case_object.layers)
|
||||
|
||||
# if user has passed quantization_spec then extend it with custom_quantization_spec
|
||||
# else use just custom_quantization_spec
|
||||
if quantization_spec:
|
||||
quantization_spec.layers.extend(custom_quantization_spec.layers)
|
||||
else:
|
||||
if len(custom_quantization_spec.layers) != 0:
|
||||
quantization_spec = custom_quantization_spec
|
||||
|
||||
# Check if config is valid and quantize model
|
||||
_validate_config(quantization_mode, quantization_spec)
|
||||
cid = _recognize_config_class_id(quantization_mode, quantization_spec)
|
||||
return _execute_quantize_model(model, cid, quantization_spec)
|
||||
@@ -0,0 +1,273 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""
|
||||
This module implements classes to configure three supported quantization modes:
|
||||
1. Full: quantize all layers with standard protocol based NVIDIA quantization scheme.
|
||||
2. Full special: quantize few layers in a specific way and remaining with standard protocol based on NVIDIA
|
||||
quantization scheme.
|
||||
3. Partial: quantize ONLY few layers.
|
||||
|
||||
Each quantization mode can quantize all supported Keras layer classes or only subset of it.
|
||||
"""
|
||||
|
||||
from abc import ABC
|
||||
import tensorflow_quantization.global_config as global_config
|
||||
import warnings
|
||||
from typing import List, Dict
|
||||
|
||||
|
||||
class BaseConfig(ABC):
|
||||
"""
|
||||
Base class from which four quantize config classes are derived.
|
||||
Default quantization recipe is Nvidia's recommendation.
|
||||
"""
|
||||
|
||||
def __new__(cls):
|
||||
instance = super().__new__(cls)
|
||||
# Add instance to global list
|
||||
global_config.add_config_object(instance)
|
||||
return instance
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.quantization_mode: str = "full"
|
||||
self.layerwise_config: dict = {} # holds special layers information.
|
||||
self.layer_classes_to_quantize: set = set()
|
||||
self.config_class_id: int = 0
|
||||
|
||||
def __str__(self) -> str:
|
||||
return (
|
||||
" quantization_mode: {quant_mode} \n "
|
||||
"layerwise_config: {layerwise_config} \n "
|
||||
"specific_layer_class: {specific_layer_class} \n "
|
||||
"config_class_id: {config_class_id} \n".format(
|
||||
quant_mode=self.quantization_mode,
|
||||
layerwise_config=self.layerwise_config,
|
||||
specific_layer_class=self.specific_layer_class,
|
||||
config_class_id=self.config_class_id,
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _validate_layer_names(
|
||||
user_passed_layer_names: List, model_layers: List
|
||||
) -> None:
|
||||
"""
|
||||
Check whether user passed layer names exists in Keras model being quantized.
|
||||
Args:
|
||||
user_passed_layer_names (List): Layer names passed by user to treat specially.
|
||||
model_layers (List): Keras model layers passed as a list.
|
||||
Returns:
|
||||
None
|
||||
Raise:
|
||||
Warning : when specific layer name is not found. Such layers are simply ignored.
|
||||
"""
|
||||
model_layer_name_set = set()
|
||||
for l in model_layers:
|
||||
model_layer_name_set.add(l.name)
|
||||
|
||||
for ul in user_passed_layer_names:
|
||||
if ul not in model_layer_name_set:
|
||||
warnings.warn(
|
||||
"layer name {} is passed by user but could not find layer with this name in model.".format(
|
||||
ul
|
||||
)
|
||||
)
|
||||
|
||||
def add_quantization_spec_object(
|
||||
self, qspec: "QuantizationSpec", original_model_layers: List
|
||||
) -> None:
|
||||
"""
|
||||
This method parses object of QuantizationSpec class and fill in `layerwise_config` dictionary
|
||||
holding information about layers that need to be treated specially.
|
||||
Specific layer classes that need to be treated specially are also here.
|
||||
Args:
|
||||
qspec (QuantizationSpec): object of QuantizationSpec class. If few layers or layer classes are to be treated
|
||||
differently, LayerConfig class objects for that layer/layer class are created internally and
|
||||
added to QuantizationSpec class.
|
||||
original_model_layers (List): Keras model layers passed as a list.
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
for layer in qspec.layers:
|
||||
if layer.is_keras_class:
|
||||
self.add_special_layer_class(layer.name)
|
||||
else:
|
||||
layer_config_dict = {"qbool_list": [False, False]}
|
||||
layer_config_dict["qbool_list"][0] = layer.quantize_input
|
||||
layer_config_dict["qbool_list"][1] = layer.quantize_weight
|
||||
if layer.quantization_index:
|
||||
layer_config_dict["qindex_list"] = layer.quantization_index
|
||||
self.add_special_layer(
|
||||
layer_name=layer.name, config_dict=layer_config_dict
|
||||
)
|
||||
# Validate whether added layers exist in the model
|
||||
self._validate_layer_names(
|
||||
list(self.layerwise_config.keys()), original_model_layers
|
||||
)
|
||||
|
||||
def add_special_layer(self, layer_name: str, config_dict: Dict) -> None:
|
||||
"""
|
||||
Add layer specific quantization information to quantize config object.
|
||||
Args:
|
||||
layer_name (str): layer name
|
||||
config_dict (Dict): Layer specific quantization parameter dictionary in the
|
||||
following format.
|
||||
There are only two accepted keys `qbool_list` and `qindex_list`.
|
||||
`qbool_list` is list of length two where each value is
|
||||
[<True/False quantize inputs>, <True/False quantize weights>]
|
||||
e.g.
|
||||
To quantize inputs and weights, `qbool_list`=[True, True]
|
||||
|
||||
`qindex_list` is a list of specific indices to quaintize for layers such as Add, Concatenate
|
||||
where more than two inputs are present.
|
||||
|
||||
Based on above information,
|
||||
1. config_dict for weighted layer with name `dense_2`, to quantize inputs and weights will be
|
||||
{'qbool_list':[True, True]} with laye_name=`dense_2`
|
||||
2. config_dict for non weighted layer with name `add_3` to quantize input at index 1 will be
|
||||
{'qbool_list':[True, False], 'qindex_list':[1]} with layer_name=`add_3`
|
||||
Returns:
|
||||
None
|
||||
Raises:
|
||||
Exception: When invalid keys are detected.
|
||||
"""
|
||||
self.layerwise_config[layer_name] = config_dict
|
||||
|
||||
def remove_layer(self, layer_name: str) -> None:
|
||||
"""
|
||||
Remove specific layer based on name from quantize config object.
|
||||
Args:
|
||||
layer_name (str): layer name
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
if layer_name in self.layerwise_config:
|
||||
del self.layerwise_config[layer_name]
|
||||
|
||||
def remove_layers(self, layers_name: List) -> None:
|
||||
"""
|
||||
Bulk remove specific layers based on names from quantize config object.
|
||||
Args:
|
||||
layers_name (List): layers names, list of strings
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
for layer_name in layers_name:
|
||||
self.remove_layer(layer_name=layer_name)
|
||||
|
||||
def get_layer_config(self) -> Dict:
|
||||
"""
|
||||
Return dictionary with information about layers to quantize for quantize
|
||||
config object.
|
||||
Args:
|
||||
None
|
||||
Returns:
|
||||
Dict: a dictionary with layerwise configuration parameters.
|
||||
"""
|
||||
return self.layerwise_config
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
"""
|
||||
Return True if no layer specific quantization information is available in quantize
|
||||
config object.
|
||||
Args:
|
||||
None
|
||||
Returns:
|
||||
bool: True if no special layers are passed else return False
|
||||
"""
|
||||
return not self.layerwise_config
|
||||
|
||||
def clear_layer_config(self) -> None:
|
||||
"""
|
||||
Clear layer config information from quanize config object
|
||||
Args:
|
||||
None
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
self.layerwise_config.clear()
|
||||
|
||||
def add_special_layer_class(self, layer_class_name: str) -> None:
|
||||
"""
|
||||
Add class name to quantize config object so that only layers with specific class are quantized.
|
||||
Args:
|
||||
layer_class_name : String that represents keras class
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
self.layer_classes_to_quantize.add(layer_class_name)
|
||||
|
||||
def clean(self):
|
||||
"""
|
||||
Clean quantize config object from global space. Calling this is important to use `quantize_model` multiple times
|
||||
within a single module.
|
||||
Args:
|
||||
None
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
global_config.remove_config_object()
|
||||
|
||||
|
||||
class FullNetworkQuantization(BaseConfig):
|
||||
"""
|
||||
Quantize all layers based on NV scheme.
|
||||
|
||||
Nvidia recommended recipe for quantization is using Q/DQ only wth inputs/weights.
|
||||
Q/DQ output support is just to compare engine performance/accuracy when other quantization
|
||||
scheme is used.
|
||||
|
||||
NV: Add Q/DQ at input and weights
|
||||
TF: Add Q/DQ at output and weights
|
||||
|
||||
This is config class with index `0` which is default.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.config_class_id = 0
|
||||
|
||||
|
||||
class FullNetworkSpecialQuantization(BaseConfig):
|
||||
"""
|
||||
Quantize few layers in specific way and remaining network in standard way based on NV scheme.
|
||||
Layers are selected based on 'names' which can be via 'model.summary()' for functional
|
||||
and sequential models.
|
||||
Subclassed model layer information can be found using `KerasModelTraveller` class from utils.
|
||||
|
||||
This is config class with index 1.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.config_class_id = 1
|
||||
|
||||
|
||||
class PartialNetworkQuantization(BaseConfig):
|
||||
"""
|
||||
Quantize only specific layers and not the entire network.
|
||||
Layers are selected based on name.
|
||||
|
||||
This is config class with index 2.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.quantization_mode = "partial"
|
||||
self.config_class_id = 2
|
||||
@@ -0,0 +1,236 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
# Copyright 2019 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.
|
||||
|
||||
import tensorflow as tf
|
||||
import tensorflow_quantization.quantizers as quantizers
|
||||
import tensorflow_quantization.global_config as cfg
|
||||
from abc import abstractmethod
|
||||
|
||||
deserialize_keras_object = tf.keras.utils.deserialize_keras_object
|
||||
serialize_keras_object = tf.keras.utils.serialize_keras_object
|
||||
|
||||
NO_WEIGHT_LAYERS = {
|
||||
"Concatenate",
|
||||
"Add",
|
||||
"AveragePooling2D",
|
||||
"GlobalAveragePooling2D",
|
||||
"MaxPooling2D",
|
||||
"BatchNormalization",
|
||||
}
|
||||
|
||||
|
||||
class BaseQuantizeWrapper(tf.keras.layers.Wrapper):
|
||||
"""Base wrapper class which all layer wrappers inherit"""
|
||||
|
||||
CHILD_WRAPPERS = {}
|
||||
|
||||
def __init_subclass__(cls, **kwargs) -> None:
|
||||
super().__init_subclass__(**kwargs)
|
||||
cls.CHILD_WRAPPERS[cls.__name__] = cls
|
||||
|
||||
def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
|
||||
"""Create a quantize emulate wrapper for a keras layer.
|
||||
This wrapper provides options to quantize inputs and weights of the layer.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
if layer is None:
|
||||
raise ValueError("`layer` cannot be None.")
|
||||
|
||||
# Check against keras.Model since it is an instance of keras.layers.Layer.
|
||||
if not isinstance(layer, tf.keras.layers.Layer) or isinstance(
|
||||
layer, tf.keras.Model
|
||||
):
|
||||
raise ValueError(
|
||||
"`layer` can only be a `tf.keras.layers.Layer` instance. "
|
||||
"You passed an instance of type: {input}.".format(
|
||||
input=layer.__class__.__name__
|
||||
)
|
||||
)
|
||||
if "name" not in kwargs:
|
||||
kwargs["name"] = self._make_layer_name(layer)
|
||||
super(BaseQuantizeWrapper, self).__init__(layer, **kwargs)
|
||||
|
||||
# get quantize config object that holds all the information about how quantization should be performed.
|
||||
quantize_config_object = cfg.get_config_object()
|
||||
|
||||
# set all initial quantization parameters to False/None
|
||||
self.quantize_inputs = False
|
||||
self.quantize_weights = False
|
||||
self.quantize_specific_input_indices = None
|
||||
|
||||
layer_class_name_t = layer.__class__.__name__ # Layer class name
|
||||
layer_name_t = layer.name # Actual layer name
|
||||
|
||||
def _configure_singular_quantize():
|
||||
self.quantize_inputs = True
|
||||
if layer_class_name_t in NO_WEIGHT_LAYERS:
|
||||
self.quantize_weights = False
|
||||
else:
|
||||
self.quantize_weights = True
|
||||
|
||||
def _configure_special_quantize(
|
||||
quantize_bool_list: list, layer_name_t: str, index_list_if_any: list = None
|
||||
):
|
||||
assert (len(quantize_bool_list)) == 2, (
|
||||
"Three boolean values (representing whether to quantize [inputs, weights]) must be provided in "
|
||||
"quantize_config for layer: {layer_name_t}. If quantization does not apply for specific part, "
|
||||
"pass None. e.g. For layer ( e.g. Concatenate, Add) with no weights, `qbool_list` to quantize "
|
||||
"input can be [True, False]".format(layer_name_t=layer_name_t)
|
||||
)
|
||||
|
||||
self.quantize_inputs = quantize_bool_list[0]
|
||||
if layer_class_name_t in NO_WEIGHT_LAYERS:
|
||||
self.quantize_weights = False
|
||||
else:
|
||||
self.quantize_weights = quantize_bool_list[1]
|
||||
|
||||
if index_list_if_any:
|
||||
self.quantize_specific_input_indices = index_list_if_any
|
||||
|
||||
if quantize_config_object.config_class_id == 0:
|
||||
# This is straight forward full network quantization
|
||||
_configure_singular_quantize()
|
||||
else:
|
||||
# Config class id 1 or 2.
|
||||
# User has provided layer (name) specific quantization information
|
||||
quantize_config_dict = quantize_config_object.get_layer_config()
|
||||
if layer_name_t in quantize_config_dict:
|
||||
# This layer needs to be quantized in specific way
|
||||
if "qindex_list" in quantize_config_dict[layer_name_t]:
|
||||
_configure_special_quantize(
|
||||
quantize_config_dict[layer_name_t]["qbool_list"],
|
||||
layer_name_t,
|
||||
quantize_config_dict[layer_name_t]["qindex_list"],
|
||||
)
|
||||
else:
|
||||
_configure_special_quantize(
|
||||
quantize_config_dict[layer_name_t]["qbool_list"], layer_name_t
|
||||
)
|
||||
else:
|
||||
_configure_singular_quantize()
|
||||
|
||||
self._track_trackable(layer, name="layer")
|
||||
|
||||
@staticmethod
|
||||
def _make_layer_name(layer):
|
||||
return "{}_{}".format("quant", layer.name)
|
||||
|
||||
@staticmethod
|
||||
def _weight_name(name):
|
||||
"""Extracts the weight name from the full TensorFlow variable name.
|
||||
For example, returns 'kernel' for 'dense_2/kernel:0'.
|
||||
Args:
|
||||
name: TensorFlow variable name.
|
||||
Returns:
|
||||
Extracted weight name.
|
||||
"""
|
||||
return name.split(":")[0].split("/")[-1]
|
||||
|
||||
def build(self, input_shape):
|
||||
super(BaseQuantizeWrapper, self).build(input_shape)
|
||||
|
||||
self.optimizer_step = self.add_weight(
|
||||
"optimizer_step",
|
||||
initializer=tf.keras.initializers.Constant(-1),
|
||||
dtype=tf.dtypes.int32,
|
||||
trainable=False,
|
||||
)
|
||||
|
||||
def compute_output_shape(self, input_shape):
|
||||
return self.layer.compute_output_shape(self.layer.input_shape)
|
||||
|
||||
def _last_value_quantizer(
|
||||
self, x, training, quantizer_vars, per_channel=False, channel_axis=-1
|
||||
):
|
||||
"""Use currying to return True/False specialized fns to the cond."""
|
||||
from tensorflow_quantization import G_NUM_BITS, G_SYMMETRIC, G_NARROW_RANGE
|
||||
|
||||
return quantizers.LastValueQuantize(
|
||||
x,
|
||||
quantizer_vars["min_var"],
|
||||
quantizer_vars["max_var"],
|
||||
per_channel=per_channel,
|
||||
channel_axis=channel_axis,
|
||||
is_training=training,
|
||||
num_bits=G_NUM_BITS,
|
||||
narrow_range=G_NARROW_RANGE,
|
||||
symmetric=G_SYMMETRIC,
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def call(self, inputs, training=None):
|
||||
raise NotImplementedError
|
||||
|
||||
def get_config(self):
|
||||
base_config = super(BaseQuantizeWrapper, self).get_config()
|
||||
config = {"quantize_config": None}
|
||||
return dict(list(base_config.items()) + list(config.items()))
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config):
|
||||
config = config.copy()
|
||||
|
||||
# BaseQuantizeWrapper may be constructed with any QuantizeConfig and the
|
||||
# wrapper itself cannot know all the possible config classes.
|
||||
# The deserialization code should ensure the QuantizeConfig is in keras
|
||||
# serialization scope.
|
||||
quantize_config = deserialize_keras_object(
|
||||
config.pop("quantize_config"), module_objects=globals(), custom_objects=None
|
||||
)
|
||||
|
||||
layer = tf.keras.layers.deserialize(config.pop("layer"))
|
||||
|
||||
return cls(layer=layer, quantize_config=quantize_config, **config)
|
||||
|
||||
@property
|
||||
def trainable(self):
|
||||
return self.layer.trainable
|
||||
|
||||
@trainable.setter
|
||||
def trainable(self, value):
|
||||
self.layer.trainable = value
|
||||
|
||||
@property
|
||||
def trainable_weights(self):
|
||||
return self.layer.trainable_weights + self._trainable_weights
|
||||
|
||||
@property
|
||||
def non_trainable_weights(self):
|
||||
return self.layer.non_trainable_weights + self._non_trainable_weights
|
||||
|
||||
@property
|
||||
def updates(self):
|
||||
return self.layer.updates + self._updates
|
||||
|
||||
@property
|
||||
def losses(self):
|
||||
return self.layer.losses + self._losses
|
||||
@@ -0,0 +1,556 @@
|
||||
#
|
||||
# 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 tensorflow.python.util import tf_inspect
|
||||
from tensorflow_quantization.quantize_wrapper_base import BaseQuantizeWrapper
|
||||
import warnings
|
||||
|
||||
"""
|
||||
Naming convention for keras `layer` quantize wrapper is
|
||||
<layer.__class__.__name__>QuantizeWrapper
|
||||
"""
|
||||
|
||||
DISABLED_LAYER_QUANTIZATION_DEFAULT = [
|
||||
"MaxPooling2D",
|
||||
"BatchNormalization",
|
||||
"Add",
|
||||
"Multiply",
|
||||
"Concatenate"
|
||||
]
|
||||
|
||||
|
||||
# ##############################################
|
||||
# ############# Weighted Layers ################
|
||||
# ##############################################
|
||||
|
||||
|
||||
class WeightedBaseQuantizeWrapper(BaseQuantizeWrapper):
|
||||
"""
|
||||
BaseQuantizeWrapper for weighted layers: Conv2D, DepthwiseConv2D, and Dense layer.
|
||||
These layers share a lot of the same code except for a few modifications. Conv2D and Dense share the same code.
|
||||
Layers that inherit this class support weight and input QDQ nodes.
|
||||
|
||||
TRT Rule:
|
||||
One Q/DQ pair should be attached to the input activation, and another Q/DQ pair should be attached to weights.
|
||||
Weights tensor is per-channel quantized:
|
||||
For the Q/DQ attached to weight tensor, set axis=0 and axis=1 for Conv and ConvTransposed respectively.
|
||||
Input tensor is per-tensor quantized.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, layer: tf.keras.layers.Layer, kernel_type: str = "kernel", **kwargs
|
||||
):
|
||||
"""
|
||||
Creates a wrapper to emulate quantization for a keras layer.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
kernel_type (str): Options=['kernel' for Conv2D/Dense, 'depthwise_kernel' for DepthwiseConv2D]
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
self.kernel_type = kernel_type
|
||||
self.channel_axis = kwargs.get("axis", -1)
|
||||
super().__init__(layer, **kwargs)
|
||||
|
||||
def build(self, input_shape):
|
||||
super().build(input_shape)
|
||||
|
||||
self._weight_vars = []
|
||||
self.input_vars = {}
|
||||
self.output_vars = {}
|
||||
self.channel_axis = -1
|
||||
if self.kernel_type == "depthwise_kernel":
|
||||
self.channel_axis = 2
|
||||
# quantize weights only applicable for weighted ops.
|
||||
# By default weights is per channel quantization
|
||||
if self.quantize_weights:
|
||||
# get kernel weights dims.
|
||||
kernel_weights = getattr(self.layer, self.kernel_type)
|
||||
min_weight = self.layer.add_weight(
|
||||
kernel_weights.name.split(":")[0] + "_min",
|
||||
shape=(kernel_weights.shape[self.channel_axis]),
|
||||
initializer=tf.keras.initializers.Constant(-6.0),
|
||||
trainable=False,
|
||||
)
|
||||
max_weight = self.layer.add_weight(
|
||||
kernel_weights.name.split(":")[0] + "_max",
|
||||
shape=(kernel_weights.shape[self.channel_axis]),
|
||||
initializer=tf.keras.initializers.Constant(6.0),
|
||||
trainable=False,
|
||||
)
|
||||
quantizer_vars = {"min_var": min_weight, "max_var": max_weight}
|
||||
self._weight_vars.append((kernel_weights, quantizer_vars))
|
||||
# Needed to ensure unquantized weights get trained as part of the wrapper.
|
||||
self._trainable_weights.append(kernel_weights)
|
||||
|
||||
# By default input is per tensor quantization
|
||||
if self.quantize_inputs:
|
||||
input_min_weight = self.layer.add_weight(
|
||||
self.layer.name + "_ip_min",
|
||||
shape=None,
|
||||
initializer=tf.keras.initializers.Constant(-6.0),
|
||||
trainable=False,
|
||||
)
|
||||
input_max_weight = self.layer.add_weight(
|
||||
self.layer.name + "_ip_max",
|
||||
shape=None,
|
||||
initializer=tf.keras.initializers.Constant(6.0),
|
||||
trainable=False,
|
||||
)
|
||||
self.input_vars["min_var"] = input_min_weight
|
||||
self.input_vars["max_var"] = input_max_weight
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
if training is None:
|
||||
training = tf.keras.backend.learning_phase()
|
||||
|
||||
# Quantize all weights, and replace them in the underlying layer.
|
||||
if self.quantize_weights:
|
||||
quantized_weights = []
|
||||
quantized_weight = self._last_value_quantizer(
|
||||
self._weight_vars[0][0],
|
||||
training,
|
||||
self._weight_vars[0][1],
|
||||
per_channel=True,
|
||||
channel_axis=self.channel_axis,
|
||||
)
|
||||
quantized_weights.append(quantized_weight)
|
||||
# Replace the original weights with QDQ weights
|
||||
setattr(self.layer, self.kernel_type, quantized_weights[0])
|
||||
|
||||
# Quantize inputs to the conv layer
|
||||
if self.quantize_inputs:
|
||||
quantized_inputs = self._last_value_quantizer(
|
||||
inputs, training, self.input_vars, per_channel=False
|
||||
)
|
||||
else:
|
||||
quantized_inputs = inputs
|
||||
|
||||
args = tf_inspect.getfullargspec(self.layer.call).args
|
||||
if "training" in args:
|
||||
outputs = self.layer.call(quantized_inputs, training=training)
|
||||
else:
|
||||
outputs = self.layer.call(quantized_inputs)
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
class Conv2DQuantizeWrapper(WeightedBaseQuantizeWrapper):
|
||||
def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
|
||||
"""
|
||||
Creates a wrapper to emulate quantization for the Conv2D keras layer.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
self.kernel_type = "kernel"
|
||||
super().__init__(layer, kernel_type=self.kernel_type, **kwargs)
|
||||
|
||||
def build(self, input_shape):
|
||||
super().build(input_shape)
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
return super().call(inputs, training=training)
|
||||
|
||||
|
||||
class DenseQuantizeWrapper(WeightedBaseQuantizeWrapper):
|
||||
def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
|
||||
"""
|
||||
Creates a wrapper to emulate quantization for the Dense keras layer.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
self.kernel_type = "kernel"
|
||||
super().__init__(layer, kernel_type=self.kernel_type, **kwargs)
|
||||
|
||||
def build(self, input_shape):
|
||||
super().build(input_shape)
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
return super().call(inputs, training=training)
|
||||
|
||||
|
||||
class DepthwiseConv2DQuantizeWrapper(WeightedBaseQuantizeWrapper):
|
||||
"""Requires TF >= 2.8.0"""
|
||||
|
||||
def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
|
||||
"""
|
||||
Creates a wrapper to emulate quantization for the DepthwiseConv2D keras layer.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
self.kernel_type = "depthwise_kernel"
|
||||
super().__init__(layer, kernel_type=self.kernel_type, **kwargs)
|
||||
|
||||
def build(self, input_shape):
|
||||
super().build(input_shape)
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
return super().call(inputs, training=training)
|
||||
|
||||
|
||||
# ##############################################
|
||||
# ########### Non-Weighted Layers ##############
|
||||
# ######### with Single Input/Output ###########
|
||||
# ##############################################
|
||||
|
||||
|
||||
class NonWeightedBaseQuantizeWrapper(BaseQuantizeWrapper):
|
||||
"""
|
||||
BaseQuantizeWrapper for non-weighted layers with Single Input/Output: AveragePooling2D, GlobalAveragePooling,
|
||||
MaxPooling2D and BatchNormalization.
|
||||
|
||||
Supports 1 input and 1 output QDQ. Similar to Concat, except that Concat supports multiple inputs.
|
||||
NonWeightedBaseQuantizeWrapper can use WeightedBaseQuantizeWrapper by giving quantize_weigths=False.
|
||||
"""
|
||||
|
||||
def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
|
||||
"""
|
||||
Creates a wrapper to emulate quantization for non-weighted keras layers.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
super().__init__(layer, **kwargs)
|
||||
|
||||
def build(self, input_shape):
|
||||
super().build(input_shape)
|
||||
|
||||
self.input_vars = {}
|
||||
|
||||
# By default input is per tensor quantization
|
||||
if self.quantize_inputs:
|
||||
input_min_weight = self.layer.add_weight(
|
||||
self.layer.name + "_ip_min",
|
||||
shape=None,
|
||||
initializer=tf.keras.initializers.Constant(-6.0),
|
||||
trainable=False,
|
||||
)
|
||||
input_max_weight = self.layer.add_weight(
|
||||
self.layer.name + "_ip_max",
|
||||
shape=None,
|
||||
initializer=tf.keras.initializers.Constant(6.0),
|
||||
trainable=False,
|
||||
)
|
||||
self.input_vars["min_var"] = input_min_weight
|
||||
self.input_vars["max_var"] = input_max_weight
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
if training is None:
|
||||
training = tf.keras.backend.learning_phase()
|
||||
|
||||
# Quantize inputs to the conv layer
|
||||
if self.quantize_inputs:
|
||||
quantized_inputs = self._last_value_quantizer(
|
||||
inputs, training, self.input_vars, per_channel=False
|
||||
)
|
||||
else:
|
||||
quantized_inputs = inputs
|
||||
|
||||
args = tf_inspect.getfullargspec(self.layer.call).args
|
||||
if "training" in args:
|
||||
outputs = self.layer.call(quantized_inputs, training=training)
|
||||
else:
|
||||
outputs = self.layer.call(quantized_inputs)
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
class AveragePooling2DQuantizeWrapper(NonWeightedBaseQuantizeWrapper):
|
||||
"""
|
||||
TRT Rule:
|
||||
Add Q/DQ to its input if the ops follows is quantized.
|
||||
Quantize average pooling will introduce small variance compared to float because of the rounding change.
|
||||
TensorRT doesn’t have Int8 in and fp32 out average pool support.
|
||||
If the op follows average pooling is not quantized, it is users choice between running average pooling
|
||||
in int8 then convert to fp32 for the following op and run average pooling in fp32.
|
||||
|
||||
Currently, we're adding QDQ to all AveragePooling2D layers.
|
||||
"""
|
||||
|
||||
def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
|
||||
"""
|
||||
Creates a wrapper to emulate quantization for the AveragePooling2D keras layer.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
super().__init__(layer, **kwargs)
|
||||
|
||||
def build(self, input_shape):
|
||||
super().build(input_shape)
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
return super().call(inputs, training=training)
|
||||
|
||||
|
||||
class GlobalAveragePooling2DQuantizeWrapper(NonWeightedBaseQuantizeWrapper):
|
||||
"""
|
||||
TRT Rule:
|
||||
No explicit rule from the TRT team. Following the same as AveragePooling2D.
|
||||
|
||||
Residual block v2: Add to MaxPool (branch1) and BN (branch2).
|
||||
|
||||
Supports 1 input and 1 output QDQ. Same as AveragePooling2DQuantizeWrapper.
|
||||
"""
|
||||
|
||||
def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
|
||||
"""
|
||||
Creates a wrapper to emulate quantization for the GlobalAveragePooling2D keras layer.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
super().__init__(layer, **kwargs)
|
||||
|
||||
def build(self, input_shape):
|
||||
super().build(input_shape)
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
return super().call(inputs, training=training)
|
||||
|
||||
|
||||
class MaxPooling2DQuantizeWrapper(NonWeightedBaseQuantizeWrapper):
|
||||
"""
|
||||
TRT Rule:
|
||||
Max pooling is precision-neutral. But unlike ReLU, input and output of max pooling will have different
|
||||
histograms which will lead to different calibration results.
|
||||
The recommendation is to let TensorRT optimize precision neutral ops.
|
||||
There are cases where adding Q/DQ before maxpool can enable additional optimization.
|
||||
|
||||
Residual block v2: Add to MaxPool (branch1) and BN (branch2).
|
||||
|
||||
Supports 1 input and 1 output QDQ. Same as AveragePooling2DQuantizeWrapper.
|
||||
"""
|
||||
|
||||
def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
|
||||
"""
|
||||
Creates a wrapper to emulate quantization for the MaxPooling2D keras layer.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
super().__init__(layer, **kwargs)
|
||||
|
||||
def build(self, input_shape):
|
||||
super().build(input_shape)
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
return super().call(inputs, training=training)
|
||||
|
||||
|
||||
class BatchNormalizationQuantizeWrapper(NonWeightedBaseQuantizeWrapper):
|
||||
"""
|
||||
TRT Rule:
|
||||
Keep batch normalization untouched, don't add Q/DQ to its input and not necessary to fold it before exporting
|
||||
graph. TensorRT supports Batch normalization folding. It can take a graph with batch normalization, fold it
|
||||
into previous convolution and create a new graph.
|
||||
If batch normalization is folded before exporting the graph, TensorRT can still import and execute the graph as
|
||||
it becomes regular convolutions.
|
||||
|
||||
Exception for Residual block v2:
|
||||
BN-ReLU-Conv2D -> need to add Q/DQ before BN in order to run in INT8.
|
||||
In order to do that, we add a check in 'quantize_model()' to check if BN's parent is a Conv layer. If it is, set
|
||||
quantize_inputs to False. The reason why we don't add this check here is to allow the user to add QDQ nodes
|
||||
before BN if they so wish.
|
||||
|
||||
Supports 1 input and 1 output QDQ. Same as AveragePooling2DQuantizeWrapper.
|
||||
"""
|
||||
|
||||
def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
|
||||
"""
|
||||
Creates a wrapper to emulate quantization for the BatchNormalization keras layer.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
super().__init__(layer, **kwargs)
|
||||
|
||||
def build(self, input_shape):
|
||||
super().build(input_shape)
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
return super().call(inputs, training=training)
|
||||
|
||||
|
||||
# ##############################################
|
||||
# ########### Non-Weighted Layers ##############
|
||||
# #### with Multiple Inputs, Single Output #####
|
||||
# ##############################################
|
||||
|
||||
|
||||
class NonWeightedBaseQuantizeWrapperForMultipleInputs(BaseQuantizeWrapper):
|
||||
"""
|
||||
BaseQuantizeWrapper for non-weighted layers with Multiple Inputs: Concat, Add, and Multiply.
|
||||
Supports multiple inputs and 1 output QDQ. Similar to AveragePooling2D, except pooling supports only a single input.
|
||||
|
||||
TRT Rule:
|
||||
Add Q/DQ to all inputs of the layer.
|
||||
"""
|
||||
|
||||
def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
|
||||
"""
|
||||
Creates a wrapper to emulate quantization for the keras layer.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
super().__init__(layer, **kwargs)
|
||||
|
||||
def _should_quantization_this_index(self, i):
|
||||
if not self.quantize_specific_input_indices:
|
||||
return True
|
||||
else:
|
||||
# This is a small list so iterating makes sense
|
||||
for e in self.quantize_specific_input_indices:
|
||||
if e == i:
|
||||
return True
|
||||
elif e >= self.num_inputs:
|
||||
warnings.warn(
|
||||
"{layer_name} has {num_inputs} inputs but quantization index {e} is passed.".format(
|
||||
layer_name=self.layer.name, num_inputs=self.num_inputs, e=e
|
||||
)
|
||||
)
|
||||
return False
|
||||
|
||||
def build(self, input_shape):
|
||||
super().build(input_shape)
|
||||
|
||||
self.input_vars = [] # list of dictionaries
|
||||
self.num_inputs = len(input_shape)
|
||||
|
||||
# By default input is per tensor quantization
|
||||
if self.quantize_inputs:
|
||||
# for concat input is list of Tensors
|
||||
layer_name_key_idx = 0
|
||||
for i in range(self.num_inputs):
|
||||
if self._should_quantization_this_index(i):
|
||||
input_min_weight = self.layer.add_weight(
|
||||
self.layer.name + "_ip{}_min".format(layer_name_key_idx),
|
||||
shape=None,
|
||||
initializer=tf.keras.initializers.Constant(-6.0),
|
||||
trainable=False,
|
||||
)
|
||||
input_max_weight = self.layer.add_weight(
|
||||
self.layer.name + "_ip{}_max".format(layer_name_key_idx),
|
||||
shape=None,
|
||||
initializer=tf.keras.initializers.Constant(6.0),
|
||||
trainable=False,
|
||||
)
|
||||
self.input_vars.append(
|
||||
{"min_var": input_min_weight, "max_var": input_max_weight}
|
||||
)
|
||||
layer_name_key_idx += 1
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
if training is None:
|
||||
training = tf.keras.backend.learning_phase()
|
||||
|
||||
# Quantize inputs to the conv layer
|
||||
quantized_inputs = inputs[:]
|
||||
if self.quantize_inputs:
|
||||
input_vars_idx = 0
|
||||
for i in range(len(inputs)):
|
||||
if self._should_quantization_this_index(i):
|
||||
quantized_inputs[i] = self._last_value_quantizer(
|
||||
inputs[i],
|
||||
training,
|
||||
self.input_vars[input_vars_idx],
|
||||
per_channel=False,
|
||||
)
|
||||
input_vars_idx += 1
|
||||
|
||||
args = tf_inspect.getfullargspec(self.layer.call).args
|
||||
if "training" in args:
|
||||
outputs = self.layer.call(quantized_inputs, training=training)
|
||||
else:
|
||||
outputs = self.layer.call(quantized_inputs)
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
class MultiplyQuantizeWrapper(NonWeightedBaseQuantizeWrapperForMultipleInputs):
|
||||
"""
|
||||
TRT Rule:
|
||||
Add Q/DQ to all inputs of Multiply layer in SE block.
|
||||
"""
|
||||
|
||||
def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
|
||||
"""
|
||||
Creates a wrapper to emulate quantization for the Multiply keras layer.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
super().__init__(layer, **kwargs)
|
||||
|
||||
def build(self, input_shape):
|
||||
super().build(input_shape)
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
return super().call(inputs, training=training)
|
||||
|
||||
|
||||
class ConcatenateQuantizeWrapper(NonWeightedBaseQuantizeWrapperForMultipleInputs):
|
||||
"""
|
||||
TRT Rule:
|
||||
Add Q/DQ to all inputs.
|
||||
Alternative: If there is Q/DQ attached to the input of the op after concat, don't add Q/DQ to input to concat,
|
||||
let TensorRT pull from the next op.
|
||||
"""
|
||||
|
||||
def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
|
||||
"""
|
||||
Creates a wrapper to emulate quantization for the Concatenate keras layer.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
super().__init__(layer, **kwargs)
|
||||
|
||||
def build(self, input_shape):
|
||||
super().build(input_shape)
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
return super().call(inputs, training=training)
|
||||
|
||||
|
||||
class AddQuantizeWrapper(NonWeightedBaseQuantizeWrapperForMultipleInputs):
|
||||
"""
|
||||
TRT Rule:
|
||||
If the add is NOT bias. Attach Q/DQ to all of its input.
|
||||
Exception: add in residual block. To trigger fusion, Attach Q/DQ to the residual being added to output of
|
||||
convolution.
|
||||
"""
|
||||
|
||||
def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
|
||||
"""
|
||||
Creates a wrapper to emulate quantization for the Add keras layer.
|
||||
Args:
|
||||
layer (tf.keras.layers.Layer): The keras layer to be quantized.
|
||||
**kwargs: Additional keyword arguments to be passed to the keras layer.
|
||||
"""
|
||||
super().__init__(layer, **kwargs)
|
||||
|
||||
def build(self, input_shape):
|
||||
super().build(input_shape)
|
||||
|
||||
def call(self, inputs, training=None):
|
||||
return super().call(inputs, training=training)
|
||||
@@ -0,0 +1,185 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
# Copyright 2019 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 module is borrowed from TFMOT repository and updated.
|
||||
It implements QDQ insertion based on "Last Value Quantization".
|
||||
"""
|
||||
|
||||
import tensorflow as tf
|
||||
|
||||
|
||||
def LastValueQuantize(
|
||||
inputs,
|
||||
min_var,
|
||||
max_var,
|
||||
per_channel=False,
|
||||
channel_axis=-1,
|
||||
name_prefix="LastValueQuant",
|
||||
is_training=True,
|
||||
num_bits=8,
|
||||
narrow_range=False,
|
||||
symmetric=False,
|
||||
):
|
||||
"""Adds a layer that collects quantization ranges as last input ranges.
|
||||
LastValueQuantize creates variables called 'min' and 'max', representing the
|
||||
interval used for quantization and clamping.
|
||||
Args:
|
||||
inputs: a tensor containing values to be quantized.
|
||||
per_channel: (Optional) a boolean specifying whether to use different
|
||||
quantization ranges per output channel.
|
||||
init_min: a float scalar, the initial value for variable min.
|
||||
init_max: a float scalar, the initial value for variable max.
|
||||
name_prefix: name_prefix for created nodes.
|
||||
is_training: Whether the op is applied to a training or eval graph.
|
||||
num_bits: Number of bits to use for quantization, must be between 2 and 8.
|
||||
narrow_range: Whether to use the narrow quantization range
|
||||
[1; 2^num_bits - 1] or wide range [0; 2^num_bits - 1].
|
||||
symmetric: If true, use symmetric quantization limits instead of training
|
||||
the minimum and maximum of each quantization range separately.
|
||||
Returns:
|
||||
a tensor containing quantized values.
|
||||
"""
|
||||
with tf.name_scope(name_prefix):
|
||||
input_shape = inputs.get_shape()
|
||||
input_dim = len(input_shape)
|
||||
if channel_axis == -1:
|
||||
channel_axis += input_dim
|
||||
|
||||
if not is_training:
|
||||
return _QuantizeAndDequantize(
|
||||
inputs,
|
||||
min_var,
|
||||
max_var,
|
||||
per_channel=per_channel,
|
||||
channel_axis=channel_axis,
|
||||
num_bits=num_bits,
|
||||
narrow_range=narrow_range,
|
||||
)
|
||||
|
||||
if per_channel:
|
||||
if input_dim == 2:
|
||||
reduce_dims = [0]
|
||||
elif input_dim == 4:
|
||||
reduce_dims = [i for i in range(input_dim) if i != channel_axis]
|
||||
|
||||
if per_channel:
|
||||
if input_dim >= 2:
|
||||
batch_min = tf.math.reduce_min(
|
||||
inputs, axis=reduce_dims, name="BatchMin"
|
||||
)
|
||||
else:
|
||||
batch_min = inputs
|
||||
else:
|
||||
batch_min = tf.math.reduce_min(inputs, name="BatchMin")
|
||||
|
||||
if per_channel:
|
||||
if input_dim >= 2:
|
||||
batch_max = tf.math.reduce_max(
|
||||
inputs, axis=reduce_dims, name="BatchMax"
|
||||
)
|
||||
else:
|
||||
batch_max = inputs
|
||||
else:
|
||||
batch_max = tf.math.reduce_max(inputs, name="BatchMax")
|
||||
|
||||
if symmetric:
|
||||
if narrow_range:
|
||||
min_max_ratio = -1
|
||||
else:
|
||||
# In two's complement notation, the negative range is slightly larger
|
||||
# than the positive range.
|
||||
min_max_ratio = -((1 << num_bits) - 2) / (1 << num_bits)
|
||||
|
||||
# TFLite requires that 0.0 if always in the [min; max] range. Because
|
||||
# batch_min <= batch_max, it follows that range_min <= 0 <= range_max.
|
||||
range_min = tf.math.minimum(batch_min, batch_max / min_max_ratio)
|
||||
range_max = tf.math.maximum(batch_max, batch_min * min_max_ratio)
|
||||
else:
|
||||
# TFLite requires that 0.0 if always in the [min; max] range.
|
||||
range_min = tf.math.minimum(batch_min, 0.0)
|
||||
range_max = tf.math.maximum(batch_max, 0.0)
|
||||
|
||||
assign_min = min_var.assign(range_min, name="AssignMinLast")
|
||||
assign_max = max_var.assign(range_max, name="AssignMaxLast")
|
||||
|
||||
return _QuantizeAndDequantize(
|
||||
inputs,
|
||||
assign_min,
|
||||
assign_max,
|
||||
per_channel=per_channel,
|
||||
channel_axis=channel_axis,
|
||||
num_bits=num_bits,
|
||||
narrow_range=narrow_range,
|
||||
)
|
||||
|
||||
|
||||
def _QuantizeAndDequantize(
|
||||
inputs, min_var, max_var, per_channel, channel_axis, num_bits, narrow_range
|
||||
):
|
||||
"""Adds a fake quantization operation.
|
||||
Depending on value of per_channel, this operation may do global quantization
|
||||
or per channel quantization. min_var and max_var should have corresponding
|
||||
shapes: [1] when per_channel == False and [d] when per_channel == True.
|
||||
Args:
|
||||
inputs: a tensor containing values to be quantized.
|
||||
min_var: a variable containing quantization range lower end(s).
|
||||
max_var: a variable containing quantization range upper end(s).
|
||||
per_channel: a boolean specifying whether to use per-channel quantization.
|
||||
num_bits: Number of bits to use for quantization, must be between 2 and 8.
|
||||
narrow_range: Whether to use the narrow quantization range
|
||||
[1; 2^num_bits - 1] or wide range [0; 2^num_bits - 1].
|
||||
Returns:
|
||||
a tensor containing quantized values.
|
||||
"""
|
||||
|
||||
if per_channel:
|
||||
|
||||
return tf.quantization.quantize_and_dequantize_v2(
|
||||
inputs,
|
||||
min_var,
|
||||
max_var,
|
||||
num_bits=num_bits,
|
||||
narrow_range=narrow_range,
|
||||
axis=channel_axis,
|
||||
range_given=True,
|
||||
)
|
||||
else:
|
||||
assert min_var.get_shape() == [] # pylint: disable=g-explicit-bool-comparison
|
||||
assert max_var.get_shape() == [] # pylint: disable=g-explicit-bool-comparison
|
||||
|
||||
return tf.quantization.quantize_and_dequantize_v2(
|
||||
inputs,
|
||||
min_var,
|
||||
max_var,
|
||||
num_bits=num_bits,
|
||||
narrow_range=narrow_range,
|
||||
range_given=True,
|
||||
)
|
||||
@@ -0,0 +1,364 @@
|
||||
#
|
||||
# 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"]
|
||||
)
|
||||
Reference in New Issue
Block a user