Fixes

nvdsinfer_custom_impl_Yolo/nvdsinfer_yolo_engine.cpp

@@ -62,6 +62,7 @@ getYoloNetworkInfo(NetworkInfo& networkInfo, const NvDsInferContextInitParams* i
   networkInfo.clusterMode = initParams->clusterMode;
   networkInfo.scaleFactor = initParams->networkScaleFactor;
   networkInfo.offsets = initParams->offsets;
+  networkInfo.workspaceSize = initParams->workspaceSize;
 
   if (initParams->networkMode == NvDsInferNetworkMode_FP32)
     networkInfo.networkMode = "FP32";
@@ -101,6 +102,8 @@ NvDsInferCreateModelParser(const NvDsInferContextInitParams* initParams)
   return new Yolo(networkInfo);
 }
 #else
+
+#if NV_TENSORRT_MAJOR >= 8
 extern "C" bool
 NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder* const builder, nvinfer1::IBuilderConfig* const builderConfig,
     const NvDsInferContextInitParams* const initParams, nvinfer1::DataType dataType, nvinfer1::ICudaEngine*& cudaEngine);
@@ -108,13 +111,29 @@ NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder* const builder, nvinfer1::IBuilder
 extern "C" bool
 NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder* const builder, nvinfer1::IBuilderConfig* const builderConfig,
     const NvDsInferContextInitParams* const initParams, nvinfer1::DataType dataType, nvinfer1::ICudaEngine*& cudaEngine)
+#else
+extern "C" bool
+NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder* const builder, const NvDsInferContextInitParams* const initParams,
+    nvinfer1::DataType dataType, nvinfer1::ICudaEngine*& cudaEngine);
+
+extern "C" bool
+NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder* const builder, const NvDsInferContextInitParams* const initParams,
+    nvinfer1::DataType dataType, nvinfer1::ICudaEngine*& cudaEngine)
+#endif
+
 {
   NetworkInfo networkInfo;
   if (!getYoloNetworkInfo(networkInfo, initParams))
     return false;
 
   Yolo yolo(networkInfo);
+
+#if NV_TENSORRT_MAJOR >= 8
   cudaEngine = yolo.createEngine(builder, builderConfig);
+#else
+  cudaEngine = yolo.createEngine(builder);
+#endif
+
   if (cudaEngine == nullptr) {
     std::cerr << "Failed to build CUDA engine" << std::endl;
     return false;

nvdsinfer_custom_impl_Yolo/yolo.cpp

@@ -39,8 +39,8 @@ Yolo::Yolo(const NetworkInfo& networkInfo) : m_InputBlobName(networkInfo.inputBl
     m_ImplicitBatch(networkInfo.implicitBatch), m_Int8CalibPath(networkInfo.int8CalibPath),
     m_DeviceType(networkInfo.deviceType), m_NumDetectedClasses(networkInfo.numDetectedClasses),
     m_ClusterMode(networkInfo.clusterMode), m_NetworkMode(networkInfo.networkMode), m_ScaleFactor(networkInfo.scaleFactor),
-    m_Offsets(networkInfo.offsets), m_InputC(0), m_InputH(0), m_InputW(0), m_InputSize(0), m_NumClasses(0), m_LetterBox(0),
-    m_NewCoords(0), m_YoloCount(0)
+    m_Offsets(networkInfo.offsets), m_WorkspaceSize(networkInfo.workspaceSize), m_InputC(0), m_InputH(0), m_InputW(0),
+    m_InputSize(0), m_NumClasses(0), m_LetterBox(0), m_NewCoords(0), m_YoloCount(0)
 {
 }
 
@@ -50,10 +50,22 @@ Yolo::~Yolo()
 }
 
 nvinfer1::ICudaEngine* 
+#if NV_TENSORRT_MAJOR >= 8
 Yolo::createEngine(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config)
+#else
+Yolo::createEngine(nvinfer1::IBuilder* builder)
+#endif
+
 {
   assert(builder);
 
+#if NV_TENSORRT_MAJOR < 8
+  nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
+  if (m_WorkspaceSize > 0) {
+    config->setMaxWorkspaceSize((size_t) m_WorkspaceSize * 1024 * 1024);
+  }
+#endif
+
   nvinfer1::NetworkDefinitionCreationFlags flags =
       1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
 
@@ -63,7 +75,13 @@ Yolo::createEngine(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config
   nvonnxparser::IParser* parser;
 
   if (m_NetworkType == "onnx") {
+
+#if NV_TENSORRT_MAJOR >= 8
     parser = nvonnxparser::createParser(*network, *builder->getLogger());
+#else
+    parser = nvonnxparser::createParser(*network, logger);
+#endif
+
     if (!parser->parseFromFile(m_OnnxWtsFilePath.c_str(), static_cast<INT>(nvinfer1::ILogger::Severity::kWARNING))) {
       std::cerr << "\nCould not parse the ONNX model\n" << std::endl;
 
@@ -72,6 +90,7 @@ Yolo::createEngine(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config
       delete network;
 #else
       parser->destroy();
+      config->destroy();
       network->destroy();
 #endif
 
@@ -89,6 +108,7 @@ Yolo::createEngine(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config
 #if NV_TENSORRT_MAJOR >= 8
       delete network;
 #else
+      config->destroy();
       network->destroy();
 #endif
 
@@ -170,6 +190,7 @@ Yolo::createEngine(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config
       if (m_NetworkType == "onnx") {
         parser->destroy();
       }
+      config->destroy();
       network->destroy();
 #endif
 
@@ -220,6 +241,7 @@ Yolo::createEngine(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config
   if (m_NetworkType == "onnx") {
     parser->destroy();
   }
+  config->destroy();
   network->destroy();
 #endif
 

nvdsinfer_custom_impl_Yolo/yolo.h

@@ -42,9 +42,16 @@
 #include "layers/reorg_layer.h"
 
 #if NV_TENSORRT_MAJOR >= 8
-  #define INT int32_t
+#define INT int32_t
 #else
-  #define INT int
+#define INT int
+
+static class Logger : public nvinfer1::ILogger {
+  void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override {
+    if (severity <= nvinfer1::ILogger::Severity::kWARNING)
+      std::cout << msg << std::endl;
+  }
+} logger;
 #endif
 
 struct NetworkInfo
@@ -64,6 +71,7 @@ struct NetworkInfo
   std::string networkMode;
   float scaleFactor;
   const float* offsets;
+  uint workspaceSize;
 };
 
 struct TensorInfo
@@ -92,7 +100,11 @@ class Yolo : public IModelParser {
 
     NvDsInferStatus parseModel(nvinfer1::INetworkDefinition& network) override;
 
+#if NV_TENSORRT_MAJOR >= 8
     nvinfer1::ICudaEngine* createEngine(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config);
+#else
+    nvinfer1::ICudaEngine* createEngine(nvinfer1::IBuilder* builder);
+#endif
 
   protected:
     const std::string m_InputBlobName;
@@ -110,6 +122,7 @@ class Yolo : public IModelParser {
     const std::string m_NetworkMode;
     const float m_ScaleFactor;
     const float* m_Offsets;
+    const uint m_WorkspaceSize;
 
     uint m_InputC;
     uint m_InputH;