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

Fix TensorRT memory leaks (#4227)

Browse source 
* fix tensorrt memory leaks

* wrap unique_pointer in a namespace to avoid conflicts

Co-authored-by: alex <act@act.com>
This commit is contained in:
alkoumpa 2020-06-20 13:37:38 +03:00 committed by GitHub
parent a541d28fb4
commit 3c633384c2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 45 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -464,11 +464,11 @@ SubGraphCollection_t TensorrtExecutionProvider::GetSupportedList(SubGraphCollect
// Get supported node list recursively
SubGraphCollection_t parser_nodes_list;
TensorrtLogger& trt_logger = GetTensorrtLogger();
auto trt_builder = unique_pointer<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(trt_logger));
auto trt_builder = tensorrt_ptr::unique_pointer<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(trt_logger));
const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
auto trt_network = unique_pointer<nvinfer1::INetworkDefinition>(trt_builder->createNetworkV2(explicitBatch));
auto trt_network = tensorrt_ptr::unique_pointer<nvinfer1::INetworkDefinition>(trt_builder->createNetworkV2(explicitBatch));
auto trt_parser = unique_pointer<nvonnxparser::IParser>(nvonnxparser::createParser(*trt_network, trt_logger));
auto trt_parser = tensorrt_ptr::unique_pointer<nvonnxparser::IParser>(nvonnxparser::createParser(*trt_network, trt_logger));
trt_parser->supportsModel(string_buf.data(), string_buf.size(), parser_nodes_list);
SubGraphCollection_t next_nodes_list;
@ -682,11 +682,11 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
// Create TensorRT engine
TensorrtLogger& trt_logger = GetTensorrtLogger();
auto trt_builder = unique_pointer<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(trt_logger));
auto trt_builder = tensorrt_ptr::unique_pointer<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(trt_logger));
const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
auto trt_network = unique_pointer<nvinfer1::INetworkDefinition>(trt_builder->createNetworkV2(explicitBatch));
auto trt_config = unique_pointer<nvinfer1::IBuilderConfig>(trt_builder->createBuilderConfig());
auto trt_parser = unique_pointer<nvonnxparser::IParser>(nvonnxparser::createParser(*trt_network, trt_logger));
auto trt_network = tensorrt_ptr::unique_pointer<nvinfer1::INetworkDefinition>(trt_builder->createNetworkV2(explicitBatch));
auto trt_config = tensorrt_ptr::unique_pointer<nvinfer1::IBuilderConfig>(trt_builder->createBuilderConfig());
auto trt_parser = tensorrt_ptr::unique_pointer<nvonnxparser::IParser>(nvonnxparser::createParser(*trt_network, trt_logger));
trt_parser->parse(string_buf.data(), string_buf.size());
trt_config->setMaxWorkspaceSize(max_workspace_size_);
@ -734,14 +734,14 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
trt_config->setFlag(nvinfer1::BuilderFlag::kFP16);
}
auto trt_engine = unique_pointer<nvinfer1::ICudaEngine>(trt_builder->buildEngineWithConfig(*trt_network, *trt_config));
auto trt_engine = tensorrt_ptr::unique_pointer<nvinfer1::ICudaEngine>(trt_builder->buildEngineWithConfig(*trt_network, *trt_config));
if (trt_engine == nullptr) {
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL,
"TensorRT EP could not build Engine for fused node: " + fused_node->Name());
}
// Build TensorRT context
auto trt_context = unique_pointer<nvinfer1::IExecutionContext>(trt_engine->createExecutionContext());
auto trt_context = tensorrt_ptr::unique_pointer<nvinfer1::IExecutionContext>(trt_engine->createExecutionContext());
if (trt_context == nullptr) {
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL,
"TensorRT EP could not build Execution Context for fused node: " + fused_node->Name());
@ -815,7 +815,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
compute_info.create_state_func = [=](ComputeContext* context, FunctionState* state) {
std::unique_ptr<TensorrtFuncState> p = onnxruntime::make_unique<TensorrtFuncState>();
*p = {context->allocate_func, context->release_func, context->allocator_handle, parsers_[context->node_name].get(),
engines_[context->node_name].get(), contexts_[context->node_name].get(), builders_[context->node_name].get(),
&engines_[context->node_name], &contexts_[context->node_name], builders_[context->node_name].get(),
networks_[context->node_name].get(), input_info_[context->node_name], output_info_[context->node_name],
input_shape_ranges_[context->node_name], output_shapes_[context->node_name], &tensorrt_mu_, &fp16_enable_,
&max_workspace_size_};
@ -845,7 +845,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
// Update shape ranges
bool dimension_update = false;
auto trt_context = trt_state->context;
auto trt_context = trt_state->context->get();
auto trt_builder = trt_state->builder;
nvinfer1::IOptimizationProfile* trt_profile = nullptr;
for (int i = 0, end = num_binding_inputs; i < end; ++i) {
@ -913,21 +913,26 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
// Regenerate engine and context
// Only one profile is generated, so no need to explicitly set optimization profile
if (dimension_update) {
auto trt_config = unique_pointer<nvinfer1::IBuilderConfig>(trt_builder->createBuilderConfig());
auto trt_config = tensorrt_ptr::unique_pointer<nvinfer1::IBuilderConfig>(trt_builder->createBuilderConfig());
trt_config->setMaxWorkspaceSize(*(trt_state->max_workspace_size_ptr));
trt_config->addOptimizationProfile(trt_profile);
if (*(trt_state->fp16_enable_ptr) && trt_builder->platformHasFastFp16()) {
trt_config->setFlag(nvinfer1::BuilderFlag::kFP16);
}
trt_state->engine = trt_builder->buildEngineWithConfig(*trt_state->network, *trt_config);
if (trt_state->engine == nullptr) {
trt_state->context->reset();
trt_state->engine->reset();
*(trt_state->engine) = tensorrt_ptr::unique_pointer<nvinfer1::ICudaEngine>(
trt_builder->buildEngineWithConfig(*trt_state->network, *trt_config));
if (trt_state->engine->get() == nullptr) {
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Build Engine.");
}
trt_state->context = trt_state->engine->createExecutionContext();
if (trt_state->context == nullptr) {
*(trt_state->context) = tensorrt_ptr::unique_pointer<nvinfer1::IExecutionContext>(
trt_state->engine->get()->createExecutionContext());
if (trt_state->context->get() == nullptr) {
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Create Context.");
}
trt_context = trt_state->context;
trt_context = trt_state->context->get();
}
// Set input shapes and assign input buffers

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

@ -38,6 +38,21 @@ class TensorrtLogger : public nvinfer1::ILogger {
}
};
namespace tensorrt_ptr {
struct TensorrtInferDeleter {
template <typename T>
void operator()(T* obj) const {
if (obj) {
obj->destroy();
}
}
};
template <typename T>
using unique_pointer = std::unique_ptr<T, TensorrtInferDeleter>;
};
// Information needed to construct trt execution providers.
struct TensorrtExecutionProviderInfo {
int device_id{0};
@ -45,12 +60,13 @@ struct TensorrtExecutionProviderInfo {
// Information to construct kernel function state.
struct TensorrtFuncState {
AllocateFunc test_allocate_func = nullptr;
DestroyFunc test_release_func = nullptr;
AllocatorHandle allocator = nullptr;
nvonnxparser::IParser* parser = nullptr;
nvinfer1::ICudaEngine* engine = nullptr;
nvinfer1::IExecutionContext* context = nullptr;
tensorrt_ptr::unique_pointer<nvinfer1::ICudaEngine> * engine = nullptr;
tensorrt_ptr::unique_pointer<nvinfer1::IExecutionContext> * context = nullptr;
nvinfer1::IBuilder* builder = nullptr;
nvinfer1::INetworkDefinition* network = nullptr;
std::vector<std::vector<int>> input_info;
@ -89,25 +105,14 @@ class TensorrtExecutionProvider : public IExecutionProvider {
bool fp16_enable_ = false;
bool dump_subgraphs_ = false;
struct InferDeleter {
template <typename T>
void operator()(T* obj) const {
if (obj) {
obj->destroy();
}
}
};
template <typename T>
using unique_pointer = std::unique_ptr<T, InferDeleter>;
OrtMutex tensorrt_mu_;
int device_id_;
std::unordered_map<std::string, unique_pointer<nvonnxparser::IParser>> parsers_;
std::unordered_map<std::string, unique_pointer<nvinfer1::ICudaEngine>> engines_;
std::unordered_map<std::string, unique_pointer<nvinfer1::IExecutionContext>> contexts_;
std::unordered_map<std::string, unique_pointer<nvinfer1::IBuilder>> builders_;
std::unordered_map<std::string, unique_pointer<nvinfer1::INetworkDefinition>> networks_;
std::unordered_map<std::string, tensorrt_ptr::unique_pointer<nvonnxparser::IParser>> parsers_;
std::unordered_map<std::string, tensorrt_ptr::unique_pointer<nvinfer1::ICudaEngine>> engines_;
std::unordered_map<std::string, tensorrt_ptr::unique_pointer<nvinfer1::IExecutionContext>> contexts_;
std::unordered_map<std::string, tensorrt_ptr::unique_pointer<nvinfer1::IBuilder>> builders_;
std::unordered_map<std::string, tensorrt_ptr::unique_pointer<nvinfer1::INetworkDefinition>> networks_;
std::unordered_map<std::string, std::vector<std::vector<int>>> input_info_;
std::unordered_map<std::string, std::vector<std::vector<int>>> output_info_;
std::unordered_map<std::string, std::unordered_map<int, std::unordered_map<int, std::pair<int64_t, int64_t>>>> input_shape_ranges_;