chore: import upstream snapshot with attribution

tools/tensorflow-quantization/tests/quantize_wrappers_test.py

@@ -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()