From d70598ec390eb69f81043f0644538383694ca9fb Mon Sep 17 00:00:00 2001
From: xiaoting <31891223+tink2123@users.noreply.github.com>
Date: Fri, 21 Jul 2023 14:21:01 +0800
Subject: [PATCH] add nms_cpu for layout_picodet (#8453)

---
 .../picodet_lcnet_x1_0_layout.yml             |  1 +
 .../picodet_lcnet_x2_5_layout.yml             |  1 +
 ppdet/modeling/architectures/picodet.py       |  8 ++++++--
 ppdet/modeling/heads/pico_head.py             | 20 ++++++++++++++++---
 4 files changed, 25 insertions(+), 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
index db9a8c683..251a3dd58 100644
--- 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
index b6d771f61..a30807fb9 100644
--- 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
index 0b87a4baa..b6f44472a 100644
--- 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 @@ class PicoDet(BaseArch):
         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
index adcd05fc6..6e04173f5 100644
--- a/ppdet/modeling/heads/pico_head.py
+++ b/ppdet/modeling/heads/pico_head.py
@@ -242,6 +242,7 @@ class PicoHead(OTAVFLHead):
         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 @@ class PicoHead(OTAVFLHead):
         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 @@ class PicoHead(OTAVFLHead):
                 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 @@ class PicoHeadV2(GFLHead):
         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
-- 
GitLab