wehub/nvidia--tensorrt
1
0
Fork 0
Code Issues Pull Requests Actions 27 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Details

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
samples/python/engine_refit_onnx_bidaf/README.md
+163
View File
@@ -0,0 +1,163 @@
# TensorRT Engine Refitting of 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)
- [Additional resources](#additional-resources)
- [License](#license)
- [Changelog](#changelog)
- [Known issues](#known-issues)
## Description
This sample shows how to refit an engine built from an ONNX model via parsers. A modified version of the [ONNX BiDAF model](https://github.com/onnx/models/tree/main/validated/text/machine_comprehension/bidirectional_attention_flow) is used as the sample model, which implements the Bi-Directional Attention Flow (BiDAF) network described in the paper [Bidirectional Attention Flow for Machine Comprehension](https://arxiv.org/abs/1611.01603).
## How does this sample work?
This sample replaces unsupported nodes (HardMax / Compress) in the original ONNX model via ONNX-graphsurgeon (in `prepare_model.py`) and build a refittable TensorRT engine.
The engine is then refitted with fake weights and correct weights, each followed by inference on sample context and query sentences in `build_and_refit_engine.py`.
## Prerequisites
Dependencies required for this sample
1. Install the dependencies for Python:
```bash
pip3 install -r requirements.txt
```
2. TensorRT
3. [ONNX-GraphSurgeon](https://github.com/NVIDIA/TensorRT/tree/main/tools/onnx-graphsurgeon)
4. Download sample data. See the "Download Sample Data" section of [the general setup guide](../README.md).
## Running the sample
The data directory needs to be specified (either via `-d /path/to/data` or environment varaiable `TRT_DATA_DIR`)
when running these scripts. An error will be thrown if not. Taking `TRT_DATA_DIR` approach in following example.
* Prepare the ONNX model. (The data directory needs to be specified.)
```bash
python3 prepare_model.py
```
The output should look similar to the following:
```
Modifying the ONNX model ...
Modified ONNX model saved as bidaf-modified.onnx
Done.
```
The script will modify the original model from [onnx/models](https://github.com/onnx/models/raw/c02f8c8699fc12273649e658b8d2a1a8e32a35d0/text/machine_comprehension/bidirectional_attention_flow/model/bidaf-9.onnx) and save an ONNX model that can be parsed and run by TensorRT.
The original ONNX model contains four CategoryMapper nodes to map the four input string arrays to int arrays.
Since TensorRT does not support string data type and CategoryMapper nodes, we dump out the four maps for the four nodes as json files (`model/CategoryMapper_{4-6}.json`) and use them to preprocess input data.
Now the four inputs become four outputs of the original CategoryMapper nodes.
And unsupported HardMax nodes and Compress nodes are replaced by ArgMax nodes and Gather nodes, respectively.
* Build a TensorRT engine, refit the engine and run inference.
`python3 build_and_refit_engine.py --weights-location GPU`
The script will build a TensorRT engine from the modified ONNX model, and then refit the engine from GPU weights and run inference on sample context and query sentences.
When running the above command for the first time, the output should look similar to the following:
```
Loading ONNX file from path bidaf-modified.onnx...
Beginning ONNX file parsing
[09/25/2023-08:48:16] [TRT] [W] ModelImporter.cpp:407: Make sure input CategoryMapper_4 has Int64 binding.
[09/25/2023-08:48:16] [TRT] [W] ModelImporter.cpp:407: Make sure input CategoryMapper_5 has Int64 binding.
[09/25/2023-08:48:16] [TRT] [W] ModelImporter.cpp:407: Make sure input CategoryMapper_6 has Int64 binding.
[09/25/2023-08:48:16] [TRT] [W] ModelImporter.cpp:407: Make sure input CategoryMapper_7 has Int64 binding.
Completed parsing of ONNX file
Network inputs:
CategoryMapper_4 <class 'numpy.int64'> (-1, 1)
CategoryMapper_5 <class 'numpy.int64'> (-1, 1, 1, 16)
CategoryMapper_6 <class 'numpy.int64'> (-1, 1)
CategoryMapper_7 <class 'numpy.int64'> (-1, 1, 1, 16)
Building an engine from file bidaf-modified.onnx; this may take a while...
Completed creating Engine
Refitting engine from GPU weights...
Engine refitted in 39.88 ms.
Doing inference...
Doing inference...
Refitting engine from GPU weights...
Engine refitted in 0.27 ms.
Doing inference...
Doing inference...
Passed
```
Note that refitting for second time will be much faster than the first time.
When running the above command again, engine will be deserialized from the plan file, the output should look similar to the following:
```
Reading engine from file bidaf.trt...
Refitting engine from GPU weights...
Engine refitted in 32.64 ms.
Doing inference...
Doing inference...
Refitting engine from GPU weights...
Engine refitted in 0.41 ms.
Doing inference...
Doing inference...
Passed
```
To refit the engine from CPU weights, change the command to be `python3 build_and_refit_engine.py --weights-location CPU`. And the output should look similar to the following
```
Reading engine from file bidaf.trt...
Refitting engine from CPU weights...
Engine refitted in 45.18 ms.
Doing inference...
Doing inference...
Refitting engine from CPU weights...
Engine refitted in 1.20 ms.
Doing inference...
Doing inference...
Passed
```
There is also an option `--version-compatible` to enable engine version compatibility. If installed, `tensorrt_dispatch` package will used for refitting and running version compatible engines instead of `tensorrt` package.
To build and refit a version compatible engine, run the command `python3 build_and_refit_engine.py --version-compatible` and the output should look similar to the above cases.
# Additional resources
The following resources provide a deeper understanding about the model used in this sample:
**Model**
- [Bidirectional Attention Flow for Machine Comprehension](https://arxiv.org/abs/1611.01603)
**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)
- [Importing A Model Using A Parser In Python](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#import_model_python)
- [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.
January 2024:
- Add support for refitting version compatible engines.
August 2023:
- Add support for refitting engines from GPU weights.
- Removed support for Python versions < 3.8.
October 2020: This sample was recreated, updated and reviewed.
# Known issues
There are no known issues in this sample.
samples/python/engine_refit_onnx_bidaf/build_and_refit_engine.py
+298
View File
@@ -0,0 +1,298 @@
#!/usr/bin/env python3
#
# 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 numpy as np
import argparse
import tensorrt as trt
sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir))
from cuda.bindings import runtime as cudart
TRT_LOGGER = trt.Logger()
def get_plan(onnx_file_path, engine_file_path, version_compatible):
"""Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it."""
def build_plan():
"""Takes an ONNX file and creates a TensorRT engine to run inference with"""
import tensorrt as trt
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
parser = trt.OnnxParser(network, TRT_LOGGER)
# Parse model file
print("Loading ONNX file from path {}...".format(onnx_file_path))
with open(onnx_file_path, "rb") as model:
print("Beginning ONNX file parsing")
if not parser.parse(model.read()):
print("ERROR: Failed to parse the ONNX file.")
for error in range(parser.num_errors):
print(parser.get_error(error))
return None
print("Completed parsing of ONNX file")
# Print input info
print("Network inputs:")
for i in range(network.num_inputs):
tensor = network.get_input(i)
print(tensor.name, trt.nptype(tensor.dtype), tensor.shape)
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.REFIT)
if version_compatible:
config.set_flag(trt.BuilderFlag.VERSION_COMPATIBLE)
for opt in [6, 10]:
profile = builder.create_optimization_profile()
input0_min = (1, 1)
input0_opt = (opt, 1)
input0_max = (15, 1)
profile.set_shape(
network.get_input(0).name,
min=input0_min,
opt=input0_opt,
max=input0_max,
)
input1_min = (1, 1, 1, 16)
input1_opt = (opt, 1, 1, 16)
input1_max = (15, 1, 1, 16)
profile.set_shape(
network.get_input(1).name,
min=input1_min,
opt=input1_opt,
max=input1_max,
)
input2_min = (1, 1)
input2_opt = (opt, 1)
input2_max = (15, 1)
profile.set_shape(
network.get_input(2).name,
min=input2_min,
opt=input2_opt,
max=input2_max,
)
input3_min = (1, 1, 1, 16)
input3_opt = (opt, 1, 1, 16)
input3_max = (15, 1, 1, 16)
profile.set_shape(
network.get_input(3).name,
min=input3_min,
opt=input3_opt,
max=input3_max,
)
config.add_optimization_profile(profile)
print(
"Building an engine from file {}; this may take a while...".format(
onnx_file_path
)
)
plan = builder.build_serialized_network(network, config)
print("Completed creating Engine")
with open(engine_file_path, "wb") as f:
f.write(plan)
return plan
if os.path.exists(engine_file_path):
# If a serialized engine exists, use it instead of building an engine.
print("Reading engine from file {}...".format(engine_file_path))
f = open(engine_file_path, "rb")
return f.read()
return build_plan()
def main():
global trt
global TRT_LOGGER
parser = argparse.ArgumentParser()
parser.add_argument(
"-l",
"--weights-location",
dest="weights_location",
default="GPU",
choices=["GPU", "CPU"],
help="The location for weights passed to refitter, either GPU/CPU, default: GPU",
)
parser.add_argument(
"--version-compatible",
dest="version_compatible",
action="store_true",
help="Build a version compatible engine for refitting",
)
args = parser.parse_args()
onnx_file_path = "bidaf-modified.onnx"
engine_file_path = "bidaf{}.trt".format("-vc" if args.version_compatible else "")
plan = get_plan(onnx_file_path, engine_file_path, args.version_compatible)
if args.version_compatible:
# Try using dispatch runtime for refitting and inference. If failed, fallback to full runtime.
try:
del sys.modules["tensorrt"]
sys.modules["tensorrt"] = __import__("tensorrt_dispatch")
sys.modules["trt"] = sys.modules["tensorrt"]
import tensorrt_dispatch as trt
print(
"Importing tensorrt_dispatch instead of full tensorrt for refitting and running vc engines."
)
except:
print(
"Failed to import tensorrt_dispatch for refitting and running vc engines. Please install the package first!"
)
sys.modules["tensorrt"] = __import__("tensorrt")
TRT_LOGGER = trt.Logger()
engine = None
with open(engine_file_path, "rb") as f:
runtime = trt.Runtime(TRT_LOGGER)
if args.version_compatible:
runtime.engine_host_code_allowed = True
engine = runtime.deserialize_cuda_engine(plan)
# should be after get_engine
from data_processing import get_inputs, preprocess
import common_runtime as common
# input
context = "A quick brown fox jumps over the lazy dog."
query = "What color is the fox?"
cw_str, _ = preprocess(context)
# get ravelled data
cw, cc, qw, qc = get_inputs(context, query)
# Do inference with TensorRT
weights_names = ["Parameter576_B_0", "W_0"]
refit_weights_dict = {
name: np.load("{}.npy".format(name)) for name in weights_names
}
fake_weights_dict = {
name: np.ones_like(weights) for name, weights in refit_weights_dict.items()
}
device_mem_dict = {}
if args.weights_location == "GPU":
for name, weights in refit_weights_dict.items():
nbytes = weights.size * weights.itemsize
device_mem_dict[name] = common.DeviceMem(nbytes)
execution_context = engine.create_execution_context()
refitter = trt.Refitter(engine, TRT_LOGGER)
# Skip weights validation since we are confident that the new weights are similar to the weights used to build engine.
refitter.weights_validation = False
# To get a list of all refittable weights' names
# in the network, use refitter.get_all_weights().
if args.weights_location == "GPU":
for name, device_mem in device_mem_dict.items():
device_weights = trt.Weights(
trt.DataType.FLOAT, device_mem.device_ptr, refit_weights_dict[name].size
)
weights_prototype = refitter.get_weights_prototype(name)
assert device_weights.dtype == weights_prototype.dtype
assert device_weights.size == weights_prototype.size
refitter.set_named_weights(name, device_weights, trt.TensorLocation.DEVICE)
for weights_dict, answer_correct in [
(fake_weights_dict, False),
(refit_weights_dict, True),
]:
import time
T1 = time.perf_counter()
device_mem_list = []
# Refit named weights via set_named_weights
for name in weights_names:
host_weights = weights_dict[name]
if args.weights_location == "CPU":
weights = host_weights
location = trt.TensorLocation.HOST
refitter.set_named_weights(name, weights, location)
else:
common.memcpy_host_to_device(device_mem_dict[name].device_ptr, host_weights)
# Get missing weights names. This should return empty lists in this case.
missing_weights = refitter.get_missing_weights()
assert (
len(missing_weights) == 0
), "Refitter found missing weights. Call set_named_weights() or set_weights() for all missing weights"
print(f"Refitting engine from {args.weights_location} weights...")
# Refit the engine with the new weights. This will return True if the refit operation succeeded.
assert refitter.refit_cuda_engine()
T2 = time.perf_counter()
print("Engine refitted in {:.2f} ms.".format((T2 - T1) * 1000))
for profile_idx in range(engine.num_optimization_profiles):
print("Doing inference...")
# Do inference
inputs, outputs, bindings = common.allocate_buffers(
engine, profile_idx
)
padding_bindings = [0] * (len(bindings) * profile_idx)
new_bindings = padding_bindings + bindings
# Use context manager for proper stream lifecycle management
with common.CudaStreamContext() as stream:
# Set host input. The common.do_inference function will copy the input to the GPU before executing.
inputs[0].host = cw
inputs[1].host = cc
inputs[2].host = qw
inputs[3].host = qc
execution_context.set_optimization_profile_async(profile_idx, stream.stream)
execution_context.set_input_shape("CategoryMapper_4", (10, 1))
execution_context.set_input_shape("CategoryMapper_5", (10, 1, 1, 16))
execution_context.set_input_shape("CategoryMapper_6", (6, 1))
execution_context.set_input_shape("CategoryMapper_7", (6, 1, 1, 16))
trt_outputs = common.do_inference(
execution_context,
engine=engine,
bindings=bindings,
inputs=inputs,
outputs=outputs,
stream=stream,
)
start = trt_outputs[0].item()
end = trt_outputs[1].item()
answer = [w.encode() for w in cw_str[start : end + 1].reshape(-1)]
assert answer_correct == (answer == [b"brown"]), answer
common.free_buffers(inputs, outputs)
for _, device_mem in device_mem_dict.items():
device_mem.free()
print("Passed")
if __name__ == "__main__":
main()
samples/python/engine_refit_onnx_bidaf/data_processing.py
+64
View File
@@ -0,0 +1,64 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 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 numpy as np
import nltk
from nltk import word_tokenize
import json
import tensorrt as trt
def preprocess(text):
try:
nltk.data.find("tokenizers/punkt_tab")
except LookupError:
nltk.download("punkt_tab")
tokens = word_tokenize(text)
# split into lower-case word tokens, in numpy array with shape of (seq, 1)
words = np.asarray([w.lower() for w in tokens]).reshape(-1, 1)
# split words into chars, in numpy array with shape of (seq, 1, 1, 16)
chars = [[c for c in t][:16] for t in tokens]
chars = [cs + [""] * (16 - len(cs)) for cs in chars]
chars = np.asarray(chars).reshape(-1, 1, 1, 16)
return words, chars
def get_map_func(filepath):
file = open(filepath)
category_map = json.load(file)
category_mapper = dict(
zip(category_map["cats_strings"], category_map["cats_int64s"])
)
default_int64 = category_map["default_int64"]
func = lambda s: category_mapper.get(s, default_int64)
return np.vectorize(func)
def get_inputs(context, query):
cw, cc = preprocess(context)
qw, qc = preprocess(query)
context_word_func = get_map_func("CategoryMapper_4.json")
context_char_func = get_map_func("CategoryMapper_5.json")
query_word_func = get_map_func("CategoryMapper_6.json")
query_char_func = get_map_func("CategoryMapper_7.json")
cw_input = context_word_func(cw).astype(trt.nptype(trt.int32)).ravel()
cc_input = context_char_func(cc).astype(trt.nptype(trt.int32)).ravel()
qw_input = query_word_func(qw).astype(trt.nptype(trt.int32)).ravel()
qc_input = query_char_func(qc).astype(trt.nptype(trt.int32)).ravel()
return cw_input, cc_input, qw_input, qc_input
samples/python/engine_refit_onnx_bidaf/download.yml
+21
View File
@@ -0,0 +1,21 @@
#
# SPDX-FileCopyrightText: Copyright (c) 2020-2024 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.
#
sample: engine_refit_onnx_bidaf
files:
- path: samples/python/engine_refit_onnx_bidaf/bidaf-original.onnx
url: https://github.com/onnx/models/raw/c02f8c8699fc12273649e658b8d2a1a8e32a35d0/text/machine_comprehension/bidirectional_attention_flow/model/bidaf-9.onnx
checksum: cf11f1eceb4731f8dd39345467fe94a1
samples/python/engine_refit_onnx_bidaf/prepare_model.py
+107
View File
@@ -0,0 +1,107 @@
#
# 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 onnx_graphsurgeon as gs
import onnx
import numpy as np
import json
import sys, os
sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir))
from downloader import getFilePath
def drop_category_mapper_nodes(graph):
new_inputs = []
for org_input in graph.inputs:
# head node, simply disconnect it with others
assert len(org_input.outputs) == 1
category_mapper_node = org_input.outputs[0]
assert category_mapper_node.op == "CategoryMapper"
assert len(category_mapper_node.outputs) == 1
new_inputs.append(category_mapper_node.outputs[0])
category_mapper_node.inputs.clear()
category_mapper_node.outputs.clear()
# Save mapping info to preprocess inputs.
with open(category_mapper_node.name + ".json", "w") as fp:
json.dump(category_mapper_node.attrs, fp)
graph.inputs = new_inputs
def replace_unsupported_ops(graph):
# replace hardmax with ArgMax
hardmaxes = [node for node in graph.nodes if node.op == "Hardmax"]
assert len(hardmaxes) == 1
hardmax = hardmaxes[0]
hardmax.op = "ArgMax"
hardmax.name = "ArgMax(org:" + hardmax.name + ")"
hardmax.attrs["axis"] = 1
hardmax.attrs["keepdims"] = 0
cast = hardmax.o()
reshape = cast.o()
hardmax.outputs = reshape.outputs
assert len(hardmax.outputs) == 1
hardmax.outputs[0].dtype = np.int64
hardmax.outputs[0].shape = [1]
compress = reshape.o()
compress.op = "Gather"
compress.name = "Gather(org:" + compress.name + ")"
compress.attrs["axis"] = 1
cast.outputs.clear()
reshape.outputs.clear()
# Remove the node from the graph completely
graph.cleanup().toposort()
def save_weights_for_refitting(graph):
# Save weights for refitting
tmap = graph.tensors()
np.save("Parameter576_B_0.npy", tmap["Parameter576_B_0"].values)
np.save("W_0.npy", tmap["W_0"].values)
def main():
org_model_file_path = getFilePath(
"samples/python/engine_refit_onnx_bidaf/bidaf-original.onnx"
)
print("Modifying the ONNX model ...")
original_model = onnx.load(org_model_file_path)
graph = gs.import_onnx(original_model)
drop_category_mapper_nodes(graph)
replace_unsupported_ops(graph)
save_weights_for_refitting(graph)
new_model = gs.export_onnx(graph)
modified_model_name = "bidaf-modified.onnx"
onnx.checker.check_model(new_model)
onnx.save(new_model, modified_model_name)
print("Modified ONNX model saved as {}".format(modified_model_name))
print("Done.")
if __name__ == "__main__":
main()
samples/python/engine_refit_onnx_bidaf/requirements.txt
+9
View File
@@ -0,0 +1,9 @@
onnx==1.18.0
nltk==3.9.1
wget==3.2
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