saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-22 19:23:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

debug option for dumping tensorrt subgraphs. (#3604)

Browse source
This commit is contained in:
George Wu 2020-04-21 11:55:30 +08:00 committed by GitHub
parent e233e6ba45
commit 1c37d5e6ec
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 19 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
View file

@ -1,6 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <fstream>
#include "core/graph/onnx_protobuf.h"
#include "tensorrt_execution_provider.h"
@ -124,6 +125,11 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv
if (!fp16_enable_env.empty()) {
fp16_enable_ = (std::stoi(fp16_enable_env) == 0 ? false : true);
}
const std::string dump_subgraphs_env = env_instance.GetEnvironmentVar(tensorrt_env_vars::kDumpSubgraphs);
if (!dump_subgraphs_env.empty()) {
dump_subgraphs_ = (std::stoi(dump_subgraphs_env) == 0 ? false : true);
}
}
TensorrtExecutionProvider::~TensorrtExecutionProvider() {}
@ -252,7 +258,7 @@ std::unique_ptr<IndexedSubGraph> TensorrtExecutionProvider::GetSubGraph(SubGraph
// Find inputs and outputs of the subgraph
std::unique_ptr<IndexedSubGraph> sub_graph = onnxruntime::make_unique<IndexedSubGraph>();
std::unordered_map<const NodeArg *, int> fused_inputs, fused_outputs, fused_outputs_to_add, graph_outputs_to_add;
std::unordered_map<const NodeArg*, int> fused_inputs, fused_outputs, fused_outputs_to_add, graph_outputs_to_add;
std::unordered_set<const NodeArg*> erased;
int input_order = 0;
int output_order = 0;
@ -318,7 +324,7 @@ std::unique_ptr<IndexedSubGraph> TensorrtExecutionProvider::GetSubGraph(SubGraph
fused_outputs.insert(graph_outputs_to_add.begin(), graph_outputs_to_add.end());
// Sort inputs and outputs by the order they were added
std::multimap<int, const NodeArg *> inputs, outputs;
std::multimap<int, const NodeArg*> inputs, outputs;
for (auto it = fused_inputs.begin(), end = fused_inputs.end(); it != end; ++it) {
inputs.insert(std::pair<int, const NodeArg*>(it->second, it->first));
}
@ -384,7 +390,7 @@ SubGraphCollection_t TensorrtExecutionProvider::GetSupportedList(SubGraphCollect
std::vector<std::string> subgraph_output_names;
for (const auto& index : group.first) {
const auto& node = graph.GetNode(node_index[index]);
std::vector<onnxruntime::NodeArg *> inputs, outputs;
std::vector<onnxruntime::NodeArg*> inputs, outputs;
for (auto input : node->InputDefs()) {
auto& n_input = graph_build.GetOrCreateNodeArg(input->Name(), input->TypeAsProto());
inputs.push_back(&n_input);
@ -661,6 +667,12 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
std::string string_buf;
model_proto.SerializeToString(&string_buf);
if (dump_subgraphs_) {
// Dump the TensorRT subgraph if enabled via ORT_TENSORRT_DUMP_SUBGRAPHS env variable.
std::fstream dump(fused_node->Name() + ".onnx", std::ios::out | std::ios::trunc | std::ios::binary);
model_proto.SerializeToOstream(&dump);
}
// Create TensorRT engine
TensorrtLogger& trt_logger = GetTensorrtLogger();
auto trt_builder = unique_pointer<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(trt_logger));
@ -1009,4 +1021,4 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
return Status::OK();
}
} // namespace onnxruntime
} // namespace onnxruntime

4
onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h
View file

@ -16,6 +16,7 @@ static const std::string kMaxPartitionIterations = "ORT_TENSORRT_MAX_PARTITION_I
static const std::string kMinSubgraphSize = "ORT_TENSORRT_MIN_SUBGRAPH_SIZE";
static const std::string kMaxWorkspaceSize = "ORT_TENSORRT_MAX_WORKSPACE_SIZE";
static const std::string kFP16Enable = "ORT_TENSORRT_FP16_ENABLE";
static const std::string kDumpSubgraphs = "ORT_TENSORRT_DUMP_SUBGRAPHS";
} // namespace tensorrt_env_vars
class TensorrtLogger : public nvinfer1::ILogger {
@ -86,6 +87,7 @@ class TensorrtExecutionProvider : public IExecutionProvider {
int max_partition_iterations_ = 1000;
int min_subgraph_size_ = 1;
bool fp16_enable_ = false;
bool dump_subgraphs_ = false;
struct InferDeleter {
template <typename T>
@ -126,7 +128,7 @@ class TensorrtExecutionProvider : public IExecutionProvider {
const onnxruntime::GraphViewer& graph, bool* early_termination) const;
void RemoveTensorRTGraphCycles(SubGraphCollection_t& supported_nodes_vector, const onnxruntime::GraphViewer& graph) const;
AllocatorPtr allocator_;
};