fix(tensorrt): update tensorrt code of tensorrt_yolo

Signed-off-by: M. Fatih Cırıt <mfc@leodrive.ai>
autowarefoundation · Dec 21, 2022 · 1d1904c · 1d1904c
1 parent 9376697
commit 1d1904c
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 4 deletions.
diff --git a/perception/tensorrt_yolo/lib/include/trt_yolo.hpp b/perception/tensorrt_yolo/lib/include/trt_yolo.hpp
@@ -144,7 +144,7 @@ class Net
   std::vector<float> preprocess(
     const cv::Mat & in_img, const int c, const int h, const int w) const;
   // Infer using pre-allocated GPU buffers {data, scores, boxes}
-  void infer(std::vector<void *> & buffers, const int batch_size);
+  void infer([[maybe_unused]] std::vector<void *> & buffers, const int batch_size);
 };
 
 bool set_cuda_device(int gpu_id)

diff --git a/perception/tensorrt_yolo/lib/src/trt_yolo.cpp b/perception/tensorrt_yolo/lib/src/trt_yolo.cpp
@@ -276,15 +276,26 @@ void Net::save(const std::string & path) const
   file.write(reinterpret_cast<const char *>(plan_->data()), plan_->size());
 }
 
-void Net::infer(std::vector<void *> & buffers, const int batch_size)
+void Net::infer([[maybe_unused]] std::vector<void *> & buffers, const int batch_size)
 {
   if (!context_) {
     throw std::runtime_error("Fail to create context");
   }
+
+#if (NV_TENSORRT_MAJOR * 10000) + (NV_TENSORRT_MINOR * 100) + NV_TENSOR_PATCH >= 80500
+  const auto input_dims = engine_->getTensorShape(engine_->getIOTensorName(0));
+  context_->setInputShape(
+    engine_->getIOTensorName(0),
+    nvinfer1::Dims4(batch_size, input_dims.d[1], input_dims.d[2], input_dims.d[3]));
+  context_->enqueueV3(stream_);
+#else
+  // Deprecated since 8.5
   auto input_dims = engine_->getBindingDimensions(0);
   context_->setBindingDimensions(
     0, nvinfer1::Dims4(batch_size, input_dims.d[1], input_dims.d[2], input_dims.d[3]));
   context_->enqueueV2(buffers.data(), stream_, nullptr);
+#endif
+
   cudaStreamSynchronize(stream_);
 }
 
@@ -316,13 +327,25 @@ bool Net::detect(const cv::Mat & in_img, float * out_scores, float * out_boxes,
 
 std::vector<int> Net::getInputDims() const
 {
-  auto dims = engine_->getBindingDimensions(0);
+#if (NV_TENSORRT_MAJOR * 10000) + (NV_TENSORRT_MINOR * 100) + NV_TENSOR_PATCH >= 80500
+  const auto dims = engine_->getTensorShape(engine_->getIOTensorName(0));
+#else
+  // Deprecated since 8.5
+  const auto dims = engine_->getBindingDimensions(0);
+#endif
   return {dims.d[1], dims.d[2], dims.d[3]};
 }
 
 int Net::getMaxBatchSize() const
 {
+#if (NV_TENSORRT_MAJOR * 10000) + (NV_TENSORRT_MINOR * 100) + NV_TENSOR_PATCH >= 80500
+  return engine_
+    ->getProfileShape(engine_->getIOTensorName(0), 0, nvinfer1::OptProfileSelector::kMAX)
+    .d[0];
+#else
+  // Deprecated since 8.5
   return engine_->getProfileDimensions(0, 0, nvinfer1::OptProfileSelector::kMAX).d[0];
+#endif
 }
 
 int Net::getInputSize() const
@@ -333,6 +356,14 @@ int Net::getInputSize() const
   return input_size;
 }
 
-int Net::getMaxDetections() const { return engine_->getBindingDimensions(1).d[1]; }
+int Net::getMaxDetections() const
+{
+#if (NV_TENSORRT_MAJOR * 10000) + (NV_TENSORRT_MINOR * 100) + NV_TENSOR_PATCH >= 80500
+  return engine_->getTensorShape(engine_->getIOTensorName(1)).d[1];
+#else
+  // Deprecated since 8.5
+  return engine_->getBindingDimensions(1).d[1];
+#endif
+}
 
 }  // namespace yolo