From 93e2d433ff07aaf95605d7d5e39779340c6b961c Mon Sep 17 00:00:00 2001
From: Feng Ni <nemonameless@qq.com>
Date: Thu, 23 Feb 2023 10:47:41 +0800
Subject: [PATCH] [cherry-pick] Fix run_benchmark (#7815)

* fix run_benchmark for small model accurate speed

* fix run_benchmark for other det models
---
 deploy/python/infer.py | 46 ++++++++++++++++++++++++++++++++++--------
 1 file changed, 38 insertions(+), 8 deletions(-)

diff --git a/deploy/python/infer.py b/deploy/python/infer.py
index e86761abc55..31e491b1278 100644
--- a/deploy/python/infer.py
+++ b/deploy/python/infer.py
@@ -181,7 +181,7 @@ def filter_box(self, result, threshold):
         filter_res = {'boxes': boxes, 'boxes_num': filter_num}
         return filter_res
 
-    def predict(self, repeats=1):
+    def predict(self, repeats=1, run_benchmark=False):
         '''
         Args:
             repeats (int): repeats number for prediction
@@ -193,6 +193,15 @@ def predict(self, repeats=1):
         '''
         # model prediction
         np_boxes_num, np_boxes, np_masks = np.array([0]), None, None
+
+        if run_benchmark:
+            for i in range(repeats):
+                self.predictor.run()
+                paddle.device.cuda.synchronize()
+            result = dict(
+                boxes=np_boxes, masks=np_masks, boxes_num=np_boxes_num)
+            return result
+
         for i in range(repeats):
             self.predictor.run()
             output_names = self.predictor.get_output_names()
@@ -272,9 +281,9 @@ def predict_image_slice(self,
                 self.det_times.preprocess_time_s.end()
 
                 # model prediction
-                result = self.predict(repeats=50)  # warmup
+                result = self.predict(repeats=50, run_benchmark=True)  # warmup
                 self.det_times.inference_time_s.start()
-                result = self.predict(repeats=repeats)
+                result = self.predict(repeats=repeats, run_benchmark=True)
                 self.det_times.inference_time_s.end(repeats=repeats)
 
                 # postprocess
@@ -370,9 +379,9 @@ def predict_image(self,
                 self.det_times.preprocess_time_s.end()
 
                 # model prediction
-                result = self.predict(repeats=50)  # warmup
+                result = self.predict(repeats=50, run_benchmark=True)  # warmup
                 self.det_times.inference_time_s.start()
-                result = self.predict(repeats=repeats)
+                result = self.predict(repeats=repeats, run_benchmark=True)
                 self.det_times.inference_time_s.end(repeats=repeats)
 
                 # postprocess
@@ -568,7 +577,7 @@ def __init__(
             output_dir=output_dir,
             threshold=threshold, )
 
-    def predict(self, repeats=1):
+    def predict(self, repeats=1, run_benchmark=False):
         '''
         Args:
             repeats (int): repeat number for prediction
@@ -577,7 +586,20 @@ def predict(self, repeats=1):
                             'cate_label': label of segm, shape:[N]
                             'cate_score': confidence score of segm, shape:[N]
         '''
-        np_label, np_score, np_segms = None, None, None
+        np_segms, np_label, np_score, np_boxes_num = None, None, None, np.array(
+            [0])
+
+        if run_benchmark:
+            for i in range(repeats):
+                self.predictor.run()
+                paddle.device.cuda.synchronize()
+            result = dict(
+                segm=np_segms,
+                label=np_label,
+                score=np_score,
+                boxes_num=np_boxes_num)
+            return result
+
         for i in range(repeats):
             self.predictor.run()
             output_names = self.predictor.get_output_names()
@@ -659,7 +681,7 @@ def postprocess(self, inputs, result):
         result = dict(boxes=np_boxes, boxes_num=np_boxes_num)
         return result
 
-    def predict(self, repeats=1):
+    def predict(self, repeats=1, run_benchmark=False):
         '''
         Args:
             repeats (int): repeat number for prediction
@@ -668,6 +690,14 @@ def predict(self, repeats=1):
                             matix element:[class, score, x_min, y_min, x_max, y_max]
         '''
         np_score_list, np_boxes_list = [], []
+
+        if run_benchmark:
+            for i in range(repeats):
+                self.predictor.run()
+                paddle.device.cuda.synchronize()
+            result = dict(boxes=np_score_list, boxes_num=np_boxes_list)
+            return result
+
         for i in range(repeats):
             self.predictor.run()
             np_score_list.clear()