Fix memory leak in TRT (#2815)

* fix memory leak issue

* revert EP_FAIL on enueueV2
This commit is contained in:
stevenlix 2020-01-10 14:07:40 -08:00 committed by Ryan Hill
parent b8114cf3ed
commit 4d8fca0a4e

View file

@ -676,7 +676,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
// Update shape ranges
bool dimension_update = false;
auto trt_builder = trt_state->builder;
auto trt_profile = trt_builder->createOptimizationProfile();
nvinfer1::IOptimizationProfile* trt_profile = nullptr;
for (int i = 0, end = num_binding_inputs; i < end; ++i) {
// TODO: check if getInput indexing is same with binding index
auto input = trt_state->network->getInput(i);
@ -714,6 +714,9 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
}
if (dimension_update) {
if (trt_profile == nullptr) {
trt_profile = trt_builder->createOptimizationProfile();
}
if (engine.isShapeBinding(i)) {
std::vector<int32_t> shapes_min(nb_dims), shapes_opt(nb_dims), shapes_max(nb_dims);
for (int j = 0, end = nb_dims; j < end; ++j) {
@ -733,7 +736,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
}
// TensorRT6 requires optimization profile to be defined for all inputs if any input dimension is symbolic
if (dynamic_shape) {
if (dimension_update && dynamic_shape) {
trt_profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMIN, dims_min);
trt_profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kOPT, dims_opt);
trt_profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, dims_max);
@ -746,10 +749,13 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
auto trt_config = unique_pointer<nvinfer1::IBuilderConfig>(trt_builder->createBuilderConfig());
trt_config->addOptimizationProfile(trt_profile);
trt_state->engine = trt_builder->buildEngineWithConfig(*trt_state->network, *trt_config);
ORT_ENFORCE(trt_state->engine != nullptr);
if (trt_state->engine == nullptr) {
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Build Engine.");
}
trt_state->context = trt_state->engine->createExecutionContext();
ORT_ENFORCE(trt_state->context != nullptr);
if (trt_state->context == nullptr) {
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Create Context.");
}
trt_context = trt_state->context;
}