chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Waiting to run
Docker Image CI / build-ubuntu2004 (push) Waiting to run
This commit is contained in:
@@ -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 NVIDIA’s 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)
|
||||
- [NVIDIA’s 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
|
||||
Reference in New Issue
Block a user