TensorRT works with custom batch size

Author: Nuzhny007

example/examples.h

@@ -595,6 +595,7 @@ class YoloDarknetExample : public VideoExample
 #else
 		std::string pathToModel = "../data/";
 #endif
+		size_t maxBatch = 1;
 		enum class YOLOModels
 		{
 			TinyYOLOv3 = 0,
@@ -628,16 +629,19 @@ class YoloDarknetExample : public VideoExample
             config.emplace("modelConfiguration", pathToModel + "yolov4-tiny.cfg");
             config.emplace("modelBinary", pathToModel + "yolov4-tiny.weights");
             config.emplace("confidenceThreshold", "0.5");
-			config.emplace("maxBatch", "4");
+			maxBatch = 4;
             break;
 
 		case YOLOModels::ScaledYOLOv4:
 			config.emplace("modelConfiguration", pathToModel + "yolov4-csp.cfg");
 			config.emplace("modelBinary", pathToModel + "yolov4-csp.weights");
 			config.emplace("confidenceThreshold", "0.5");
-			config.emplace("maxBatch", "2");
+			maxBatch = 2;
 			break;
 		}
+		if (maxBatch < m_batchSize)
+			maxBatch = m_batchSize;
+		config.emplace("maxBatch", std::to_string(m_batchSize));
         config.emplace("classNames", pathToModel + "coco.names");
         config.emplace("maxCropRatio", "-1");
 
@@ -787,7 +791,7 @@ class YoloTensorRTExample : public VideoExample
 #else
 		std::string pathToModel = "../data/";
 #endif
-
+		size_t maxBatch = 1;
         enum class YOLOModels
         {
             TinyYOLOv3 = 0,
@@ -805,7 +809,7 @@ class YoloTensorRTExample : public VideoExample
             config.emplace("confidenceThreshold", "0.5");
             config.emplace("inference_precison", "FP32");
             config.emplace("net_type", "YOLOV3_TINY");
-			config.emplace("maxBatch", "4");
+			maxBatch = 4;
 			config.emplace("maxCropRatio", "2");
             break;
 
@@ -815,17 +819,17 @@ class YoloTensorRTExample : public VideoExample
             config.emplace("confidenceThreshold", "0.7");
             config.emplace("inference_precison", "FP32");
             config.emplace("net_type", "YOLOV3");
-			config.emplace("maxBatch", "2");
+			maxBatch = 2;
 			config.emplace("maxCropRatio", "-1");
             break;
 
         case YOLOModels::YOLOv4:
             config.emplace("modelConfiguration", pathToModel + "yolov4.cfg");
             config.emplace("modelBinary", pathToModel + "yolov4.weights");
             config.emplace("confidenceThreshold", "0.8");
-            config.emplace("inference_precison", "FP16");
+            config.emplace("inference_precison", "FP32");
             config.emplace("net_type", "YOLOV4");
-			config.emplace("maxBatch", "1");
+			maxBatch = 1;
 			config.emplace("maxCropRatio", "-1");
             break;
 
@@ -835,7 +839,7 @@ class YoloTensorRTExample : public VideoExample
             config.emplace("confidenceThreshold", "0.5");
             config.emplace("inference_precison", "FP32");
             config.emplace("net_type", "YOLOV4_TINY");
-			config.emplace("maxBatch", "4");
+			maxBatch = 4;
 			config.emplace("maxCropRatio", "1");
             break;
 
@@ -845,11 +849,13 @@ class YoloTensorRTExample : public VideoExample
             config.emplace("confidenceThreshold", "0.5");
             config.emplace("inference_precison", "FP32");
             config.emplace("net_type", "YOLOV5");
-            config.emplace("maxBatch", "1");
+			maxBatch = 1;
             config.emplace("maxCropRatio", "-1");
             break;
         }
-
+		if (maxBatch < m_batchSize)
+			maxBatch = m_batchSize;
+		config.emplace("maxBatch", std::to_string(m_batchSize));
 		config.emplace("classNames", pathToModel + "coco.names");
 
 		config.emplace("white_list", std::to_string((objtype_t)ObjectTypes::obj_person));

src/Detector/YoloTensorRTDetector.cpp

@@ -47,6 +47,7 @@ bool YoloTensorRTDetector::Init(const config_t& config)
 	auto maxBatch = config.find("maxBatch");
 	if (maxBatch != config.end())
         m_batchSize = std::max(1, std::stoi(maxBatch->second));
+	m_localConfig.batch_size = m_batchSize;
 
 	m_localConfig.file_model_cfg = modelConfiguration->second;
 	m_localConfig.file_model_weights = modelBinary->second;

src/Detector/tensorrt_yolo/class_detector.h

@@ -23,13 +23,13 @@ namespace tensor_rt
 
 	enum ModelType
 	{
-		YOLOV2 = 0,
-		YOLOV3,
-		YOLOV2_TINY,
-		YOLOV3_TINY,
-		YOLOV4,
-		YOLOV4_TINY,
-                YOLOV5
+        YOLOV2 = 0,
+        YOLOV3,
+        YOLOV2_TINY,
+        YOLOV3_TINY,
+        YOLOV4,
+        YOLOV4_TINY,
+        YOLOV5
 	};
 
 	enum Precision
@@ -41,9 +41,9 @@ namespace tensor_rt
 
 	struct Config
 	{
-		std::string file_model_cfg = "configs/yolov3.cfg";
+		std::string file_model_cfg = "yolov4.cfg";
 
-		std::string file_model_weights = "configs/yolov3.weights";
+		std::string file_model_weights = "yolov4.weights";
 
 		float detect_thresh = 0.9f;
 
@@ -53,6 +53,8 @@ namespace tensor_rt
 
 		int	gpu_id = 0;
 
+		uint32_t batch_size = 1;
+
 		std::string calibration_image_list_file_txt = "configs/calibration_images.txt";
 	};
 

src/Detector/tensorrt_yolo/class_yolo_detector.hpp

@@ -116,6 +116,7 @@ class YoloDectector
 		_infer_param.calibImagesPath = "";
 		_infer_param.probThresh = _config.detect_thresh;
 		_infer_param.nmsThresh = 0.5;
+		_infer_param.batchSize = _config.batch_size;
 	}
 
 	void build_net()

src/Detector/tensorrt_yolo/yolo.cpp

@@ -28,6 +28,7 @@ Yolo::Yolo( const NetworkInfo& networkInfo, const InferParams& inferParams) :
 	m_NMSThresh(inferParams.nmsThresh),
 	m_PrintPerfInfo(inferParams.printPerfInfo),
 	m_PrintPredictions(inferParams.printPredictionInfo),
+	m_BatchSize(inferParams.batchSize),
 	m_Logger(Logger()),
 	m_Network(nullptr),
 	m_Builder(nullptr),

src/Detector/tensorrt_yolo/yolo.h

@@ -64,12 +64,13 @@ struct NetworkInfo
  */
 struct InferParams
 {
-    bool printPerfInfo;
-    bool printPredictionInfo;
+    bool printPerfInfo = false;
+    bool printPredictionInfo = false;
     std::string calibImages;
     std::string calibImagesPath;
-    float probThresh;
-    float nmsThresh;
+    float probThresh = 0.5f;
+    float nmsThresh = 0.5f;
+	uint32_t batchSize = 1;
 };
 
 /**