chore: import upstream snapshot with attribution

2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
@@ -0,0 +1,67 @@
+## About
+This script presents a QAT end-to-end workflow (TF2-to-ONNX) for [ResNet models](https://keras.io/api/applications/resnet/) in `tf.keras.applications`.
+
+### Contents
+[Requirements](#requirements) • [Workflow](#workflow) • [Results](#results)  
+
+## Requirements
+Install base requirements and prepare data. Please refer to [examples' README](../README.md).
+
+## Workflow
+
+### Step 1: Model Quantization and Fine-tuning
+
+Please run the following to quantize, fine-tune, and save the final graph in SavedModel format (checkpoints are also saved).
+
+```sh
+python run_qat_workflow.py
+```
+
+### Step 2: Conversion to ONNX
+Step 1 already does the conversion from SavedModel to ONNX automatically. For manual steps, please see step 3 in [EfficientNet's README](../efficientnet_b0/README.md).
+
+### Step 3: TensorRT Deployment
+Please refer to the [examples' README](../README.md).
+
+## Results
+Results obtained on NVIDIA's A100 GPU and TensorRT 8.4 EA.
+
+### ResNet50-v1
+
+| Model    | TF (%)      | TF latency (ms, bs=1) | TRT(%) | TRT latency (ms, bs=1) |
+|----------|-------------|-----------------------|--------|------------------------|
+| Baseline | 75.05       | 7.95                  | 75.05  | 1.96                   |
+| PTQ      | -           | -                     | 74.96  | 0.46                   |
+| **QAT**  | 75.11 (ep5) | -                     | 75.12  | 0.45                   |
+
+### ResNet50-v2
+
+| Model    | TF (%)       | TF latency (ms, bs=1) | TRT(%)  | TRT latency (ms, bs=1) |
+|----------|--------------|-----------------------|---------|------------------------|
+| Baseline | 75.36        | 6.16                  | 75.37   | 2.35                   |
+| PTQ      | -            | -                     | 75.48   | 0.57                   |
+| **QAT**  | 75.59 (ep5)  | -                     | 75.65   | 0.57                   |
+
+### ResNet101-v1
+
+| Model    | TF (%)       | TF latency (ms, bs=1) | TRT(%) | TRT latency (ms, bs=1) |
+|----------|--------------|-----------------------|--------|------------------------|
+| Baseline | 76.47        | 15.92                 | 76.48  | 3.84                   |
+| PTQ      | -            | -                     | 76.32  | 0.84                   |
+| **QAT**  | 76.33 (ep30) | -                     | 76.26  | 0.84                   |
+
+### ResNet101-v2
+
+| Model    | TF (%) | TF latency (ms, bs=1) | TRT(%) | TRT latency (ms, bs=1) |
+|----------|--------|-----------------------|--------|------------------------|
+| Baseline | 76.89  | 14.13                 | 76.88  | 4.55                   |
+| PTQ      | -      | -                     | 76.94  | 1.05                   |
+| **QAT**  | 77.20  | -                     | 77.15  | 1.05                   |
+> QAT fine-tuning hyper-parameters for ResNet101-v2: `bs=32` (`bs=64` was OOM).
+
+### Notes
+- QAT fine-tuning hyper-parameters:
+  - Optimizer: `piecewise_sgd`, `lr_schedule=[(1.0,1),(0.1,2),(0.01,7)]` (default)
+  - Hyper-parameters: `bs=64, ep=10, lr=0.001`.
+  - Added QDQ nodes in Residual connection.
+- PTQ calibration: `bs=64`.
@@ -0,0 +1,188 @@
+#
+# 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 tensorflow as tf
+
+from tensorflow_quantization.quantize import quantize_model
+from tensorflow_quantization.utils import convert_saved_model_to_onnx
+from tensorflow_quantization.custom_qdq_cases import ResNetV1QDQCase, ResNetV2QDQCase
+
+from examples.utils import ensure_dir
+from examples.data.data_loader import load_data
+from examples.utils_finetuning import (
+    get_finetuned_weights_dirname,
+    fine_tune,
+    compile_model,
+)
+from utils import get_resnet_model
+import gc
+import numpy as np
+import random
+import sys
+import logging
+
+
+RESNET_DEPTH = "50"  # Options=[50, 101]
+RESNET_VERSION = "v1"  # Options=[v1, v2]
+
+# For reference: the baseline model used bs=256, train_epochs=90, train_steps=10.
+HYPERPARAMS = {
+    # ################ Data loading ################
+    "tfrecord_data_dir": "/media/Data/imagenet_data/tf_records",
+    "batch_size": 64,
+    "train_data_size": None,  # If None, consider all data, otherwise, consider subset.
+    "val_data_size": None,  # If None, consider all data, otherwise, consider subset.
+    # ############## Fine-tuning ##################
+    "epochs": 2,
+    "steps_per_epoch": 500,  # Set as 'None' if you want to use the default number of steps. If you use this, make sure the number of steps is <= the number of shards (total number of samples / batch_size). Otherwise, an error will occur.
+    "base_lr": 0.001,
+    "optimizer": "piecewise_sgd",  # Options={sgd, piecewise_sgd, adam}
+    "save_root_dir": "./weights/resnet{}{}_test".format(
+        RESNET_DEPTH, RESNET_VERSION
+    ),  # DIR is updated to reflect hyperparams
+    # ############## Enable/disable tasks ##################
+    "finetune_qat_model": True,  # If True, finetune QAT model. Otherwise, just quantize and load weights if existent.
+    "rewrite_weights_qat_finetuning": True,  # If True, rewrites existing fine-tuned weights. Otherwise, just load weights if they exist.
+    "evaluate_baseline_model": True,
+    "evaluate_qat_model": True,
+    "save_baseline_model": True,
+    "seed": 42,
+}
+
+# Set seed for reproducible results
+os.environ["PYTHONHASHSEED"] = str(HYPERPARAMS["seed"])
+random.seed(HYPERPARAMS["seed"])
+np.random.seed(HYPERPARAMS["seed"])
+tf.random.set_seed(HYPERPARAMS["seed"])
+
+# Create logger and save to out.log
+LOGGER = logging.getLogger()
+LOGGER.setLevel(logging.INFO)
+
+
+def main():
+    # ------------- Initial settings -------------
+    # Create directory to save the fine-tuned weights + add relevant hyperparameters in the name
+    qat_save_finetuned_weights = get_finetuned_weights_dirname(HYPERPARAMS)
+    ensure_dir(qat_save_finetuned_weights)
+
+    # Add terminal and file handlers to logger
+    output_file_handler = logging.FileHandler(
+        os.path.join(qat_save_finetuned_weights, "out.log"), mode="w"
+    )
+    stdout_handler = logging.StreamHandler(sys.stdout)
+    LOGGER.addHandler(output_file_handler)
+    LOGGER.addHandler(stdout_handler)
+
+    # Load data
+    train_batches, val_batches = load_data(
+        HYPERPARAMS, model_name="resnet_{}".format(RESNET_VERSION)
+    )
+
+    # ------------- Baseline model -------------
+    LOGGER.info("------------- Baseline model -------------")
+
+    # Instantiate Baseline model
+    model = get_resnet_model(
+        resnet_depth=RESNET_DEPTH,
+        resnet_version=RESNET_VERSION,
+    )
+
+    if HYPERPARAMS["evaluate_baseline_model"]:
+        # Compile model (needed to evaluate model)
+        compile_model(model)
+        _, baseline_model_accuracy = model.evaluate(val_batches)
+        LOGGER.info("Baseline val accuracy: {}".format(baseline_model_accuracy))
+
+    if HYPERPARAMS["save_baseline_model"]:
+        tf.keras.models.save_model(
+            model, os.path.join(HYPERPARAMS["save_root_dir"], "saved_model_baseline")
+        )
+        convert_saved_model_to_onnx(
+            saved_model_dir=os.path.join(
+                HYPERPARAMS["save_root_dir"], "saved_model_baseline"
+            ),
+            onnx_model_path=os.path.join(
+                HYPERPARAMS["save_root_dir"], "model_baseline.onnx"
+            ),
+        )
+
+    # ------------- QAT model -------------
+    # Quantize model
+    LOGGER.info("\n------------- QAT model -------------")
+    if "v1" == RESNET_VERSION:
+        q_model = quantize_model(model, custom_qdq_cases=[ResNetV1QDQCase()])
+    else:
+        q_model = quantize_model(model, custom_qdq_cases=[ResNetV2QDQCase()])
+
+    finetuned_qat_weights_path = os.path.join(
+        qat_save_finetuned_weights, "checkpoints_best"
+    )
+    # Performs fine-tuning if `rewrite` is enabled or if fine-tuned weights don't exist yet
+    #   (1st time fine-tuning model).
+    if HYPERPARAMS["finetune_qat_model"] and (
+        HYPERPARAMS["rewrite_weights_qat_finetuning"]
+        or not os.path.exists(finetuned_qat_weights_path)
+    ):
+        # Fine-tuning + saving new checkpoints
+        LOGGER.info("\nFine-tuning model...")
+        fine_tune(
+            q_model,
+            train_batches,
+            val_batches,
+            qat_save_finetuned_weights,
+            HYPERPARAMS,
+            LOGGER,
+        )
+        LOGGER.info("Fine-tuning done!")
+
+    # Loads best weights if they exist
+    if os.path.exists(finetuned_qat_weights_path + ".index"):
+        LOGGER.info("Loading fine-tuned weights...")
+        q_model.load_weights(finetuned_qat_weights_path).expect_partial()
+        LOGGER.info("Loaded complete!")
+    compile_model(q_model)
+
+    if HYPERPARAMS["evaluate_qat_model"]:
+        LOGGER.info("\nEvaluating QAT model...")
+        _, qat_model_accuracy = q_model.evaluate(val_batches)
+        LOGGER.info("QAT val accuracy: {}".format(qat_model_accuracy))
+
+    # Save quantized model
+    LOGGER.info("\nSaving QAT model")
+    tf.keras.models.save_model(
+        q_model, os.path.join(qat_save_finetuned_weights, "saved_model")
+    )
+
+    # Clear GPU and invoke Garbage Collector to avoid script ending during ONNX conversion
+    tf.keras.backend.clear_session()
+    gc.collect()
+    del model
+    del q_model
+
+    # Convert SavedModel to ONNX
+    LOGGER.info("\nONNX conversion...")
+    convert_saved_model_to_onnx(
+        saved_model_dir=os.path.join(qat_save_finetuned_weights, "saved_model"),
+        onnx_model_path=os.path.join(qat_save_finetuned_weights, "model.onnx"),
+    )
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,22 @@
+## Scripts for TRT Deployment
+For both baseline and QAT, change:
+- `RESNET_DEPTH` for 50 or 101,
+- `RESNET_VERSION` for v1 or v2,
+- `BS` for which batch sizes you wish to evaluate the engine on.
+
+#### Baseline
+```
+./scripts/deploy_engine_baseline.sh
+```
+> Change `ROOT_DIR` to where your ONNX file is.
+
+#### QAT
+```
+./scripts/deploy_engine_qat.sh
+```
+> Change `QAT_SUBDIR` and `ROOT_DIR` to where your ONNX file is.
+
+### Only accuracy
+```
+./scripts/infer_engine.sh
+```
@@ -0,0 +1,47 @@
+#!/usr/bin/env bash
+
+# Single run:
+# ../../engine_builder/build_engine_single.py --root_dir=/home/nvidia/PycharmProjects/tensorflow-quantization/examples/resnet/weights/resnet50v1 --onnx=model_baseline_dynamic.onnx --engine=model_baseline_dynamic.engine --input=224,224,3 --min_bs=1 --max_bs=1 --opt_bs=1 --precision=fp32
+#
+
+RESNET_DEPTH=50
+RESNET_VERSION=v1
+ROOT_DIR=../weights/resnet${RESNET_DEPTH}${RESNET_VERSION}
+LOGS_SUBDIR=baseline_engines_trtSource
+LOGS_DIR=${ROOT_DIR}/${LOGS_SUBDIR}
+mkdir $LOGS_DIR
+
+echo "1/3. Building engine"
+# bs=32 OOM in workstation
+ONNX=model_baseline_dynamic.onnx
+ENGINE=${LOGS_SUBDIR}/model_baseline_dynamic_bs{min1,opt8,max16}.engine
+python ../../../engine_builder/build_engine_single.py --root_dir=$ROOT_DIR \
+                                                      --onnx=$ONNX \
+                                                      --engine=$ENGINE \
+                                                      --input=224,224,3 \
+                                                      --min_bs=1 --opt_bs=8 --max_bs=16 \
+                                                      --precision=fp32
+wait
+
+for BS in 8 16;  do # 8 32 128; do
+  echo "Model evaluation..."
+  echo "############### bs=${BS} ###############"
+
+  # Latency calculation from built engine
+  echo "2/3. Latency evaluation"
+  trtexec --device=0 \
+          --loadEngine=${ROOT_DIR}/${ENGINE} \
+          --shapes=input_1:0:${BS}x224x224x3 \
+          --workspace=2048 \
+          --separateProfileRun \
+          --dumpProfile \
+          --explicitBatch &> ${LOGS_DIR}/trtexec_latency_bs${BS}.log
+  wait
+
+  echo "3/3. Accuracy evaluation"
+  python ../infer_engine.py --engine=${ROOT_DIR}/${ENGINE} \
+                            --log_file=engine_accuracy_bs${BS}.log \
+                            --model_name=resnet_$RESNET_VERSION \
+                            -b=$BS
+  wait
+done
@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+
+# Single run:
+# ../../engine_builder/build_engine_single.py --root_dir=/home/nvidia/PycharmProjects/tensorflow-quantization/examples/resnet/weights/resnet50v1 --onnx=model_baseline_dynamic.onnx --engine=model_baseline_dynamic.engine --input=224,224,3 --min_bs=1 --max_bs=1 --opt_bs=1 --precision=fp32
+#
+
+RESNET_DEPTH=50
+RESNET_VERSION=v1
+QAT_SUBDIR=qat_tfrecord_ep10_steps500_l2False_baselr0.0001_piecewise_sgd_bs128
+ROOT_DIR=../weights/resnet${RESNET_DEPTH}${RESNET_VERSION}/${QAT_SUBDIR}
+LOGS_SUBDIR=engines_trtSource
+LOGS_DIR=${ROOT_DIR}/${LOGS_SUBDIR}
+mkdir $LOGS_DIR
+
+echo "1/3. Building engine"
+# bs=32 OOM in workstation
+ONNX=model_dynamic.onnx
+ENGINE=${LOGS_SUBDIR}/model_baseline_bs{min1,opt8,max128}.engine
+python ../../../engine_builder/build_engine_single.py --root_dir=$ROOT_DIR \
+                                                      --onnx=$ONNX \
+                                                      --engine=$ENGINE \
+                                                      --input=224,224,3 \
+                                                      --min_bs=1 --opt_bs=8 --max_bs=128 \
+                                                      --precision=int8
+wait
+
+for BS in 1 8 128;  do # 8 32 128; do
+  echo "Model evaluation..."
+  echo "############### bs=${BS} ###############"
+
+  # Latency calculation from built engine
+  echo "2/3. Latency evaluation"
+  trtexec --device=0 \
+          --loadEngine=${ROOT_DIR}/${ENGINE} \
+          --shapes=input_1:0:${BS}x224x224x3 \
+          --workspace=1024 \
+          --separateProfileRun \
+          --dumpProfile \
+          --explicitBatch \
+          --int8 &> ${LOGS_DIR}/trtexec_latency_bs${BS}.log
+  wait
+
+  echo "3/3. Accuracy evaluation"
+  python ../infer_engine.py --engine=${ROOT_DIR}/${ENGINE} \
+                            --log_file=engine_accuracy_bs${BS}.log \
+                            --model_name=resnet_$RESNET_VERSION \
+                            -b=$BS
+  wait
+done
@@ -0,0 +1,17 @@
+ROOT_DIR="/home/nvidia/PycharmProjects/tensorrt_qat/examples/resnet/"
+
+RESNET_DEPTH="50"
+RESNET_VERSION="v1"
+MODEL_TYPE="baseline"  # "qat"
+PRECISION="fp32"  # "int8"
+
+ENGINES_DIR="engines_gtc_trt8.4_gittrt/${MODEL_TYPE}"
+LOGS_DIR="logs_gtc_trt8.4_gittrt/${MODEL_TYPE}"
+
+for BS in 1; do
+  SUBDIR="resnet${RESNET_DEPTH}${RESNET_VERSION}_${PRECISION}_${BS}_sparsity_disable_DLA_disabled"
+
+  python ../infer_engine.py --engine=${ROOT_DIR}/${ENGINES_DIR}/${SUBDIR}.plan \
+                            --log_file=${ROOT_DIR}/${LOGS_DIR}/${SUBDIR}_accuracy.log \
+                            --model_name=resnet_$RESNET_VERSION -b=1
+done
@@ -0,0 +1,113 @@
+#
+# 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.quantize import QuantizationSpec
+from tensorflow_quantization.custom_qdq_cases import ResNetV1QDQCase, ResNetV2QDQCase
+from examples.resnet.utils import get_resnet_model
+from tests.onnx_graph_qdq_validator import validate_quantized_model
+from tensorflow_quantization.utils import CreateAssetsFolders
+import pytest
+
+# Create a directory to save test models
+test_assets = CreateAssetsFolders("test_qdq_node_placement")
+
+
+def test_resnet50v1_quantize_full():
+    """
+    ResNet-v1: Full model quantization
+    """
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Instantiate Baseline model
+    nn_model_original = get_resnet_model(resnet_depth="50", resnet_version="v1")
+
+    custom_qdq_cases = [ResNetV1QDQCase()]
+    q_model, validated = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name,
+        custom_qdq_cases=custom_qdq_cases
+    )
+    assert validated, "ONNX QDQ validation for full network quantization failed!"
+    # necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()
+
+
+def test_resnet50v1_quantize_full_special():
+    """
+    ResNet-v1: Full model quantization with the first Conv layer (conv2_block1_1_conv) not being quantized.
+    """
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Instantiate Baseline model
+    nn_model_original = get_resnet_model(resnet_depth="50", resnet_version="v1")
+
+    # Full quantization
+    custom_qdq_cases = [ResNetV1QDQCase()]
+    # Special case
+    qspec = QuantizationSpec()
+    qspec.add(name="conv2_block1_1_conv", quantize_input=False, quantize_weight=False)
+
+    q_model, validated = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name,
+        qspec=qspec, custom_qdq_cases=custom_qdq_cases
+    )
+
+    assert validated, "ONNX QDQ validation for full network quantization failed!"
+    # necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()
+
+
+def test_resnet50v2_quantize_full():
+    """
+    ResNet-v2: Full model quantization
+    """
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Instantiate Baseline model
+    nn_model_original = get_resnet_model(resnet_depth="50", resnet_version="v2")
+
+    custom_qdq_cases = [ResNetV2QDQCase()]
+    q_model, validated = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name,
+        custom_qdq_cases=custom_qdq_cases
+    )
+    assert validated, "ONNX QDQ validation for full network quantization failed!"
+    # necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()
+
+
+def test_resnet50v2_quantize_full_special():
+    """
+    ResNet-v2: Full model quantization with the first MaxPool layer (pool1_pool) not being quantized.
+    """
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Instantiate Baseline model
+    nn_model_original = get_resnet_model(resnet_depth="50", resnet_version="v2")
+
+    custom_qdq_cases = [ResNetV2QDQCase()]
+    qspec = QuantizationSpec()
+    qspec.add(name="pool1_pool", quantize_input=False, quantize_weight=False)
+    q_model, validated = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name,
+        qspec=qspec, custom_qdq_cases=custom_qdq_cases
+    )
+    assert validated, "ONNX QDQ validation for full network quantization failed!"
+    # necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()
@@ -0,0 +1,45 @@
+#
+# 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 examples.data.data_loader import _DEFAULT_IMAGE_SIZE, _NUM_CHANNELS
+from examples.utils import get_tfkeras_model
+
+
+def get_resnet_model(resnet_depth: str = "50", resnet_version: str = "v1") -> tf.keras.Model:
+    """
+    Creates a native tf.keras ResNet model.
+
+    Args:
+        resnet_depth (str): ResNet depth. Options=[50 (default), 101, 152].
+        resnet_version (str): ResNet version. Options=[v1 (default), v2].
+
+    Returns:
+        model (tf.keras.Model): model corresponding to 'resnet_depth' and 'resnet_version'.
+    """
+
+    shape = (
+        _DEFAULT_IMAGE_SIZE["resnet_{}".format(resnet_version)],
+        _DEFAULT_IMAGE_SIZE["resnet_{}".format(resnet_version)],
+        _NUM_CHANNELS,
+    )
+
+    model_name = "resnet_" + resnet_depth + resnet_version
+    model = get_tfkeras_model(model_name=model_name, shape=shape)
+
+    return model