chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
This commit is contained in:
@@ -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,
|
||||
)
|
||||
Reference in New Issue
Block a user