[incubator-mxnet] branch v1.x updated: Add TRT verbose mode (#19100)

[ptrendx]

This is an automated email from the ASF dual-hosted git repository.

ptrendx pushed a commit to branch v1.x
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git


The following commit(s) were added to refs/heads/v1.x by this push:
     new 2d077db  Add TRT verbose mode (#19100)
2d077db is described below

commit 2d077db1c92dbf7db979e15609fddf5f371277c0
Author: Serge Panev <spa...@nvidia.com>
AuthorDate: Thu Sep 10 13:12:44 2020 -0700

    Add TRT verbose mode (#19100)
    
    Signed-off-by: Serge Panev <spa...@nvidia.com>
---
 src/operator/subgraph/tensorrt/tensorrt.cc | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/operator/subgraph/tensorrt/tensorrt.cc 
b/src/operator/subgraph/tensorrt/tensorrt.cc
index e872e92..8395fb4 100644
--- a/src/operator/subgraph/tensorrt/tensorrt.cc
+++ b/src/operator/subgraph/tensorrt/tensorrt.cc
@@ -28,6 +28,8 @@
 
 #include "./tensorrt-inl.h"
 
+#include <NvInfer.h>
+
 namespace mxnet {
 namespace op {
 
@@ -311,7 +313,13 @@ OpStatePtr TRTCreateState(const nnvm::NodeAttrs& attrs, 
Context ctx,
   graph.attrs["dtype"]        = std::make_shared<nnvm::any>(std::move(dtypes));
   graph.attrs["shape"]        = std::make_shared<nnvm::any>(std::move(shapes));
   auto onnx_graph = op::nnvm_to_onnx::ConvertNnvmGraphToOnnx(graph, 
&params_map);
-  auto trt_tuple = ::onnx_to_tensorrt::onnxToTrtCtx(onnx_graph, 
max_batch_size, 1 << 30);
+  uint32_t verbose = dmlc::GetEnv("MXNET_TENSORRT_VERBOSE", 0);
+  auto log_lvl = nvinfer1::ILogger::Severity::kWARNING;
+  if (verbose != 0) {
+    log_lvl = nvinfer1::ILogger::Severity::kVERBOSE;
+  }
+
+  auto trt_tuple = ::onnx_to_tensorrt::onnxToTrtCtx(onnx_graph, 
max_batch_size, 1 << 30, log_lvl);
   return OpStatePtr::Create<TRTEngineParam>(std::move(std::get<0>(trt_tuple)),
                                             std::move(std::get<1>(trt_tuple)),
                                             std::move(std::get<2>(trt_tuple)),