wehub/nvidia--tensorrt
1
0
Fork 0
Code Issues Pull Requests Actions 25 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/tensorflow-quantization/tests/custom_qdq_cases_test.py
+216
View File
@@ -0,0 +1,216 @@
#
# 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 network_pool import pippin_28_28
from tensorflow_quantization import custom_qdq_cases
import pytest
from tensorflow_quantization import quantize_model
from tensorflow_quantization.utils import convert_saved_model_to_onnx
from tensorflow_quantization.utils import CreateAssetsFolders
import tensorflow as tf
test_assets = CreateAssetsFolders("test_custom_qdq_cases")
def test_resnet_residual_qdq_case():
model = pippin_28_28()
test_assets.add_folder("pipin_28_28")
tf.keras.models.save_model(model, test_assets.pipin_28_28.fp32_saved_model)
convert_saved_model_to_onnx(
saved_model_dir=test_assets.pipin_28_28.fp32_saved_model,
onnx_model_path=test_assets.pipin_28_28.fp32_onnx_model,
)
resnet_residual_qdq = custom_qdq_cases.ResNetV1QDQCase()
r = resnet_residual_qdq.case(model, None)
expected_qdq_insertion = {
"add": 1,
"add_1": "any",
}
assert (
len(r.layers) == 2
), "There should be 2 custom layers, but found {}".format(len(r.layers))
for l in r.layers:
if l.name not in expected_qdq_insertion:
raise Exception(
"Layer {} is not expected to be treated as custom layer".format(l.name)
)
else:
if l.quantization_index != None:
if expected_qdq_insertion[l.name] == "any":
continue
assert (
l.quantization_index[0] == expected_qdq_insertion[l.name]
), "For layer {l_name}, only {expected_qdq} indices should be quantized".format(
l_name=l.name, expected_qdq=expected_qdq_insertion[l.name]
)
def assert_add_bn_expected_layers(
r, expected_add_layer_behavior, expected_bn_layer_behavior, expected_mp_layer_behavior
):
assert len(r.layers) == (
len(expected_add_layer_behavior) + len(expected_bn_layer_behavior) + len(expected_mp_layer_behavior)
), "Not all expected layers are captured for ResNet custom QDQ case."
for l in r.layers:
assert (
l.name in expected_add_layer_behavior
or l.name in expected_bn_layer_behavior
or l.name in expected_mp_layer_behavior
), "layer {} is not expected to be captured for ResNet custom QDQ case".format(
l.name
)
if "add" in l.name:
if expected_add_layer_behavior[l.name] == "any":
continue
assert l.quantization_index[0] == expected_add_layer_behavior[l.name], (
"For layer {l_name}, expected quantization index is {expected_add_behavior} but index {l_quant_id} "
"is captured in ResNet custom QDQ case.".format(
l_name=l.name,
expected_add_behavior=expected_add_layer_behavior[l.name],
l_quant_idx=l.quantization_index[0],
)
)
def test_resnet50_residual_qdq_case():
resnet50 = tf.keras.applications.resnet50.ResNet50(weights=None)
test_assets.add_folder("resnet50_v1")
tf.keras.models.save_model(resnet50, test_assets.resnet50_v1.fp32_saved_model)
convert_saved_model_to_onnx(
saved_model_dir=test_assets.resnet50_v1.fp32_saved_model,
onnx_model_path=test_assets.resnet50_v1.fp32_onnx_model,
)
resnet_custom_qdq_case = custom_qdq_cases.ResNetV1QDQCase()
r = resnet_custom_qdq_case.case(resnet50, None)
for r_1 in r.layers:
print("\"{}\",".format(r_1.name))
expected_add_layer_behavior = {
"conv2_block1_add": "any",
"conv2_block2_add": 0,
"conv2_block3_add": 0,
"conv3_block1_add": "any",
"conv3_block2_add": 0,
"conv3_block3_add": 0,
"conv3_block4_add": 0,
"conv4_block1_add": "any",
"conv4_block2_add": 0,
"conv4_block3_add": 0,
"conv4_block4_add": 0,
"conv4_block5_add": 0,
"conv4_block6_add": 0,
"conv5_block1_add": "any",
"conv5_block2_add": 0,
"conv5_block3_add": 0,
}
# Empty, no BatchNorm layers should be quantized in ResNet-v1
expected_bn_layer_behavior = {}
# MaxPool quantization is actually not needed in ResNet-v1
expected_mp_layer_behavior = {}
assert_add_bn_expected_layers(
r, expected_add_layer_behavior, expected_bn_layer_behavior, expected_mp_layer_behavior
)
q_resnet50 = quantize_model(resnet50, custom_qdq_cases=[resnet_custom_qdq_case])
tf.keras.models.save_model(q_resnet50, test_assets.resnet50_v1.int8_saved_model)
convert_saved_model_to_onnx(
saved_model_dir=test_assets.resnet50_v1.int8_saved_model,
onnx_model_path=test_assets.resnet50_v1.int8_onnx_model,
)
def test_resnet50v2_bn_qdq_case():
resnet50_v2 = tf.keras.applications.resnet_v2.ResNet50V2(weights=None)
test_assets.add_folder("resnet50_v2")
tf.keras.models.save_model(resnet50_v2, test_assets.resnet50_v2.fp32_saved_model)
convert_saved_model_to_onnx(
saved_model_dir=test_assets.resnet50_v2.fp32_saved_model,
onnx_model_path=test_assets.resnet50_v2.fp32_onnx_model,
)
resnet_custom_qdq_case = custom_qdq_cases.ResNetV2QDQCase()
r = resnet_custom_qdq_case.case(resnet50_v2, None)
for r_1 in r.layers:
print("\"{}\",".format(r_1.name))
expected_add_layer_behavior = {
"conv2_block1_out": 0,
"conv2_block2_out": 0,
"conv2_block3_out": 0,
"conv3_block1_out": 0,
"conv3_block2_out": 0,
"conv3_block3_out": 0,
"conv3_block4_out": 0,
"conv4_block1_out": 0,
"conv4_block2_out": 0,
"conv4_block3_out": 0,
"conv4_block4_out": 0,
"conv4_block5_out": 0,
"conv4_block6_out": 0,
"conv5_block1_out": 0,
"conv5_block2_out": 0,
"conv5_block3_out": 0,
}
# ResNet-v2 quantizes BatchNorms that are not connected to Conv layers
expected_bn_layer_behavior = {
"conv2_block1_preact_bn",
"conv2_block2_preact_bn",
"conv2_block3_preact_bn",
"conv3_block1_preact_bn",
"conv3_block2_preact_bn",
"conv3_block3_preact_bn",
"conv3_block4_preact_bn",
"conv4_block1_preact_bn",
"conv4_block2_preact_bn",
"conv4_block3_preact_bn",
"conv4_block4_preact_bn",
"conv4_block5_preact_bn",
"conv4_block6_preact_bn",
"conv5_block1_preact_bn",
"conv5_block2_preact_bn",
"conv5_block3_preact_bn",
"post_bn",
}
# ResNet-v2 quantizes all MaxPool layers
expected_mp_layer_behavior = {
"pool1_pool",
"max_pooling2d",
"max_pooling2d_1",
"max_pooling2d_2",
}
assert_add_bn_expected_layers(
r, expected_add_layer_behavior, expected_bn_layer_behavior, expected_mp_layer_behavior
)
q_resnet50_v2 = quantize_model(
resnet50_v2, custom_qdq_cases=[resnet_custom_qdq_case]
)
tf.keras.models.save_model(q_resnet50_v2, test_assets.resnet50_v2.int8_saved_model)
convert_saved_model_to_onnx(
saved_model_dir=test_assets.resnet50_v2.int8_saved_model,
onnx_model_path=test_assets.resnet50_v2.int8_onnx_model,
)
tools/tensorflow-quantization/tests/network_pool.py
+277
View File
@@ -0,0 +1,277 @@
#
# 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 contains tiny networks used for testing across different modules.
They are named after famous Hobbits for obvious reasons.
"""
import tensorflow as tf
##################################################
###### Tiny, VGG like network ####################
##################################################
def bilbo_28_28():
"""
Network with VGG like architecture.
"""
input_img = tf.keras.layers.Input(shape=(28, 28), name="nn_input")
x = tf.keras.layers.Reshape(target_shape=(28, 28, 1), name="reshape_0")(input_img)
x = tf.keras.layers.Conv2D(filters=516, kernel_size=(3, 3), name="conv_0")(x)
x = tf.keras.layers.ReLU(name="relu_0")(x)
x = tf.keras.layers.Conv2D(filters=252, kernel_size=(3, 3), name="conv_1")(x)
x = tf.keras.layers.ReLU(name="relu_1")(x)
x = tf.keras.layers.Conv2D(filters=126, kernel_size=(3, 3), name="conv_2")(x)
x = tf.keras.layers.ReLU(name="relu_2")(x)
x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3), name="conv_3")(x)
x = tf.keras.layers.ReLU(name="relu_3")(x)
x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), name="conv_4")(x)
x = tf.keras.layers.ReLU(name="relu_4")(x)
x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3, 3), name="conv_5")(x)
x = tf.keras.layers.ReLU(name="relu_5")(x)
x = tf.keras.layers.Conv2D(filters=8, kernel_size=(3, 3), name="conv_6")(x)
x = tf.keras.layers.ReLU(name="relu_6")(x)
x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2), name="max_pool_0")(x)
x = tf.keras.layers.Flatten(name="flatten_0")(x)
x = tf.keras.layers.Dense(100, name="dense_0")(x)
x = tf.keras.layers.ReLU(name="relu_7")(x)
x = tf.keras.layers.Dense(10, name="dense_1")(x)
return tf.keras.Model(input_img, x, name="Bilbo")
#####################################################
###### Tiny, ResNet like network ####################
#####################################################
def identity_block_plain(input_tensor):
"""
Identity block with no shortcut convolution
"""
y = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), padding="same")(
input_tensor
)
y = tf.keras.layers.ReLU()(y)
y = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3), padding="same")(y)
out = tf.keras.layers.Add()([y, input_tensor])
out = tf.keras.layers.ReLU()(out)
return out
def identity_block_short_conv_plain(input_tensor):
"""
Identity block with shortcut convolution
"""
y = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), padding="same")(
input_tensor
)
y = tf.keras.layers.ReLU()(y)
y = tf.keras.layers.Conv2D(
filters=24, kernel_size=(3, 3), strides=(2, 2), padding="same"
)(y)
ds_input = tf.keras.layers.Conv2D(
filters=24, kernel_size=(3, 3), strides=(2, 2), padding="same"
)(input_tensor)
out = tf.keras.layers.Add()([y, ds_input])
out = tf.keras.layers.ReLU()(out)
return out
def frodo_32_32():
"""
Dummy network with resnet like architecture.
"""
input_img = tf.keras.layers.Input(shape=(32, 32, 3))
x = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3))(input_img)
x = tf.keras.layers.ReLU()(x)
x = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3))(x)
x = tf.keras.layers.ReLU()(x)
x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
x = identity_block_plain(x)
x = identity_block_short_conv_plain(x)
x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
x = tf.keras.layers.Flatten()(x)
x = tf.keras.layers.Dense(100)(x)
x = tf.keras.layers.ReLU()(x)
x = tf.keras.layers.Dense(10)(x)
return tf.keras.Model(input_img, x, name="Frodo")
def sam_32_32():
"""
Dummy network with resnet like architecture.
"""
input_img = tf.keras.layers.Input(shape=(32, 32, 3))
x = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3))(input_img)
x = tf.keras.layers.ReLU()(x)
x = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3))(x)
x = tf.keras.layers.ReLU()(x)
x = identity_block_plain(x)
x = identity_block_plain(x)
x = identity_block_short_conv_plain(x)
x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
x = tf.keras.layers.Flatten()(x)
x = tf.keras.layers.Dense(100)(x)
x = tf.keras.layers.ReLU()(x)
x = tf.keras.layers.Dense(10)(x)
return tf.keras.Model(input_img, x, name="Sam")
##############################################
###### Popular network blocks ################
##############################################
def relu_bn(input):
"""
Block with BN+ReLU
"""
bn = tf.keras.layers.BatchNormalization()(input)
relu = tf.keras.layers.ReLU()(bn)
return relu
def bn(input):
return tf.keras.layers.BatchNormalization()(input)
def relu(input):
return tf.keras.layers.ReLU()(input)
def inception_block(input_tensor):
"""
Inception block from GoogleNet
"""
b1x1 = tf.keras.layers.Conv2D(filters=12, kernel_size=(1, 1), padding="same")(
input_tensor
)
b1x1 = relu_bn(b1x1)
b5x5 = tf.keras.layers.Conv2D(filters=12, kernel_size=(1, 1), padding="same")(
input_tensor
)
b5x5 = tf.keras.layers.Conv2D(filters=24, kernel_size=(5, 5), padding="same")(b5x5)
b5x5 = relu_bn(b5x5)
b3x3 = tf.keras.layers.Conv2D(filters=12, kernel_size=(1, 1), padding="same")(
input_tensor
)
b3x3 = tf.keras.layers.Conv2D(filters=20, kernel_size=(3, 3), padding="same")(b3x3)
b3x3 = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3), padding="same")(b3x3)
b3x3 = relu_bn(b3x3)
out = tf.keras.layers.Concatenate()([b1x1, b5x5])
return out
def identity_block_bn(input_tensor):
"""
Identity block with no shortcut convolution
"""
y = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), padding="same")(
input_tensor
)
y = relu_bn(y)
y = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3), padding="same")(y)
y = bn(y)
out = tf.keras.layers.Add()([y, input_tensor])
out = relu(out)
return out
def identity_block_short_conv_bn(input_tensor):
"""
Identity block with shortcut convolution
"""
y = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), padding="same")(
input_tensor
)
y = relu_bn(y)
y = tf.keras.layers.Conv2D(
filters=24, kernel_size=(3, 3), strides=(2, 2), padding="same"
)(y)
y = bn(y)
ds_input = tf.keras.layers.Conv2D(
filters=24, kernel_size=(3, 3), strides=(2, 2), padding="same"
)(input_tensor)
ds_input = bn(ds_input)
out = tf.keras.layers.Add()([y, ds_input])
out = relu(out)
return out
def otho_28_28():
input_img = tf.keras.layers.Input(shape=(28, 28), name="input_0")
r = tf.keras.layers.Reshape(target_shape=(28, 28, 1), name="reshape_0")(input_img)
x = tf.keras.layers.Conv2D(filters=2, kernel_size=(3, 3), name="conv_0")(r)
x = tf.keras.layers.ReLU(name="relu_0")(x)
x = tf.keras.layers.Conv2D(filters=2, kernel_size=(3, 3), name="conv_1")(x)
x = tf.keras.layers.ReLU(name="relu_1")(x)
x = tf.keras.layers.Flatten(name="flatten_0")(x)
return tf.keras.Model(input_img, x, name="Otho")
def lotho_28_28():
input_img = tf.keras.layers.Input(shape=(28, 28), name="input_0")
r = tf.keras.layers.Reshape(target_shape=(28, 28, 1), name="reshape_0")(input_img)
x = tf.keras.layers.DepthwiseConv2D(kernel_size=(3, 3), name="dconv_0")(r)
x = tf.keras.layers.ReLU(name="relu_0")(x)
x = tf.keras.layers.DepthwiseConv2D(kernel_size=(3, 3), name="dconv_1")(x)
x = tf.keras.layers.ReLU(name="relu_1")(x)
x = tf.keras.layers.Flatten(name="flatten_0")(x)
return tf.keras.Model(input_img, x, name="Lotho")
def lobelia_28_28():
input_img = tf.keras.layers.Input(shape=(28, 28), name="input_0")
r = tf.keras.layers.Reshape(target_shape=(28, 28, 1), name="reshape_0")(input_img)
x = tf.keras.layers.Flatten(name="flatten_0")(r)
x = tf.keras.layers.Dense(100, name="dense_0")(x)
x = tf.keras.layers.ReLU(name="relu_0")(x)
x = tf.keras.layers.Dense(10, name="dense_1")(x)
return tf.keras.Model(input_img, x, name="Lobelia")
def merry_28_28():
input_img = tf.keras.layers.Input(shape=(28, 28))
x = tf.keras.layers.Reshape(target_shape=(28, 28, 1))(input_img)
x = tf.keras.layers.Conv2D(filters=8, kernel_size=(3, 3))(x)
x = relu_bn(x)
x = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3))(x)
x = relu_bn(x)
x = inception_block(x)
x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
x = tf.keras.layers.Flatten()(x)
x = tf.keras.layers.Dense(100)(x)
x = tf.keras.layers.ReLU()(x)
x = tf.keras.layers.Dense(10)(x)
return tf.keras.Model(input_img, x, name="Merry")
def pippin_28_28():
input_img = tf.keras.layers.Input(shape=(28, 28))
x = tf.keras.layers.Reshape(target_shape=(28, 28, 1))(input_img)
x = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3))(x)
x = relu_bn(x)
x = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3))(x)
x = relu_bn(x)
x = identity_block_bn(x)
x = identity_block_short_conv_bn(x)
x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
x = tf.keras.layers.Flatten()(x)
x = tf.keras.layers.Dense(100)(x)
x = tf.keras.layers.ReLU()(x)
x = tf.keras.layers.Dense(10)(x)
return tf.keras.Model(input_img, x, name="Pippin")
tools/tensorflow-quantization/tests/onnx_graph_qdq_validator.py
+551
View File
@@ -0,0 +1,551 @@
#
# 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 onnx
import onnx_graphsurgeon as gs
import tensorflow as tf
from tensorflow_quantization.quantize import LayerConfig, quantize_model
from typing import List, Tuple
from tensorflow_quantization.utils import convert_saved_model_to_onnx
import copy
EXPECTED_QDQ_INSERTION = [
LayerConfig(name="Conv2D", is_keras_class=True),
LayerConfig(name="Dense", is_keras_class=True),
LayerConfig(name="DepthwiseConv2D", is_keras_class=True),
LayerConfig(
name="Concatenate",
is_keras_class=True,
quantize_weight=False,
quantization_index=["all"],
),
LayerConfig(
name="AveragePooling2D", is_keras_class=True, quantize_weight=False
),
LayerConfig(
name="GlobalAveragePooling2D", is_keras_class=True, quantize_weight=False
)
]
class ONNXQDQValidator:
"""
Validate ONNX file for correct QDQ insertion.
All onnx-graphsurgeon terminologies are used in the explanations.
"""
def __init__(self) -> None:
self.expected_qdq_layer_behavior = {}
self.graph = None
self.data_format = tf.keras.backend.image_data_format()
@staticmethod
def _extract_layer_names_from_class_type(
expected_qdq_behavior, original_keras_model
):
"""Checks if expected_qdq_behavior has items where is_keras_class=True and extract all layers relevant to it.
Also checks if the user didn't specifically name that layer in expected_qdq_behavior.
"""
def layer_is_class_type(class_specs, origin_layer):
for c in class_specs:
if origin_layer.__class__.__name__ == c.name:
return c
return None
def skip_layer_name(layer_name):
for layer in expected_qdq_behavior:
if layer_name == layer.name:
return True
return False
expected_qdq_behavior_class = [
layer for layer in expected_qdq_behavior if layer.is_keras_class
]
expected_qdq_behavior_layers = [
layer for layer in expected_qdq_behavior
# Skip if quantize_input and quantize_weight=False
if not layer.is_keras_class and (layer.quantize_input or layer.quantize_weight)
]
if original_keras_model is not None:
for original_layer in original_keras_model.layers:
class_type = layer_is_class_type(
expected_qdq_behavior_class, original_layer
)
if class_type is not None and not skip_layer_name(original_layer.name):
# Skip if quantize_input and quantize_weight=False
if class_type.quantize_input or class_type.quantize_weight:
expected_qdq_behavior_layers.append(
LayerConfig(
name=original_layer.name,
quantize_input=class_type.quantize_input,
quantize_weight=class_type.quantize_weight,
quantization_index=class_type.quantization_index,
)
)
return expected_qdq_behavior_layers
def _collect_layer_names(self, expected_qdq_behavior):
"""
Populates the global variable 'self.expected_qdq_layer_behavior', a dictionary in the format:
key (string) : Layer name after quantization wrapper is applied.
value (list) : List with layer specific parameters.
value[0] (bool) = True if this layer is a keras class.
value[1] (bool) = True if input to this layer should be quantized.
value[2] (bool) = True if layer weight should be quantized.
value[3] (list) = List of quantization index, if any.
value[4] (bool) = Set to False initially but when quantization of this layer is verified, set to True.
"""
for layer in expected_qdq_behavior:
self.expected_qdq_layer_behavior["quant_" + layer.name] = []
self.expected_qdq_layer_behavior["quant_" + layer.name].append(
layer.is_keras_class
)
self.expected_qdq_layer_behavior["quant_" + layer.name].append(
layer.quantize_input
)
self.expected_qdq_layer_behavior["quant_" + layer.name].append(
layer.quantize_weight
)
self.expected_qdq_layer_behavior["quant_" + layer.name].append(
layer.quantization_index if layer.quantization_index is not None else []
)
self.expected_qdq_layer_behavior["quant_" + layer.name].append(False)
def _load_onnx_graph(self, onnx_model_path):
self.graph = gs.import_onnx(onnx.load(onnx_model_path))
def _get_tf_name_of_node(self, onnx_node):
splitted_node_name = onnx_node.name.split("/")
if len(splitted_node_name) > 1:
# This is other node than QuantizeLinear or DequantizeLinear
node_op = onnx_node.op
# Most layers have their name in position -2
# List: Conv, BatchNormalization, Relu, Add, MatMul, Softmax, Pad, MaxPool, GlobalAveragePool
# Exceptions: Squeeze, Transpose, Reshape
if node_op == "Squeeze" or node_op == "Transpose" or node_op == "Reshape":
return splitted_node_name[-1]
# Quantized layers
for exp_qdq_layer_name in self.expected_qdq_layer_behavior.keys():
if exp_qdq_layer_name + "/" in onnx_node.name:
return exp_qdq_layer_name
# Other layers
return splitted_node_name[-2]
else:
return None
def _get_input_tensor_parent(self, onnx_node, input_idx):
"""
Get input Tensors parent recursively.
Here we want to know id DequantizeLinear is tensors parent.
Recursively we go up the graph since reshape layers are added while onnx conversion between QDQ and node.
"""
current_node_ip_tensor = onnx_node.inputs[input_idx]
try:
current_node_ip_tensor_parent = current_node_ip_tensor.inputs[0]
except IndexError:
# Example for weight Tensor, parent is None
return None
while (
current_node_ip_tensor_parent.op == "Transpose"
or current_node_ip_tensor_parent.op == "Reshape"
): # and self.data_format == "channels_last":
# When image data format is 'channels_last' or Conv is of type 'Depthwise', Transpose and/or Reshape
# layers are added between QDQ and target layer. Always select input at index 0 since it's the
# variable coming from the previous node. Other indices, if present, are constant inputs to the node.
current_node_ip_tensor = current_node_ip_tensor_parent.inputs[0]
try:
current_node_ip_tensor_parent = current_node_ip_tensor.inputs[0]
except IndexError:
# We can't move upwards anymore in the graph
break
return current_node_ip_tensor_parent.op
def _weighted_qdq_behavior(self, node, tf_node_name):
"""
For weighted layers such as MatMul/Conv2D and DepthwiseConv2D, only quantize_input and quantize_weight options
are valid.
Input has length of 2 usually.
Index 0 is Variable i.e. output of previous op
Index 1 is Constant i.e. weight
NOTE: In general, for node with more than one inputs, index 0 is variable coming out of previous node.
Other indices are constant inputs to the node.
"""
# case 1. Only one of quantize_input or quantize_weight is True
if (
not self.expected_qdq_layer_behavior[tf_node_name][1]
or not self.expected_qdq_layer_behavior[tf_node_name][2]
):
# subcase 1. When quantize_input=False
if not self.expected_qdq_layer_behavior[tf_node_name][1]:
if self._get_input_tensor_parent(node, 0) == "DequantizeLinear":
print(
"[E] quantize_input=False but still input is quantized for weighted layer `{}`".format(
tf_node_name
)
)
return False
# send update that correct quantization is found for intended layer.
self.expected_qdq_layer_behavior[tf_node_name][-1] = True
# subcase 2. When quantize_weight=False
if not self.expected_qdq_layer_behavior[tf_node_name][2]:
if self._get_input_tensor_parent(node, 1) == "DequantizeLinear":
print(
"[E] quantize_weight=False but still weight is quantized for weighted layer `{}`".format(
tf_node_name
)
)
return False
# send update that correct quantization is found for intended layer.
self.expected_qdq_layer_behavior[tf_node_name][-1] = True
else:
# case 2. Both quantize_input=True, quantize_weight=True
# Every input should be output of DequantizeLinear op
parent_check = []
for idx in range(len(node.inputs)):
input_tensor_parent = self._get_input_tensor_parent(node, idx)
if input_tensor_parent != "DequantizeLinear":
parent_check.append(0)
else:
parent_check.append(1)
# Check if both input and weight are quantized (2 inputs == 'DequantizeLinear')
# This takes into consideration that Conv sometimes has a BiasAdd input, which is not quantized.
if sum(parent_check) < 2:
print(
"[E] quantize_weight=True and quantize_input=True but still not all inputs are quantized for "
"weighted layer `{}`".format(tf_node_name)
)
return False
# send update that correct quantization is found for intended layer.
self.expected_qdq_layer_behavior[tf_node_name][-1] = True
return True
def _pool_qdq_behavior(self, node, tf_node_name):
"""
Pool layer has just one input which is variable coming from the previous op.
"""
if self._get_input_tensor_parent(node, 0) != "DequantizeLinear":
print(
"[E] Variable input for MaxPool layer `{}` is not quantized.".format(
tf_node_name
)
)
return False
# send update that correct quantization is found for intended layer.
self.expected_qdq_layer_behavior[tf_node_name][-1] = True
return True
def _bn_qdq_behavior(self, node, tf_node_name):
"""
BN has one variable input and four (scale, beta, mean, var) constant inputs.
Remember variable input is always at index 0
For quantization, just check QDQ nodes insertion in variable input.
"""
# Check if the parent node is not DequantizeLinear and input should be quantized.
# Reason: in the ResNet CustomQDQCase, BN is only quantized when preceded by Conv. Otherwise, quantize_input
# (and quantize_weight) is set to False.
quantize_input = self.expected_qdq_layer_behavior[tf_node_name][1]
if (
self._get_input_tensor_parent(node, 0) != "DequantizeLinear"
and quantize_input
):
print(
"[E] Variable input for BatchNormalization layer `{}` is not quantized.".format(
tf_node_name
)
)
return False
# send update that correct quantization is found for intended layer.
self.expected_qdq_layer_behavior[tf_node_name][-1] = True
return True
def _multi_input_qdq_behavior(self, node, tf_node_name):
"""
For layers with multiple inputs, we need to check whether each intended layer is quantized.
"""
# There is quantization index list, check if provided indices are output of DequantizeLinear
for _, e in enumerate(self.expected_qdq_layer_behavior[tf_node_name][3]):
if e in ["any", "all"]:
all_inputs = len(node.inputs)
q_inputs = 0
for inp_idx in range(all_inputs):
if node.i(inp_idx).op == "DequantizeLinear":
q_inputs += 1
if e == "any" and q_inputs != 1:
print(
"[E] quantization_index=['{}'] thus only one input should be quantized, but {} out of {} "
"inputs are quantized for layer `{}`".format(e, q_inputs, all_inputs, tf_node_name)
)
return False
elif e == "all" and q_inputs != all_inputs:
print(
"[E] quantization_index=['{}'] thus all inputs should be quantized, but {} out of {} "
"inputs are quantized for layer `{}`".format(e, q_inputs, all_inputs, tf_node_name)
)
return False
# send update that correct quantization is found for intended layer.
self.expected_qdq_layer_behavior[tf_node_name][-1] = True
else:
if node.i(e).op != "DequantizeLinear":
print(
"[E] Input at index {e} in layer `{tf_node_name}` should be quantized but it is not.".format(
e=e, tf_node_name=tf_node_name
)
)
return False
# send update that correct quantization is found for intended layer.
self.expected_qdq_layer_behavior[tf_node_name][-1] = True
return True
def _non_quantized_layer_qdq_behavior(self, node, tf_node_name):
"""
Squeeze layer should not be quantized.
"""
if self._get_input_tensor_parent(node, 0) == "DequantizeLinear":
print(
"[E] Variable input for {node_op} layer `{tf_node_name}` is quantized.".format(
node_op=node.op, tf_node_name=tf_node_name
)
)
return False
# send update that correct quantization is found for intended layer.
self.expected_qdq_layer_behavior[tf_node_name][-1] = True
return True
def _qdq_monitor(self, node, tf_node_name):
m = {
"MatMul": self._weighted_qdq_behavior,
"Conv": self._weighted_qdq_behavior,
"MaxPool": self._pool_qdq_behavior,
"AveragePool": self._pool_qdq_behavior,
"GlobalAveragePool": self._pool_qdq_behavior,
"BatchNormalization": self._bn_qdq_behavior,
"Concat": self._multi_input_qdq_behavior,
"Add": self._multi_input_qdq_behavior,
"Mul": self._multi_input_qdq_behavior,
}
if node.op not in m: # Squeeze, Softmax, ...
m[node.op] = self._non_quantized_layer_qdq_behavior
return m[node.op](node, tf_node_name)
def check_onnx_node(self, node_name):
for node in self.graph.nodes:
if node.name == node_name:
print(node)
def _unintended_layer_quantize_check_pass(self):
"""
Check whether un-intended layer is quantized.
If any un-intended layer is quantized, checking fails immediately.
"""
for node in self.graph.nodes:
tf_node_name = self._get_tf_name_of_node(node)
if tf_node_name and "quant" in tf_node_name:
if tf_node_name not in self.expected_qdq_layer_behavior:
print(
"[E] layer `{}` should not be quantized.".format(tf_node_name)
)
return False
return True
def _intended_layer_quantize_check_pass(self):
"""
Checks if the layers exists and whether all expected layers are quantized.
If any intended layer is not quantized, checking fails immediately.
"""
for k, v in self.expected_qdq_layer_behavior.items():
check_quant_layer_exists = any([k + "/" in node.name for node in self.graph.nodes])
check_original_layer_exists = any([k.replace("quant_", "") + "/" in node.name for node in self.graph.nodes])
if not check_quant_layer_exists:
if check_original_layer_exists:
print("[E] layer `{}` should have been quantized but wasn't.".format(k.replace("quant_", "")))
return False
else:
print("[W] layer `{}` does not exist.".format(k))
continue
elif not v[-1]:
print("[E] layer `{}` should be quantized but it did not.".format(k))
return False
return True
def _qdq_insertion_check_pass(self):
"""
Validate QDQ insertion.
"""
check_status = True
for node in self.graph.nodes:
tf_node_name = self._get_tf_name_of_node(node)
if tf_node_name and "quant" in tf_node_name:
check_status = check_status and self._qdq_monitor(node, tf_node_name)
if not check_status:
return check_status
return check_status
def validate(
self, onnx_model_path, expected_qdq_behavior, original_keras_model=None
):
self._load_onnx_graph(onnx_model_path)
expected_qdq_behavior = self._extract_layer_names_from_class_type(
expected_qdq_behavior, original_keras_model
)
# Populate 'self.expected_qdq_layer_behavior'
self._collect_layer_names(expected_qdq_behavior)
ulcp = self._unintended_layer_quantize_check_pass()
if not ulcp:
print("[I] Unintended layer quantization check failed.")
return False
qicp = self._qdq_insertion_check_pass()
if not qicp:
print("[I] Quantize insertion check failed.")
return False
ilcp = self._intended_layer_quantize_check_pass()
if not ilcp:
print("[I] Intended layer quantization check failed.")
return False
return True
def get_expected_qdq_insertion(
nn_model_original: tf.keras.Model,
qspec_test: "QuantizationSpec" = None,
custom_qdq_cases: List["CustomQDQInsertionCase"] = None,
quantization_mode: str = "full",
expected_qdq_insertion_user: List[LayerConfig] = None
) -> List[LayerConfig]:
"""
Gets expected QDQ insertion.
Args:
nn_model_original (tf.keras.Model): baseline model (non-quantized), needed to obtain all layers quantized with
Custom QDQ Case.
qspec_test (QuantizationSpec): Quantization specification to test the quantized model with.
custom_qdq_cases (List[CustomQDQInsertionCase]): indicates layers with custom QDQ placements
(i.e., ResidualConnectionQDQCase).
quantization_mode (str): quantization mode, can be "full" or "partial".
expected_qdq_insertion_user (List[LayerConfig]): List of layer configs specified by the user. If 'None', use
the default quantization behavior.
Returns:
expected_qdq_insertion (List[LayerConfig]): list with expected QDQ node placements.
"""
# 1. Establish QDQ node placement behavior for all relevant classes
if expected_qdq_insertion_user is not None:
# User-specified QDQ behavior
expected_qdq_insertion = expected_qdq_insertion_user
else:
if quantization_mode == "partial":
# No classes are quantized by default
expected_qdq_insertion = []
else:
# Default quantization behavior
expected_qdq_insertion = copy.deepcopy(EXPECTED_QDQ_INSERTION)
# 2. Extend quantization behavior with the user's specifications
if qspec_test is not None:
# Only add layers that are being quantized (don't add when `quantize_input` or 'quantize_weight`=False)
expected_qdq_insertion.extend(qspec_test.layers)
# 3. Extend quantization behavior with the Custom QDQ Cases
if custom_qdq_cases is not None:
for custom_qdq_case in custom_qdq_cases:
qspec_case_object = custom_qdq_case.case(nn_model_original, qspec=qspec_test)
expected_qdq_insertion.extend(qspec_case_object.layers)
# 4. Check if Multiple Input classes have empty or None 'quantization_index'. If so, update it to
# 'quantization_index=["all"]'.
for exp_insertion in expected_qdq_insertion:
if exp_insertion.is_keras_class and exp_insertion.name in ['Add', 'Multiply', 'Concatenate']:
if not exp_insertion.quantization_index: # None or []
exp_insertion.quantization_index = ["all"]
return expected_qdq_insertion
# ###############################################
# ######### Full QAT workflow test ##############
# ###############################################
def validate_quantized_model(
test_assets: "CreateAssetsFolders",
nn_model_original: tf.keras.Model,
quantization_mode: str = "full",
qspec: "QuantizationSpec" = None,
custom_qdq_cases: List["CustomQDQInsertionCase"] = None,
test_name: str = "test",
expected_qdq_insertion: List["LayerConfig"] = None
) -> Tuple[tf.keras.Model, bool]:
"""
Full test workflow: quantization, obtain expected QDQ node placements, check node placements against expected.
Args:
test_assets (CreateAssetsFolders): Folder organizer.
nn_model_original (tf.keras.Model): Keras model.
quantization_mode (str): quantization mode, can be "full" or "partial".
qspec (QuantizationSpec): QuantizationSpec for model quantization.
custom_qdq_cases (List[CustomQDQInsertionCase]): list of custom QDQ cases for model quantization.
test_name (str): name for this test workflow.
expected_qdq_insertion (List[LayerConfig]): expected QDQ insertion classes and/or layers.
Returns:
q_model (tf.keras.Model): quantized model.
validated (bool): indicates whether the quantized ONNX file is correct or not (according to QDQ node placements).
"""
# Create test folders
test_assets.add_folder(test_name)
test_assets_attr = getattr(test_assets, test_name)
# Save baseline model
tf.keras.models.save_model(nn_model_original, test_assets_attr.fp32_saved_model)
convert_saved_model_to_onnx(
saved_model_dir=test_assets_attr.fp32_saved_model,
onnx_model_path=test_assets_attr.fp32_onnx_model,
)
# Quantize model
q_model = quantize_model(
model=nn_model_original,
quantization_mode=quantization_mode,
quantization_spec=copy.deepcopy(qspec),
custom_qdq_cases=custom_qdq_cases
)
# Save quantized model
tf.keras.models.save_model(q_model, test_assets_attr.int8_saved_model)
convert_saved_model_to_onnx(
saved_model_dir=test_assets_attr.int8_saved_model,
onnx_model_path=test_assets_attr.int8_onnx_model,
)
# Validate QDQ node placements in ONNX file
expected_qdq_insertion = get_expected_qdq_insertion(
tf.keras.models.clone_model(nn_model_original),
qspec_test=copy.deepcopy(qspec),
quantization_mode=quantization_mode,
custom_qdq_cases=custom_qdq_cases,
expected_qdq_insertion_user=expected_qdq_insertion
)
v = ONNXQDQValidator()
validated = v.validate(
test_assets_attr.int8_onnx_model, expected_qdq_insertion, original_keras_model=nn_model_original
)
return q_model, validated
tools/tensorflow-quantization/tests/quantize_config_test.py
+69
View File
@@ -0,0 +1,69 @@
#
# 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_quantization import quantize_config
import tensorflow_quantization.global_config as global_config
from tensorflow_quantization import QuantizationSpec
from network_pool import bilbo_28_28
def test_global_object_creation():
fnq = quantize_config.FullNetworkQuantization()
assert (
len(global_config.G_CONFIG_OBJECT) == 1
), "quantization config class object is not added to the global list"
assert isinstance(
global_config.G_CONFIG_OBJECT[0], quantize_config.FullNetworkQuantization
)
fnq.clean()
tf.keras.backend.clear_session()
def test_quantization_config_layer_names_add():
model = bilbo_28_28()
fnq = quantize_config.FullNetworkQuantization()
qspec = QuantizationSpec()
qspec.add(name="conv_0")
qspec.add(name="conv_2")
qspec.add(name="conv_4")
fnq.add_quantization_spec_object(qspec, model.layers)
assert (
"conv_0" in fnq.layerwise_config
), "There seems to be an issue with layer name addition in `add_special_layers` function"
assert (
"conv_2" in fnq.layerwise_config
), "There seems to be an issue with layer name addition in `add_special_layers` function"
assert (
"conv_4" in fnq.layerwise_config
), "There seems to be an issue with layer name addition in `add_special_layers` function"
fnq.clean()
tf.keras.backend.clear_session()
def test_quantization_config_layer_class_add():
model = bilbo_28_28()
fnq = quantize_config.FullNetworkQuantization()
qspec = QuantizationSpec()
qspec.add(name="Dense", is_keras_class=True)
fnq.add_quantization_spec_object(qspec, model.layers)
assert (
"Dense" in fnq.layer_classes_to_quantize
), "There seems to be an issue with layer class name addition in `add_special_layers` function"
fnq.clean()
tf.keras.backend.clear_session()
tools/tensorflow-quantization/tests/quantize_qdq_insertion_test.py
+158
View File
@@ -0,0 +1,158 @@
#
# 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 sys
import tensorflow as tf
from tensorflow_quantization import QuantizationSpec
from tensorflow_quantization.custom_qdq_cases import ResNetV1QDQCase
from network_pool import frodo_32_32
from onnx_graph_qdq_validator import validate_quantized_model
from tensorflow_quantization.utils import CreateAssetsFolders
import pytest
test_assets = CreateAssetsFolders("test_quantize_qdq_insertion")
# ############################################
# ######### Full Quantize Test ###############
# ############################################
def test_quantize_full():
this_function_name = sys._getframe().f_code.co_name
nn_model_original = frodo_32_32()
q_model, vr = validate_quantized_model(
test_assets, nn_model_original, test_name=this_function_name
)
assert vr, "ONNX QDQ Validation for full network quantization failed!"
# Necessary to clear model layer names from the memory
tf.keras.backend.clear_session()
def test_quantize_full_residual():
this_function_name = sys._getframe().f_code.co_name
nn_model_original = frodo_32_32()
q_model, vr = validate_quantized_model(
test_assets, nn_model_original, custom_qdq_cases=[ResNetV1QDQCase()], test_name=this_function_name
)
assert vr, "ONNX QDQ Validation for quantizing full network with special residual failed!"
tf.keras.backend.clear_session()
# ############################################
# ######### Full Special Quantize Test #######
# ############################################
def test_quantize_full_special_layer():
this_function_name = sys._getframe().f_code.co_name
nn_model_original = frodo_32_32()
# Create a Quantization Spec (dictionary telling how `add` layer should be treated differently).
qspec = QuantizationSpec()
qspec.add(name="add", quantization_index=[0])
q_model, vr = validate_quantized_model(
test_assets, nn_model_original, qspec=qspec, test_name=this_function_name
)
assert vr, "QDQ Validation for full network but one special layer quantization failed!"
tf.keras.backend.clear_session()
# ##########################################
# ######### Partial Quantize Test ##########
# ##########################################
def test_quantize_partial():
this_function_name = sys._getframe().f_code.co_name
nn_model_original = frodo_32_32()
# Create a qspec dictionary to quantize only two layers named 'conv2d_2' and 'dense'
qspec = QuantizationSpec()
qspec.add(name="conv2d_2")
qspec.add(name="dense")
q_model, vr = validate_quantized_model(
test_assets, nn_model_original, quantization_mode="partial", qspec=qspec, test_name=this_function_name
)
assert vr, "ONNX QDQ Validation for partial network quantization failed!"
tf.keras.backend.clear_session()
# ####################################################
# ######### Subset layers Test - Full quantize #######
# ####################################################
def test_quantize_specific_class_maxpool():
this_function_name = sys._getframe().f_code.co_name
nn_model_original = frodo_32_32()
# Create a list with keras layer classes to quantize
qspec = QuantizationSpec()
qspec.add(name="MaxPooling2D", is_keras_class=True)
q_model, vr = validate_quantized_model(
test_assets, nn_model_original, qspec=qspec, test_name=this_function_name
)
assert vr, "ONNX QDQ Validation for specific class `Dense` quantization failed!"
tf.keras.backend.clear_session()
def test_quantize_specific_class_add():
this_function_name = sys._getframe().f_code.co_name
nn_model_original = frodo_32_32()
# Create a list with keras layer classes to quantize
qspec = QuantizationSpec()
qspec.add(name="Add", is_keras_class=True)
q_model, vr = validate_quantized_model(
test_assets, nn_model_original, qspec=qspec, test_name=this_function_name
)
assert vr, "ONNX QDQ Validation for quantizing specific class `Add` failed!"
tf.keras.backend.clear_session()
# ####################################################
# ####### Subset layers Test - Partial quantize ######
# ####################################################
def test_quantize_specific_class_conv2d_partial():
this_function_name = sys._getframe().f_code.co_name
nn_model_original = frodo_32_32()
# Create a list with keras layer classes to quantize
qspec = QuantizationSpec()
qspec.add(name="Conv2D", is_keras_class=True)
q_model, vr = validate_quantized_model(
test_assets, nn_model_original, quantization_mode="partial", qspec=qspec, test_name=this_function_name
)
assert vr, "ONNX QDQ Validation for quantizing specific class `Conv2D` and `conv2d_1` layer failed!"
tf.keras.backend.clear_session()
tools/tensorflow-quantization/tests/quantize_test.py
+206
View File
@@ -0,0 +1,206 @@
#
# 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 contains test cases for `quantize_model` feature.
`quantize_model` feature quantizes all supported layers in the given Keras model with `NVIDIA` quantization scheme.
Tests if weights were copied correctly after quantization and end-to-end training accuracy.
"""
import tensorflow as tf
from tensorflow_quantization import quantize
from tensorflow_quantization import quantize_model
from network_pool import lobelia_28_28
from network_pool import bilbo_28_28
import pytest
import tensorflow_quantization
from tensorflow_quantization.utils import (
CreateAssetsFolders,
convert_saved_model_to_onnx,
)
def _print_model_weights_shapes(model):
"""
Print shapes of all weights
Args:
model: Keras model
"""
print([model.get_weights()[i].shape for i in range(len(model.get_weights()))])
def test_clone_numerics_quantize_whole_model(debug=False):
"""
Checks whether weights are copied correctly when a dummy model is quantized.
"""
model = lobelia_28_28()
if debug:
_print_model_weights_shapes(model)
om_l0_test_weights = model.get_weights()[0][10, :5]
om_l1_test_weights = model.get_weights()[2][10, :5]
# Quantize model
q_model = quantize_model(model)
if debug:
_print_model_weights_shapes(q_model)
qm_l0_test_weights = q_model.get_weights()[1][10, :5]
qm_l1_test_weights = q_model.get_weights()[8][10, :5]
assert all([a == b for a, b in zip(om_l0_test_weights, qm_l0_test_weights)])
assert all([a == b for a, b in zip(om_l1_test_weights, qm_l1_test_weights)])
tf.keras.backend.clear_session()
def test_adding_one_layer_at_a_time():
qspec = quantize.QuantizationSpec()
qspec.add(name="conv2d_1")
qspec.add(name="Dense", is_keras_class=True)
assert isinstance(
qspec.layers[0], quantize.LayerConfig
), "LayerConfig object is not created for newly added layer."
assert (
len(qspec.layers) == 2
), "New layers are not added to layer list of QuantizationSpec."
def test_adding_layer_name_list():
qspec = quantize.QuantizationSpec()
layer_name = ["conv2d", "conv2d_1", "conv2d_7", "dense"]
layer_qip = [True, False, True, False]
layer_idx = [None, [0], None, None]
qspec.add(name=layer_name, quantize_input=layer_qip, quantization_index=layer_idx)
assert (
len(qspec.layers) == 4
), "Four layers are not added to qspec object as expected."
def train_quantize_fine_tune(exp_folder: "Folder", perform_four_bit_quantization: bool = False) -> None:
"""
Train, quantize and fine-tune Keras model using NVIDIA's QAT wrapper library.
Args:
exp_folder (Folder): Base experiment folder object.
perform_four_bit_quantization (bool): If True, 4 bit quantization is performed. 8 bit quantization is default.
Returns:
None
"""
# Load MNIST dataset
mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()
# Normalize the input image so that each pixel value is between 0 to 1.
train_images = train_images / 255.0
test_images = test_images / 255.0
nn_model_original = bilbo_28_28()
# Train original classification model
nn_model_original.compile(
optimizer="adam",
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=["accuracy"],
)
nn_model_original.fit(
train_images, train_labels, batch_size=128, epochs=5, validation_split=0.1
)
# get baseline model accuracy
_, baseline_model_accuracy = nn_model_original.evaluate(
test_images, test_labels, verbose=0
)
print("Baseline test accuracy:", baseline_model_accuracy)
tf.keras.models.save_model(nn_model_original, exp_folder.fp32_saved_model)
convert_saved_model_to_onnx(
saved_model_dir=exp_folder.fp32_saved_model,
onnx_model_path=exp_folder.fp32_onnx_model,
)
if perform_four_bit_quantization:
tensorflow_quantization.G_NUM_BITS = 4
# quantize entire model using `quantize_model` feature
q_model = quantize_model(nn_model_original)
# fine tune annotated model
q_model.compile(
optimizer="adam",
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=["accuracy"],
)
q_model.fit(
train_images, train_labels, batch_size=32, epochs=5, validation_split=0.1
)
# Get quantized accuracy
_, q_aware_model_accuracy = q_model.evaluate(test_images, test_labels, verbose=0)
print("Quant test accuracy:", q_aware_model_accuracy)
assert (
q_aware_model_accuracy >= baseline_model_accuracy or
abs(baseline_model_accuracy - q_aware_model_accuracy) * 100 <= 2.0
), "QAT accuracy is not acceptable: {:.2f} vs {:.2f} for baseline".format(
q_aware_model_accuracy * 100, baseline_model_accuracy * 100
)
# save quantized model and convert to ONNX
tf.keras.models.save_model(q_model, exp_folder.int8_saved_model)
convert_saved_model_to_onnx(
saved_model_dir=exp_folder.int8_saved_model,
onnx_model_path=exp_folder.int8_onnx_model,
)
def test_end_to_end_workflow():
"""
Test end-to-end QAT workflow using the `quantize_model` function.
The following steps are included:
1. Create a dummy model (baseline)
2. Train model on Fashion MNIST dataset
3. Calculate baseline FP32 model accuracy
4. Perform 4 bit (default) quantization and fine-tuning
5. Convert QAT model to ONNX
"""
test_assets = CreateAssetsFolders("test_quantize_end_to_end")
test_assets.add_folder("test_end_to_end_workflow")
train_quantize_fine_tune(test_assets.test_end_to_end_workflow)
tf.keras.backend.clear_session()
@pytest.mark.skip(reason="Just used to test 4 bit quantization feature.")
def test_end_to_end_workflow_4bit():
"""
Test end-to-end QAT workflow using the `quantize_model` function for 4 bit quantization.
The following steps are included:
1. Create a dummy model (baseline)
2. Train model on Fashion MNIST dataset
3. Calculate baseline FP32 model accuracy
4. Perform 4 bit quantization and fine-tuning
5. Convert QAT model to ONNX
"""
test_assets = CreateAssetsFolders("test_quantize_end_to_end")
test_assets.add_folder("test_end_to_end_workflow_4bit")
train_quantize_fine_tune(test_assets.test_end_to_end_workflow_4bit, perform_four_bit_quantization=True)
tf.keras.backend.clear_session()
tools/tensorflow-quantization/tests/quantize_wrapper_base_test.py
+55
View File
@@ -0,0 +1,55 @@
#
# 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.quantize_wrapper_base import BaseQuantizeWrapper
import copy
import pytest
EXPECTED_WRAPPERS = [
"WeightedBaseQuantizeWrapper",
"Conv2DQuantizeWrapper",
"DenseQuantizeWrapper",
"DepthwiseConv2DQuantizeWrapper",
"NonWeightedBaseQuantizeWrapper",
"AveragePooling2DQuantizeWrapper",
"GlobalAveragePooling2DQuantizeWrapper",
"MaxPooling2DQuantizeWrapper",
"BatchNormalizationQuantizeWrapper",
"NonWeightedBaseQuantizeWrapperForMultipleInputs",
"MultiplyQuantizeWrapper",
"ConcatenateQuantizeWrapper",
"AddQuantizeWrapper",
]
def test_old_wrappers_registration():
all_wrappers = BaseQuantizeWrapper.CHILD_WRAPPERS
assert EXPECTED_WRAPPERS == list(all_wrappers.keys())
def test_new_wrapper_registration():
class TestWrapper(BaseQuantizeWrapper):
def __init__(self, layer, **kwargs):
super().__init__(layer, **kwargs)
all_wrappers = BaseQuantizeWrapper.CHILD_WRAPPERS
expected = copy.deepcopy(EXPECTED_WRAPPERS)
expected.append("TestWrapper")
assert expected == list(all_wrappers.keys())
tools/tensorflow-quantization/tests/quantize_wrappers_test.py
+632
View File
@@ -0,0 +1,632 @@
#
# 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 sys
import tensorflow as tf
from tensorflow_quantization import QuantizationSpec
from tensorflow_quantization.quantize import LayerConfig
from onnx_graph_qdq_validator import validate_quantized_model
from tensorflow_quantization.utils import CreateAssetsFolders
from network_pool import (
otho_28_28,
lotho_28_28,
lobelia_28_28,
merry_28_28,
pippin_28_28,
)
import pytest
# Create a directory to save wrapper test data
test_assets = CreateAssetsFolders("test_quantize_wrappers")
# ###################################################
# ####### Conv2D layer wrapper tests ################
# ###################################################
def test_conv2d_wrapper_quant_full_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = otho_28_28()
# Quantization
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="conv_0"),
LayerConfig(name="conv_1"),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
def test_conv2d_wrapper_quant_partial_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = otho_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(name="conv_1")
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="conv_1"),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
def test_conv2d_wrapper_quant_partial_only_input_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = otho_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(name="conv_0", quantize_weight=False)
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="conv_0", quantize_weight=False),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
def test_conv2d_wrapper_quant_partial_only_weight_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = otho_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(name="conv_0", quantize_input=False)
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="conv_0", quantize_input=False),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
# ###################################################
# ####### DepthwiseConv2D layer wrapper tests #######
# ###################################################
def test_depthwise_conv2d_wrapper_quant_full_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = lotho_28_28()
# Quantization
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="dconv_0"),
LayerConfig(name="dconv_1"),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
def test_depthwise_conv2d_wrapper_quant_partial_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = lotho_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(name="dconv_1")
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="dconv_1"),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
def test_depthwise_conv2d_wrapper_quant_partial_only_input_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = lotho_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(name="dconv_1", quantize_weight=False)
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="dconv_1", quantize_weight=False),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
def test_depthwise_conv2d_wrapper_quant_partial_only_weight_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = lotho_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(name="dconv_1", quantize_input=False)
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="dconv_1", quantize_input=False),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
# ###################################################
# ####### Dense layer wrapper tests #################
# ###################################################
def test_dense_wrapper_quant_full_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = lobelia_28_28()
# Quantization
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="dense_0"),
LayerConfig(name="dense_1"),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
def test_dense_wrapper_quant_partial_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = lobelia_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(name="dense_0")
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="dense_0"),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
def test_dense_wrapper_quant_partial_only_input_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = lobelia_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(name="dense_0", quantize_weight=False)
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="dense_0", quantize_weight=False),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
def test_dense_wrapper_quant_partial_only_weight_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = lobelia_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(name="dense_1", quantize_input=False)
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="dense_1", quantize_input=False),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
# ###################################################
# ####### Concatenation layer wrapper tests #########
# ###################################################
def test_concat_wrapper_quant_full_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = merry_28_28()
# Quantization
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="conv2d"),
LayerConfig(name="conv2d_1"),
LayerConfig(name="conv2d_3"),
LayerConfig(name="conv2d_4"),
LayerConfig(name="conv2d_2"),
LayerConfig(name="dense"),
LayerConfig(name="dense_1"),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
def test_concat_wrapper_quant_full_quant_bn_concat_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = merry_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(name="batch_normalization_3")
qspec.add(name="concatenate", quantization_index=[0, 1])
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="conv2d"),
LayerConfig(name="conv2d_1"),
LayerConfig(name="conv2d_3"),
LayerConfig(name="conv2d_4"),
LayerConfig(name="batch_normalization_3"),
LayerConfig(name="conv2d_2"),
LayerConfig(name="concatenate", quantization_index=[0, 1]),
LayerConfig(name="dense"),
LayerConfig(name="dense_1"),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
def test_concat_wrapper_quant_specific_index_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = merry_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(
name="concatenate",
quantize_input=True,
quantize_weight=False,
quantization_index=[0],
)
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="concatenate", quantization_index=[0]),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
# ###################################################
# ####### Add layer wrapper tests ###################
# ###################################################
# Use KerasModelLayersSurgeon() from utils to find layer names.
def test_add_wrapper_quant_full_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = pippin_28_28()
# Quantization
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="conv2d"),
LayerConfig(name="conv2d_1"),
LayerConfig(name="conv2d_2"),
LayerConfig(name="conv2d_3"),
LayerConfig(name="conv2d_4"),
LayerConfig(name="conv2d_6"),
LayerConfig(name="conv2d_5"),
LayerConfig(name="dense"),
LayerConfig(name="dense_1"),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
def test_add_wrapper_quant_partial_specific_index_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = pippin_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(
name="add", quantize_input=True, quantize_weight=False, quantization_index=[1]
)
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="add", quantization_index=[1])
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
def test_add_wrapper_quant_full_specific_index_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = pippin_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(
name="add", quantize_input=True, quantize_weight=False, quantization_index=[1]
)
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="conv2d"),
LayerConfig(name="conv2d_1"),
LayerConfig(name="conv2d_2"),
LayerConfig(name="conv2d_3"),
LayerConfig(name="add", quantization_index=[1]),
LayerConfig(name="conv2d_4"),
LayerConfig(name="conv2d_6"),
LayerConfig(name="conv2d_5"),
LayerConfig(name="dense"),
LayerConfig(name="dense_1"),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
# ########################################################
# ############ Test subset layer class selection #########
# ########################################################
def test_subset_layer_class_selection_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = pippin_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(name="Conv2D", is_keras_class=True)
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="conv2d"),
LayerConfig(name="conv2d_1"),
LayerConfig(name="conv2d_2"),
LayerConfig(name="conv2d_3"),
LayerConfig(name="conv2d_4"),
LayerConfig(name="conv2d_6"),
LayerConfig(name="conv2d_5"),
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode='partial', qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
# ########################################################
# ############ Test missing layer name warning ###########
# ########################################################
def test_missing_layer_name_warning_nv():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = pippin_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(
name="add", quantize_input=True, quantize_weight=False, quantization_index=[1]
)
qspec.add(name="wrong_layer")
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="add", quantization_index=[1])
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode='partial', qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
# ########################################################
# ##### Test Add,Concat out of range index warning #######
# ########################################################
@pytest.mark.skip(
reason="When quantization index out of range does not give error but still wraps \
add layer without quantizing any input"
)
def test_out_of_range_index():
# Create experiment specific directory
this_function_name = sys._getframe().f_code.co_name
# Baseline model
nn_model = pippin_28_28()
# Quantization
qspec = QuantizationSpec()
qspec.add(
name="add", quantize_input=True, quantize_weight=False, quantization_index=[3]
)
qspec.add(name="dense")
# (Optional) QDQ node placement check
# Here just to explicitly show the user which layers are quantized.
expected_qdq_insertion = [
LayerConfig(name="add"), LayerConfig(name="dense")
]
q_model, vr = validate_quantized_model(
test_assets, nn_model, quantization_mode='partial', qspec=qspec, test_name=this_function_name,
expected_qdq_insertion=expected_qdq_insertion
)
assert vr, "ONNX QDQ Validation failed!"
tf.keras.backend.clear_session()
tools/tensorflow-quantization/tests/run_all_tests.sh
Executable
+30
View File
@@ -0,0 +1,30 @@
#!/bin/bash
# clean
rm -rf wrappers_test_saved_models
rm -rf quantize_model_test_saved_models
rm -rf utils_test_saved_models
rm -rf qdq_test_saved_models
rm -rf __pycache__
rm -rf custom_qdq_models
# Run quantize_config tests
python -m pytest quantize_config_test.py -rP
# Run QDQ insertion tests
python -m pytest quantize_qdq_insertion_test.py -rP
# Run wrappers tests
python -m pytest quantize_wrappers_test.py -rP
# Run wrappers base tests
python -m pytest quantize_wrapper_base_test.py -rP
# Run end to end training test
python -m pytest quantize_test.py -rP
# Run special qdq insertion tests
python -m pytest custom_qdq_cases_test.py -rP
# Run utils test
python -m pytest utils_test.py -rP
tools/tensorflow-quantization/tests/utils_test.py
+101
View File
@@ -0,0 +1,101 @@
#
# 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 os
import sys
import tensorflow_quantization.utils as utils
import tensorflow as tf
from tensorflow_quantization import quantize_model
from tensorflow_quantization.utils import (
CreateAssetsFolders,
convert_saved_model_to_onnx,
)
from network_pool import sam_32_32
import pytest
test_assets = CreateAssetsFolders("test_utils")
def test_keras_traveller():
kmt = utils.KerasModelTraveller()
model = sam_32_32()
layer_names = kmt.get_layer_names(keras_model=model)
expected_layer_names = [
"input_1",
"conv2d",
"re_lu",
"conv2d_1",
"re_lu_1",
"conv2d_2",
"re_lu_2",
"conv2d_3",
"add",
"re_lu_3",
"conv2d_4",
"re_lu_4",
"conv2d_5",
"add_1",
"re_lu_5",
"conv2d_6",
"re_lu_6",
"conv2d_7",
"conv2d_8",
"add_2",
"re_lu_7",
"max_pooling2d",
"flatten",
"dense",
"re_lu_8",
"dense_1",
]
assert layer_names == expected_layer_names, "Keras model traveller failed."
tf.keras.backend.clear_session()
def test_convert_to_onnx():
test_assets.add_folder("test_convert_to_onnx")
model = sam_32_32()
q_model = quantize_model(model)
# Create experiment specific directory
tf.keras.models.save_model(
q_model, test_assets.test_convert_to_onnx.int8_saved_model
)
convert_saved_model_to_onnx(
saved_model_dir=test_assets.test_convert_to_onnx.int8_saved_model,
onnx_model_path=test_assets.test_convert_to_onnx.int8_onnx_model,
)
tf.keras.backend.clear_session()
def test_find_my_predecessors():
resnet50 = tf.keras.applications.resnet.ResNet50(weights=None)
r = utils.find_my_predecessors(resnet50, "conv2_block1_add")
assert r[0]["class"] == "BatchNormalization"
assert r[0]["name"] == "conv2_block1_0_bn"
assert r[1]["class"] == "BatchNormalization"
assert r[1]["name"] == "conv2_block1_3_bn"
def test_find_my_successors():
resnet50 = tf.keras.applications.resnet.ResNet50(weights=None)
r = utils.find_my_successors(resnet50, "pool1_pool")
assert r[0]["class"] == "Conv2D"
assert r[0]["name"] == "conv2_block1_1_conv"
assert r[1]["class"] == "Conv2D"
assert r[1]["name"] == "conv2_block1_0_conv"