chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Waiting to run

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
@@ -0,0 +1,91 @@
# Introduction To Building and Refitting Weight-stripped Engines from ONNX Models
**Table Of Contents**
- [Description](#description)
- [How does this sample work?](#how-does-this-sample-work)
- [Prerequisites](#prerequisites)
- [Running the sample](#running-the-sample)
* [Sample `--help` options](#sample-help-options)
- [Additional resources](#additional-resources)
- [License](#license)
- [Changelog](#changelog)
- [Known issues](#known-issues)
## Description
This sample, sample_weight_stripping, is a Python sample which uses TensorRT to build a weight-stripped engine and later refit to a full engine for inference.
## How does this sample work?
This sample demonstrates how to build a weight-stripped engine from an ONNX model file using TensorRT Python API which can reduce the saved engine size. Later, the weight-stripped engine is refitted by parser refitter with the original ONNX model as input. The refitted full engine is used for inference and guarantees no performance and accuracy loss. In this sample, we use ResNet50 to showcase our features.
## Prerequisites
1. Install the dependencies for Python.
```bash
pip3 install -r requirements.txt
```
2. Preparing sample data
See [Preparing sample data](../../README.md#preparing-sample-data) in the main samples README.
## Running the sample
1. Build and save both normal engine and weight-stripped engine:
```
python3 build_engines.py --output_stripped_engine=stripped_engine.trt --output_normal_engine=normal_engine.trt
```
After running this step, you can see two saved TensorRT engines. `stripped_engine.trt` contains a stripped engine (~2.3MB) and `normal_engine.trt` contains a normal engine with all weights included (~51MB). By using stripped engine build, we can greatly reduce the size of the saved engine file.
**Note:** If the TensorRT sample data is not installed in the default location, for example `/usr/src/tensorrt/data/`, the model directory must be specified. For example: `--stripped_onnx=/path/to/my/data/` sets the model path for building weight-stripped engine and `--original_onnx=/path/to/my/data/` sets the model path for building normal engine. In most of the cases, they can use the same ONNX model.
2. Refit the weight-stripped engine and perform inference with the weight-stripped engine and the normal engine:
```
python3 refit_engine_and_infer.py --stripped_engine=stripped_engine.trt -normal_engine=normal_engine.trt
```
3. Verify that the sample ran successfully. If the sample runs successfully you should see output similar to the following. The prediction results of the refitted stripped engine is the same as the normal engine. There is no performance loss.
```
Normal engine inference time on 100 cases: 0.1066 seconds
Refitted stripped engine inference time on 100 cases: 0.0606 seconds
Normal engine correctly recognized data/samples/resnet50/tabby_tiger_cat.jpg as tiger cat
Refitted stripped engine correctly recognized data/samples/resnet50/tabby_tiger_cat.jpg as tiger cat
```
### Sample --help options
To see the full list of available options and their descriptions, use the `-h` or `--help` command line option.
# Additional resources
The following resources provide a deeper understanding about importing a model into TensorRT using Python:
**ResNet-50**
- [Deep Residual Learning for Image Recognition](https://arxiv.org/pdf/1512.03385.pdf)
**Documentation**
- [Introduction To NVIDIAs TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
- [NVIDIAs TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
# License
For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
# Changelog
October 2025
Migrate to strongly typed APIs.
August 2025
Removed support for Python versions < 3.10.
February 2024
Initial release of this sample.
# Known issues
There are no known issues in this sample.
@@ -0,0 +1,115 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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 argparse
import math
import time
import datetime
import tensorrt as trt
sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir))
import common
# You can set the logger severity higher to suppress messages (or lower to display more messages).
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
def convert_size(size_bytes):
if size_bytes == 0:
return "0B"
size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
i = int(math.floor(math.log(size_bytes, 1024)))
p = math.pow(1024, i)
s = round(size_bytes / p, 2)
return "%s %s" % (s, size_name[i])
def main(args):
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
with open(args.original_onnx, 'rb') as onnx_model:
parser.parse(onnx_model.read())
with builder.create_builder_config() as config:
config.set_flag(trt.BuilderFlag.STRIP_PLAN)
cache = config.create_timing_cache(b"")
config.set_timing_cache(cache, ignore_mismatch = False)
profile = builder.create_optimization_profile()
profile.set_shape("gpu_0/data_0", min=[1, 3, 224, 224], opt=[1, 3, 224, 224], max=[1, 3, 224, 224])
config.add_optimization_profile(profile)
def build_and_save_engine(builder, network, config, output):
start_time = time.time()
engine_bytes = builder.build_serialized_network(network, config)
assert engine_bytes is not None
with open(output, 'wb') as f:
f.write(engine_bytes)
total_time = time.time() - start_time
print("built and saved {} in time {}".format(output, str(datetime.timedelta(seconds=int(total_time)))))
# build weight-stripped engine and generate timing cache.
build_and_save_engine(builder, network, config, args.output_stripped_engine)
# build normal engine with the same timing cache.
config.flags &= ~(1 << int(trt.BuilderFlag.STRIP_PLAN))
build_and_save_engine(builder, network, config, args.output_normal_engine)
def get_default_model_file():
# Set the data path to the directory that contains the ONNX model.
_, data_files = common.find_sample_data(
description="Runs a ResNet50 network with a TensorRT inference engine.",
subfolder="resnet50",
find_files=["ResNet50.onnx"],
)
onnx_model_file = data_files[0]
return onnx_model_file
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--stripped_onnx", default=None, type=str,
help="The ONNX model file to load for building stripped engine.")
parser.add_argument("--original_onnx", default=None, type=str,
help="The ONNX model file to load for building normal engine.")
parser.add_argument("--output_stripped_engine", default='stripped_engine.trt', type=str,
help="The output path for the weight-stripped TRT engine.")
parser.add_argument("--output_normal_engine", default='normal_engine.trt', type=str,
help="The output path for the full TRT engine.")
args, _ = parser.parse_known_args()
onnx_model_file = get_default_model_file()
if args.stripped_onnx is None:
args.stripped_onnx = onnx_model_file
if args.original_onnx is None:
args.original_onnx = onnx_model_file
if not os.path.exists(args.stripped_onnx):
parser.print_help()
print(f"--stripped_onnx {args.stripped_onnx} does not exist.")
sys.exit(1)
if not os.path.exists(args.original_onnx):
parser.print_help()
print(f"--original_onnx {args.original_onnx} does not exist.")
sys.exit(1)
main(args)
File diff suppressed because one or more lines are too long
@@ -0,0 +1,173 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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 argparse
import random
import sys
import time
import datetime
import numpy as np
import tensorrt as trt
from PIL import Image
sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir))
import common
# You can set the logger severity higher to suppress messages (or lower to display more messages).
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
class ModelData(object):
MODEL_PATH = "ResNet50.onnx"
INPUT_SHAPE = (3, 224, 224)
# We can convert TensorRT data types to numpy types with trt.nptype()
DTYPE = trt.float32
def load_stripped_engine_and_refit(input_file, onnx_model_path):
runtime = trt.Runtime(TRT_LOGGER)
with open(input_file, 'rb') as engine_file:
engine = runtime.deserialize_cuda_engine(engine_file.read())
refitter = trt.Refitter(engine, TRT_LOGGER)
parser_refitter = trt.OnnxParserRefitter(refitter, TRT_LOGGER)
assert parser_refitter.refit_from_file(onnx_model_path)
assert refitter.refit_cuda_engine()
return engine
def load_normal_engine(input_file):
runtime = trt.Runtime(TRT_LOGGER)
with open(input_file, 'rb') as engine_file:
engine = runtime.deserialize_cuda_engine(engine_file.read())
return engine
def load_normalized_test_case(test_image, pagelocked_buffer):
# Converts the input image to a CHW Numpy array
def normalize_image(image):
# Resize, antialias and transpose the image to CHW.
c, h, w = ModelData.INPUT_SHAPE
image_arr = (
np.asarray(image.resize((w, h), Image.LANCZOS))
.transpose([2, 0, 1])
.astype(trt.nptype(ModelData.DTYPE))
.ravel()
)
# This particular ResNet50 model requires some preprocessing, specifically, mean normalization.
return (image_arr / 255.0 - 0.45) / 0.225
# Normalize the image and copy to pagelocked memory.
np.copyto(pagelocked_buffer, normalize_image(Image.open(test_image)))
return test_image
def main(args):
# Set the data path to the directory that contains the trained models and test images for inference.
_, data_files = common.find_sample_data(
description="Runs a ResNet50 network with a TensorRT inference engine.",
subfolder="resnet50",
find_files=[
"binoculars.jpeg",
"reflex_camera.jpeg",
"tabby_tiger_cat.jpg",
ModelData.MODEL_PATH,
"class_labels.txt",
],
)
# Get test images, models and labels.
test_images = data_files[0:3]
onnx_model_file, labels_file = data_files[3:]
labels = open(labels_file, "r").read().split("\n")
# Load a TensorRT engine.
engine = load_normal_engine(args.normal_engine)
refitted_engine = load_stripped_engine_and_refit(args.stripped_engine, onnx_model_file)
# Allocate buffers
inputs, outputs, bindings = common.allocate_buffers(engine)
inputs_1, outputs_1, bindings_1 = common.allocate_buffers(refitted_engine)
# Contexts are used to perform inference.
context = engine.create_execution_context()
context_1 = refitted_engine.create_execution_context()
# Load a normalized test case into the host input page-locked buffer.
test_image = random.choice(test_images)
test_case = load_normalized_test_case(test_image, inputs[0].host)
test_case_1 = load_normalized_test_case(test_image, inputs_1[0].host)
# Run the engine. The output will be a 1D tensor of length 1000, where each value represents the
# probability that the image corresponds to that label
# Use context manager for proper stream lifecycle management - Normal engine
trt_outputs = []
with common.CudaStreamContext() as stream:
start_time = time.time()
for i in range(100): # count time for 100 times of inference
trt_outputs = common.do_inference(context, engine=engine, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
total_time = time.time() - start_time
print("Normal engine inference time on 100 cases: {:.4f} seconds".format(total_time))
# Use context manager for proper stream lifecycle management - Refitted engine
trt_outputs_refitted = []
with common.CudaStreamContext() as stream_1:
start_time = time.time()
for i in range(100):
trt_outputs_refitted = common.do_inference(context_1, engine=refitted_engine, bindings=bindings_1, inputs=inputs_1, outputs=outputs_1, stream=stream_1)
total_time = time.time() - start_time
print("Refitted stripped engine inference time on 100 cases: {:.4f} seconds".format(total_time))
# We use the highest probability as our prediction. Its index corresponds to the predicted label.
pred = labels[np.argmax(trt_outputs[0])]
if "_".join(pred.split()) in os.path.splitext(os.path.basename(test_case))[0]:
print("Normal engine correctly recognized " + test_case + " as " + pred)
else:
print("Normal engine incorrectly recognized " + test_case + " as " + pred)
exit(1)
pred_refitted = labels[np.argmax(trt_outputs_refitted[0])]
if "_".join(pred_refitted.split()) in os.path.splitext(os.path.basename(test_case_1))[0]:
print("Refitted stripped engine correctly recognized " + test_case + " as " + pred_refitted)
else:
print("Refitted stripped engine incorrectly recognized " + test_case + " as " + pred_refitted)
exit(1)
return trt_outputs, trt_outputs_refitted
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--stripped_engine", default='stripped_engine.trt', type=str,
help="The stripped engine file to load.")
parser.add_argument("--normal_engine", default='normal_engine.trt', type=str,
help="The normal engine file to load.")
args, _ = parser.parse_known_args()
if not os.path.exists(args.stripped_engine):
parser.print_help()
print(f"--stripped_engine {args.stripped_engine} does not exist.")
sys.exit(1)
if not os.path.exists(args.normal_engine):
parser.print_help()
print(f"--normal_engine {args.normal_engine} does not exist.")
sys.exit(1)
trt_outputs, trt_outputs_refitted = main(args)
print("The MSE of the final layer output is", np.square(np.subtract(trt_outputs, trt_outputs_refitted)).mean())
@@ -0,0 +1,7 @@
Pillow==11.3.0
cuda-python==12.9.0
pywin32; platform_system == "Windows"
pyyaml==6.0.3
requests==2.32.4
tqdm==4.66.4
numpy==1.26.4