未验证 提交 9bd20aa0 编写于 作者: Z zlsh80826 提交者: GitHub

Add a persistent ibuilder to speedup unit test (#48906)

上级 a9ed6f96
...@@ -852,6 +852,15 @@ class TRTEngineManager { ...@@ -852,6 +852,15 @@ class TRTEngineManager {
size_t max_ctx_mem_size_{0}; size_t max_ctx_mem_size_{0};
std::unordered_map<PredictorID, AllocationPtr> context_memorys_; std::unordered_map<PredictorID, AllocationPtr> context_memorys_;
std::unordered_map<std::string, std::unique_ptr<TensorRTEngine>> engines_; std::unordered_map<std::string, std::unique_ptr<TensorRTEngine>> engines_;
// createInferBuilder loads trt kernels and take a few second
// But as long as one IBuilder lives, trt kernel will not be unloaded
// Hence, a persistent IBuilder to avoid TensorRT unload/reload kernels
std::unique_ptr<nvinfer1::IBuilder, std::function<void(nvinfer1::IBuilder*)>>
holder{createInferBuilder(&NaiveLogger::Global()), [](auto* ptr) {
if (ptr) {
ptr->destroy();
}
}};
}; };
} // namespace tensorrt } // namespace tensorrt
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册