From a853016850a3a0a173544d7072bd8cc115dd279e Mon Sep 17 00:00:00 2001 From: wangguanzhong Date: Tue, 10 May 2022 19:12:33 +0800 Subject: [PATCH] refine device when assign on cpu (#5924) --- ppdet/modeling/post_process.py | 6 +++--- ppdet/modeling/proposal_generator/target.py | 4 +++- ppdet/modeling/proposal_generator/target_layer.py | 6 ++---- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/ppdet/modeling/post_process.py b/ppdet/modeling/post_process.py index b50874c3a..27890c17e 100644 --- a/ppdet/modeling/post_process.py +++ b/ppdet/modeling/post_process.py @@ -194,8 +194,7 @@ class MaskPostProcess(object): super(MaskPostProcess, self).__init__() self.binary_thresh = binary_thresh self.export_onnx = export_onnx - self.assign_on_cpu = assign_on_cpu and paddle.device.is_compiled_with_cuda( - ) + self.assign_on_cpu = assign_on_cpu def paste_mask(self, masks, boxes, im_h, im_w): """ @@ -240,6 +239,7 @@ class MaskPostProcess(object): """ num_mask = mask_out.shape[0] origin_shape = paddle.cast(origin_shape, 'int32') + device = paddle.device.get_device() if self.export_onnx: h, w = origin_shape[0][0], origin_shape[0][1] @@ -269,7 +269,7 @@ class MaskPostProcess(object): im_w] = pred_mask id_start += bbox_num[i] if self.assign_on_cpu: - paddle.set_device('gpu') + paddle.set_device(device) return pred_result diff --git a/ppdet/modeling/proposal_generator/target.py b/ppdet/modeling/proposal_generator/target.py index 62bbe8b6f..fd04f0522 100644 --- a/ppdet/modeling/proposal_generator/target.py +++ b/ppdet/modeling/proposal_generator/target.py @@ -74,9 +74,11 @@ def label_box(anchors, is_crowd=None, assign_on_cpu=False): if assign_on_cpu: + device = paddle.device.get_device() paddle.set_device("cpu") iou = bbox_overlaps(gt_boxes, anchors) - paddle.set_device("gpu") + paddle.set_device(device) + else: iou = bbox_overlaps(gt_boxes, anchors) n_gt = gt_boxes.shape[0] diff --git a/ppdet/modeling/proposal_generator/target_layer.py b/ppdet/modeling/proposal_generator/target_layer.py index da9b0c5e2..3b5a09601 100644 --- a/ppdet/modeling/proposal_generator/target_layer.py +++ b/ppdet/modeling/proposal_generator/target_layer.py @@ -68,8 +68,7 @@ class RPNTargetAssign(object): self.negative_overlap = negative_overlap self.ignore_thresh = ignore_thresh self.use_random = use_random - self.assign_on_cpu = assign_on_cpu and paddle.device.is_compiled_with_cuda( - ) + self.assign_on_cpu = assign_on_cpu def __call__(self, inputs, anchors): """ @@ -150,8 +149,7 @@ class BBoxAssigner(object): self.use_random = use_random self.cascade_iou = cascade_iou self.num_classes = num_classes - self.assign_on_cpu = assign_on_cpu and paddle.device.is_compiled_with_cuda( - ) + self.assign_on_cpu = assign_on_cpu def __call__(self, rpn_rois, -- GitLab