Skip to content

Commit 4442153

Browse files
kevinch-nvkevinch-nv
authored
ONNX-TensorRT 10.6 GA Release (#1000)
Signed-off-by: Kevin Chen <[email protected]>
1 parent 886aff9 commit 4442153

20 files changed

+355
-223
lines changed

Diff for: ‎.gitmodules

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
[submodule "third_party/onnx"]
22
path = third_party/onnx
33
url = https://github.com/onnx/onnx.git
4-
branch = v1.16.0
4+
branch = v1.17.0

Diff for: ‎CMakeLists.txt

+10-4
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ add_definitions("-DSOURCE_LENGTH=${SOURCE_LENGTH}")
2828
# Version information
2929
#--------------------------------------------------
3030
set(ONNX2TRT_MAJOR 10)
31-
set(ONNX2TRT_MINOR 5)
31+
set(ONNX2TRT_MINOR 6)
3232
set(ONNX2TRT_PATCH 0)
3333
set(ONNX2TRT_VERSION "${ONNX2TRT_MAJOR}.${ONNX2TRT_MINOR}.${ONNX2TRT_PATCH}" CACHE STRING "ONNX2TRT version")
3434

@@ -106,6 +106,12 @@ find_path(TENSORRT_INCLUDE_DIR NvInfer.h
106106
PATH_SUFFIXES include)
107107
MESSAGE(STATUS "Found TensorRT headers at ${TENSORRT_INCLUDE_DIR}")
108108

109+
# TensorRT Python Headers
110+
find_path(TENSORRT_PYTHON_INCLUDE_DIR plugin.h
111+
HINTS ${TENSORRT_ROOT}
112+
PATH_SUFFIXES python/include/impl)
113+
message(NOTICE "Found TensorRT Python headers at ${TENSORRT_PYTHON_INCLUDE_DIR}")
114+
109115
# Output dynamic library names depends on platform:
110116
if (MSVC)
111117
set(nvonnxparser_lib_name "nvonnxparser_${ONNX2TRT_MAJOR}")
@@ -119,7 +125,7 @@ set(nvonnxparser_lib_name_static "nvonnxparser_static")
119125
# Importer library
120126
# --------------------------------
121127
add_library(${nvonnxparser_lib_name} SHARED ${IMPORTER_SOURCES})
122-
target_include_directories(${nvonnxparser_lib_name} PUBLIC ${ONNX_INCLUDE_DIRS} ${TENSORRT_INCLUDE_DIR} ${CUDA_INCLUDE_DIR})
128+
target_include_directories(${nvonnxparser_lib_name} PUBLIC ${ONNX_INCLUDE_DIRS} ${TENSORRT_INCLUDE_DIR} ${TENSORRT_PYTHON_INCLUDE_DIR} ${CUDA_INCLUDE_DIR})
123129
target_link_libraries(${nvonnxparser_lib_name} PUBLIC onnx_proto ${PROTOBUF_LIBRARY})
124130
set_target_properties(${nvonnxparser_lib_name} PROPERTIES
125131
VERSION ${ONNX2TRT_VERSION}
@@ -131,7 +137,7 @@ set_target_properties(${nvonnxparser_lib_name} PROPERTIES
131137
RUNTIME_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
132138
)
133139
add_library(${nvonnxparser_lib_name_static} STATIC ${IMPORTER_SOURCES})
134-
target_include_directories(${nvonnxparser_lib_name_static} PUBLIC ${ONNX_INCLUDE_DIRS} ${TENSORRT_INCLUDE_DIR} ${CUDA_INCLUDE_DIR})
140+
target_include_directories(${nvonnxparser_lib_name_static} PUBLIC ${ONNX_INCLUDE_DIRS} ${TENSORRT_INCLUDE_DIR} ${TENSORRT_PYTHON_INCLUDE_DIR} ${CUDA_INCLUDE_DIR})
135141
target_link_libraries(${nvonnxparser_lib_name_static} PUBLIC onnx_proto ${PROTOBUF_LIBRARY})
136142
set_target_properties(${nvonnxparser_lib_name_static} PROPERTIES
137143
ARCHIVE_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
@@ -143,7 +149,7 @@ set_target_properties(${nvonnxparser_lib_name_static} PROPERTIES
143149
# --------------------------------
144150
if(BUILD_ONNXIFI)
145151
add_library(trt_onnxify SHARED ${ONNXIFI_SOURCES})
146-
target_include_directories(trt_onnxify PUBLIC ${CUDA_INCLUDE_DIR} ${ONNX_INCLUDE_DIRS} ${TENSORRT_INCLUDE_DIR})
152+
target_include_directories(trt_onnxify PUBLIC ${CUDA_INCLUDE_DIR} ${ONNX_INCLUDE_DIRS} ${TENSORRT_INCLUDE_DIR} ${TENSORRT_PYTHON_INCLUDE_DIR})
147153
target_link_libraries(trt_onnxify PUBLIC ${nvonnxparser_lib_name_static} ${CMAKE_THREAD_LIBS_INIT} ${CMAKE_DL_LIBS})
148154
endif()
149155

Diff for: ‎ConditionalHelpers.cpp

+15-95
Original file line numberDiff line numberDiff line change
@@ -10,32 +10,17 @@
1010
namespace onnx2trt
1111
{
1212

13-
using NodeName = std::string;
14-
using LayerName = std::string;
15-
using InputIndex = int32_t;
16-
17-
// A SubgraphPortsMap maps either the inputs or outputs ports of each node in an ONNX graph.
18-
using SubgraphPortsMap = std::unordered_map<nvinfer1::ITensor*, std::set<InputIndex>>;
19-
20-
// An InputsMap tracks which IIfConditionalInputLayer we've added to a layer's inputs,
21-
// so that we can reuse them if needed.
22-
using InputsMap = std::unordered_map<LayerName, nvinfer1::IIfConditionalInputLayer*>;
23-
24-
// Search for a network Layer name in a SubgraphPortsMap using partial (prefix) name matching.
25-
// ONNX nodes are matched to network layers using prefix-matching because an ONNX node may have
26-
// several network layers associcated with it.
13+
// Search for a network Layer name in a SubgraphPortsMap.
2714
SubgraphPortsMap::const_iterator findLayer(const SubgraphPortsMap& inputs, const std::string layerName)
2815
{
29-
return std::find_if(inputs.begin(), inputs.end(), [&](const auto& item) {
30-
std::string const key = item.first->getName();
31-
return layerName.compare(0, key.size(), key) == 0;
32-
});
16+
return std::find_if(
17+
inputs.begin(), inputs.end(), [&](const auto& item) { return layerName == item.first->getName(); });
3318
}
3419

3520
// Add an ConditionalInputLayer between `layer` and its inputs.
3621
// I.e. input[inIdx] -> layer ==> input[inIdx] -> ConditionalInputLayer -> layer.
3722
void addConditionalInputLayer(ImporterContext* ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
38-
nvinfer1::ILayer& layer, int32_t inIdx)
23+
nvinfer1::ILayer& layer, int32_t inIdx, ::ONNX_NAMESPACE::NodeProto const* node)
3924
{
4025
auto input = layer.getInput(inIdx);
4126
if (input == nullptr)
@@ -57,7 +42,7 @@ void addConditionalInputLayer(ImporterContext* ctx, nvinfer1::IIfConditional* co
5742
inputLayer = N_CHECK(conditional->addInput(*input));
5843
inputsMap[name] = inputLayer;
5944
const std::string inputLayerName(name);
60-
ctx->registerLayer(inputLayer, inputLayerName + "_InputLayer", nullptr);
45+
ctx->registerLayer(inputLayer, inputLayerName + "_InputLayer", node);
6146
// Note: Since multiple conditionals may use the same external tensor, check unique names for output tensors of
6247
// IfConditionalInputLayers to avoid tensor name duplication.
6348
ctx->registerTensor(
@@ -100,7 +85,7 @@ void importSubgraph(ImporterContext* ctx, ::ONNX_NAMESPACE::GraphProto const& su
10085

10186
// Add an IConditionalInputLayer to `layer`'s inputs, if they don't already exist.
10287
void addConditionalInputIfNeeded(ImporterContext* ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
103-
nvinfer1::ILayer& layer, SubgraphPortsMap subgraphInputsMap)
88+
nvinfer1::ILayer& layer, SubgraphPortsMap subgraphInputsMap, ::ONNX_NAMESPACE::NodeProto const* node)
10489
{
10590
// Return all of the layer's inputs that are external to the subgraph that
10691
// that the layer belongs to.
@@ -120,36 +105,29 @@ void addConditionalInputIfNeeded(ImporterContext* ctx, nvinfer1::IIfConditional*
120105
for (auto inIdx : inIndices)
121106
{
122107
LOG_VERBOSE("Adding Input layer for " << layer.getName());
123-
addConditionalInputLayer(ctx, conditional, inputsMap, layer, inIdx);
108+
addConditionalInputLayer(ctx, conditional, inputsMap, layer, inIdx, node);
124109
}
125110
}
126111

127112
// Add IConditionalInputLayers to `layer`'s inputs.
128113
void addIfInputLayers(ImporterContext* ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
129-
const std::vector<nvinfer1::ILayer*>& newLayers)
114+
const std::vector<nvinfer1::ILayer*>& newLayers, ::ONNX_NAMESPACE::NodeProto const* node)
130115
{
131116
// Find all of the tensors entering the subgraph.
132-
// The node-names are from the ONNX context.
133-
using InputIndex = int32_t;
134-
std::unordered_map<nvinfer1::ITensor*, std::set<InputIndex>> subgraphInputsMap;
135-
getSubgraphInputs(newLayers, subgraphInputsMap);
117+
SubgraphPortsMap externalInputs;
118+
getSubgraphInputs(newLayers, externalInputs);
136119

137120
// Add a ConditionalInputLayer in front of each input that is external to the subgraph.
138121
for (const auto& layer : newLayers)
139122
{
140-
addConditionalInputIfNeeded(ctx, conditional, inputsMap, *layer, subgraphInputsMap);
123+
addConditionalInputIfNeeded(ctx, conditional, inputsMap, *layer, externalInputs, node);
141124
}
142125
}
143126

144-
// Given a subgraph, find all of its external inputs/outputs (tensors entering/exiting the subgraph).
145-
void getSubgraphTensors(const std::vector<nvinfer1::ILayer*>& newLayers,
146-
std::unordered_map<nvinfer1::ITensor*, std::set<int32_t>>& externalOutputs, bool extractOutputs,
147-
const std::vector<std::string>* reportedOutputs = nullptr)
127+
// Given a subgraph, find all of its external inputs (tensors entering the subgraph).
128+
void getSubgraphInputs(const std::vector<nvinfer1::ILayer*>& newLayers, SubgraphPortsMap& externalInputs)
148129
{
149-
using NodeName = std::string;
150-
using TensorName = std::string;
151130
using PortIndex = int32_t;
152-
using Port = std::pair<NodeName, PortIndex>;
153131
using TensorsSet = std::unordered_set<nvinfer1::ITensor*>;
154132
TensorsSet outputTensors;
155133
TensorsSet inputTensors;
@@ -176,20 +154,14 @@ void getSubgraphTensors(const std::vector<nvinfer1::ILayer*>& newLayers,
176154
}
177155

178156
using TensorsVec = std::vector<nvinfer1::ITensor*>;
179-
auto getOutputs = [&](nvinfer1::ILayer const* l, TensorsVec& res) {
180-
getTensors(l, false, [&](nvinfer1::ITensor* t) { res.emplace_back(t); });
181-
};
182-
183157
auto getInputs = [&](nvinfer1::ILayer const* l, TensorsVec& res) {
184158
getTensors(l, true, [&](nvinfer1::ITensor* t) { res.emplace_back(t); });
185159
};
186160

187161
// Retrieve the list of tensors either exiting or entering the subgraph.
188-
std::unordered_map<nvinfer1::ITensor*, std::vector<Port>> externalPortsMap;
189162
auto filterTensors = [&](TensorsSet const& tensors, auto getNodeAccessor) {
190163
for (nvinfer1::ILayer const* l : newLayers)
191164
{
192-
const auto& nodeName = l->getName();
193165
PortIndex i = 0;
194166

195167
TensorsVec nodeAccessor;
@@ -202,66 +174,14 @@ void getSubgraphTensors(const std::vector<nvinfer1::ILayer*>& newLayers,
202174
}
203175
if (tensors.count(tensor) == 0)
204176
{
205-
TensorName tensorName = tensor->getName();
206-
auto prefixFound = false;
207-
if (reportedOutputs)
208-
{
209-
// reportedOutputs are the names of the outputs as reported by the ONNX parser and help
210-
// us further filter the output tensors.
211-
// Exiting tensors := {outputs} - {inputs} - {unreported tensors}
212-
// An example: a Split node is internal to a subgraph and has 4 outputs, but only two are
213-
// connected to the rest of the graph. To prevent mistaking the 2 unused outputs as subgraph
214-
// outputs, we look for them in reportedOutputs which leads us to ignore the 2 tensors.
215-
const auto iter = std::find_if(
216-
reportedOutputs->begin(), reportedOutputs->end(), [&](const auto& outputName) {
217-
// Prefix name matching.
218-
return tensorName.compare(0, outputName.size(), outputName) == 0;
219-
});
220-
prefixFound = iter != reportedOutputs->end();
221-
}
222-
if (!reportedOutputs || prefixFound)
223-
{
224-
externalPortsMap[tensor].push_back(std::make_pair(nodeName, i));
225-
}
177+
externalInputs[l].insert(i);
226178
}
227179
i++;
228180
}
229181
}
230182
};
231183

232-
if (extractOutputs)
233-
{
234-
filterTensors(inputTensors, getOutputs);
235-
}
236-
else
237-
{
238-
filterTensors(outputTensors, getInputs);
239-
}
240-
241-
// Create the user's view of the external inputs, which uses the node-name as the key for
242-
// looking up input/output port index.
243-
for (auto const& input : externalPortsMap)
244-
{
245-
for (const Port& inPort : input.second)
246-
{
247-
auto* tensor = input.first;
248-
auto const portIndex = inPort.second;
249-
externalOutputs[tensor].insert(portIndex);
250-
}
251-
}
252-
}
253-
254-
void getSubgraphOutputs(const std::vector<nvinfer1::ILayer*>& newLayers,
255-
std::unordered_map<nvinfer1::ITensor*, std::set<int32_t>>& externalOutputs,
256-
const std::vector<std::string>& reportedOutputs)
257-
{
258-
getSubgraphTensors(newLayers, externalOutputs, true, &reportedOutputs);
259-
}
260-
261-
void getSubgraphInputs(const std::vector<nvinfer1::ILayer*>& newLayers,
262-
std::unordered_map<nvinfer1::ITensor*, std::set<int32_t>>& externalInputs)
263-
{
264-
getSubgraphTensors(newLayers, externalInputs, false);
184+
filterTensors(outputTensors, getInputs);
265185
}
266186

267187
} // namespace onnx2trt

Diff for: ‎ConditionalHelpers.hpp

+12-13
Original file line numberDiff line numberDiff line change
@@ -18,28 +18,27 @@
1818
namespace onnx2trt
1919
{
2020

21+
using NodeName = std::string;
22+
using LayerName = std::string;
23+
using InputIndex = int32_t;
24+
25+
// A SubgraphPortsMap maps inputs' ports of each layer in an ONNX graph.
26+
using SubgraphPortsMap = std::unordered_map<const nvinfer1::ILayer*, std::unordered_set<InputIndex>>;
27+
2128
// Given a subgraph, find all of its external inputs (tensors entering the subgraph).
22-
// The result is returned in `subgraphInputs`, which is a map indexed by ITensor (a tensor entering the subgraph) and
23-
// with values indicating a set of external input indices.
24-
void getSubgraphInputs(std::vector<nvinfer1::ILayer*> const& newLayers,
25-
std::unordered_map<nvinfer1::ITensor*, std::set<int32_t>>& subgraphInputs);
26-
27-
// Given a subgraph, find all of its external outputs (tensors exiting the subgraph).
28-
// The result is returned in `subgraphInputs`, which is a map indexed by ITensor (a tensor exiting the subgraph) and
29-
// with values indicating a set of external outputs indices.
30-
void getSubgraphOutputs(const std::vector<nvinfer1::ILayer*>& newLayers,
31-
std::unordered_map<nvinfer1::ITensor*, std::set<int32_t>>& subgraphOutputs,
32-
const std::vector<std::string>& reportedOutputs);
29+
void getSubgraphInputs(const std::vector<nvinfer1::ILayer*>& newLayers, SubgraphPortsMap& externalInputs);
3330

3431
// Take a snapshot of the network before and after parsing the subgraph and return a list
3532
// of newly added network layers.
3633
void importSubgraph(ImporterContext* ctx, ::ONNX_NAMESPACE::GraphProto const& subgraph,
3734
std::vector<nvinfer1::ILayer*>& newLayers, std::vector<TensorOrWeights>& subgraphTensors);
3835

39-
using InputsMap = std::unordered_map<std::string, nvinfer1::IIfConditionalInputLayer*>;
36+
// An InputsMap tracks which IIfConditionalInputLayer we've added to a layer's inputs,
37+
// so that we can reuse them if needed.
38+
using InputsMap = std::unordered_map<LayerName, nvinfer1::IIfConditionalInputLayer*>;
4039

4140
// Add IIfConditionalInputLayers to the inputs of the subgraph indicated by `subgraph`.
4241
void addIfInputLayers(ImporterContext* ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
43-
const std::vector<nvinfer1::ILayer*>& newLayers);
42+
const std::vector<nvinfer1::ILayer*>& newLayers, ::ONNX_NAMESPACE::NodeProto const* node);
4443

4544
} // namespace onnx2trt

Diff for: ‎ImporterContext.hpp

+11-2
Original file line numberDiff line numberDiff line change
@@ -138,8 +138,17 @@ class ImporterContext
138138
//! Map holding FunctionProtos
139139
StringMap<::ONNX_NAMESPACE::FunctionProto> mLocalFunctions;
140140

141+
//! Data type to keep track of a local function in mLocalFunctionStack.
142+
//! It is a tuple of three elements: (1) function name (2) node name and (3) function attributes.
143+
struct LocalFunctionMetadata
144+
{
145+
std::string functionName;
146+
std::string nodeName;
147+
StringMap<::ONNX_NAMESPACE::AttributeProto const*> attrs;
148+
};
149+
141150
//! Vector to hold current local function names and attributes
142-
std::vector<std::pair<std::string, StringMap<::ONNX_NAMESPACE::AttributeProto const*>>> mLocalFunctionStack;
151+
std::vector<LocalFunctionMetadata> mLocalFunctionStack;
143152

144153
//! Vector to hold the local function names at each error
145154
std::vector<std::vector<std::string>> mLocalFunctionErrors;
@@ -325,7 +334,7 @@ class ImporterContext
325334
{
326335
return mLocalFunctions;
327336
}
328-
std::vector<std::pair<std::string, StringMap<::ONNX_NAMESPACE::AttributeProto const*>>>& localFunctionStack()
337+
std::vector<LocalFunctionMetadata>& localFunctionStack()
329338
{
330339
return mLocalFunctionStack;
331340
}

Diff for: ‎ModelImporter.cpp

+14-5
Original file line numberDiff line numberDiff line change
@@ -221,19 +221,27 @@ void parseNode(
221221
}
222222
}
223223

224-
ONNXTRT_CHECK_NODE((node.output().size() <= static_cast<int32_t>(outputs.size())),
224+
int32_t nonEmptyOutputs
225+
= std::count_if(node.output().begin(), node.output().end(), [](std::string const& str) { return !str.empty(); });
226+
ONNXTRT_CHECK_NODE(nonEmptyOutputs == static_cast<int32_t>(outputs.size()),
225227
"Node has more output tensors than TRT expected, expected output size is "
226-
<< outputs.size() << ", actual output size is " << node.output().size() << ".",
228+
<< outputs.size() << ", actual output size is " << nonEmptyOutputs << ".",
227229
node, nodeIdx, ErrorCode::kINVALID_GRAPH);
228230

229231
// Set output names and register outputs with the context.
230232
std::ostringstream ssOutputs{};
231233
ssOutputs << nodeName << " [" << node.op_type() << "] outputs: ";
232-
for (int32_t i = 0; i < node.output().size(); ++i)
234+
for (int32_t i = 0, trtCnt = 0; i < node.output().size(); ++i)
233235
{
234236
auto const& outputName = node.output(i);
235-
auto& output = outputs.at(i);
236-
ssOutputs << "[" << outputName << " -> " << output.shape() << "[" << output.getType() << "]" << "], ";
237+
// Empty strings denote null-tensor outputs. Ignore these.
238+
if (outputName.empty())
239+
{
240+
continue;
241+
}
242+
auto& output = outputs.at(trtCnt);
243+
ssOutputs << "[" << outputName << " -> " << output.shape() << "[" << output.getType() << "]"
244+
<< "], ";
237245
// Note: This condition is to allow ONNX outputs to be ignored
238246
// Always register output weights (even empty ones) as it may be mapped to an unused input
239247
if ((output || output.is_weights()) && !outputName.empty())
@@ -255,6 +263,7 @@ void parseNode(
255263
ONNXTRT_CHECK_NODE(legalUINT8, "TensorRT does not support UINT8 types for intermediate tensors!", node,
256264
nodeIdx, ErrorCode::kUNSUPPORTED_NODE);
257265
}
266+
trtCnt++;
258267
}
259268
LOG_VERBOSE(ssOutputs.str());
260269
}

0 commit comments

Comments
 (0)