add nms_cpu for layout_picodet (#8453)

PaddlePaddle · Jul 21, 2023 · d70598e · d70598e
1 parent 484510a
commit d70598e
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 5 deletions.
diff --git a/configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml b/configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml
@@ -17,6 +17,7 @@ PicoDet:
   backbone: LCNet
   neck: CSPPAN
   head: PicoHead
+  nms_cpu: True
 
 LCNet:
   scale: 1.0

diff --git a/configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x2_5_layout.yml b/configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x2_5_layout.yml
@@ -10,6 +10,7 @@ PicoDet:
   backbone: LCNet
   neck: CSPPAN
   head: PicoHead
+  nms_cpu: True
 
 LCNet:
   scale: 2.5

diff --git a/ppdet/modeling/architectures/picodet.py b/ppdet/modeling/architectures/picodet.py
@@ -36,13 +36,14 @@ class PicoDet(BaseArch):
 
     __category__ = 'architecture'
 
-    def __init__(self, backbone, neck, head='PicoHead'):
+    def __init__(self, backbone, neck, head='PicoHead', nms_cpu=False):
         super(PicoDet, self).__init__()
         self.backbone = backbone
         self.neck = neck
         self.head = head
         self.export_post_process = True
         self.export_nms = True
+        self.nms_cpu = nms_cpu
 
     @classmethod
     def from_config(cls, cfg, *args, **kwargs):
@@ -69,7 +70,10 @@ def _forward(self):
         else:
             scale_factor = self.inputs['scale_factor']
             bboxes, bbox_num = self.head.post_process(
-                head_outs, scale_factor, export_nms=self.export_nms)
+                head_outs,
+                scale_factor,
+                export_nms=self.export_nms,
+                nms_cpu=self.nms_cpu)
             return bboxes, bbox_num
 
     def get_loss(self, ):

diff --git a/ppdet/modeling/heads/pico_head.py b/ppdet/modeling/heads/pico_head.py
@@ -242,6 +242,7 @@ def __init__(self,
         self.nms_pre = nms_pre
         self.cell_offset = cell_offset
         self.eval_size = eval_size
+        self.device = paddle.device.get_device()
 
         self.use_sigmoid = self.loss_vfl.use_sigmoid
         if self.use_sigmoid:
@@ -397,7 +398,11 @@ def _generate_anchors(self, feats=None):
         stride_tensor = paddle.concat(stride_tensor)
         return anchor_points, stride_tensor
 
-    def post_process(self, head_outs, scale_factor, export_nms=True):
+    def post_process(self,
+                     head_outs,
+                     scale_factor,
+                     export_nms=True,
+                     nms_cpu=False):
         pred_scores, pred_bboxes = head_outs
         if not export_nms:
             return pred_bboxes, pred_scores
@@ -409,7 +414,12 @@ def post_process(self, head_outs, scale_factor, export_nms=True):
                 axis=-1).reshape([-1, 1, 4])
             # scale bbox to origin image size.
             pred_bboxes /= scale_factor
-            bbox_pred, bbox_num, _ = self.nms(pred_bboxes, pred_scores)
+            if nms_cpu:
+                paddle.set_device("cpu")
+                bbox_pred, bbox_num, _ = self.nms(pred_bboxes, pred_scores)
+                paddle.set_device(self.device)
+            else:
+                bbox_pred, bbox_num, _ = self.nms(pred_bboxes, pred_scores)
             return bbox_pred, bbox_num
 
 
@@ -767,7 +777,11 @@ def _generate_anchors(self, feats=None):
         stride_tensor = paddle.concat(stride_tensor)
         return anchor_points, stride_tensor
 
-    def post_process(self, head_outs, scale_factor, export_nms=True):
+    def post_process(self,
+                     head_outs,
+                     scale_factor,
+                     export_nms=True,
+                     nms_cpu=False):
         pred_scores, pred_bboxes = head_outs
         if not export_nms:
             return pred_bboxes, pred_scores