【PIR API adaptor No.161、162】Migrate paddle.vision.ops.nms paddle.nn.functional.one_hot into pir (#58735)

ooooo-create · web-flow · commit e8eb917ea889 · 2023-12-06T11:33:46.000+08:00
diff --git a/python/paddle/tensor/search.py b/python/paddle/tensor/search.py
@@ -99,7 +99,7 @@ def argsort(x, axis=-1, descending=False, name=None):
               [1, 1, 0, 2],
               [0, 2, 1, 1]]])
     """
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         _, ids = _C_ops.argsort(x, axis, descending)
         return ids
     else:
diff --git a/python/paddle/vision/ops.py b/python/paddle/vision/ops.py
@@ -1937,7 +1937,7 @@ def nms(
     """
 
     def _nms(boxes, iou_threshold):
-        if in_dygraph_mode():
+        if in_dynamic_or_pir_mode():
             return _C_ops.nms(boxes, iou_threshold)
 
         else:
@@ -1971,10 +1971,12 @@ def _nms(boxes, iou_threshold):
         categories is not None
     ), "if category_idxs is given, categories which is a list of unique id of all categories is necessary"
 
-    mask = paddle.zeros_like(scores, dtype=paddle.int32)
+    mask = paddle.zeros_like(scores, dtype='int32')
 
     for category_id in categories:
-        cur_category_boxes_idxs = paddle.where(category_idxs == category_id)[0]
+        cur_category_boxes_idxs = paddle.where(
+            paddle.equal(category_idxs, paddle.to_tensor(category_id))
+        )[0]
         shape = cur_category_boxes_idxs.shape[0]
         cur_category_boxes_idxs = paddle.reshape(
             cur_category_boxes_idxs, [shape]
@@ -1999,7 +2001,7 @@ def _nms(boxes, iou_threshold):
 
         updates = paddle.ones_like(
             cur_category_boxes_idxs[cur_category_keep_boxes_sub_idxs],
-            dtype=paddle.int32,
+            dtype='int32',
         )
         mask = paddle.scatter(
             mask,
diff --git a/test/legacy_test/test_nms_op.py b/test/legacy_test/test_nms_op.py
@@ -90,7 +90,7 @@ def init_dtype_type(self):
         pass
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
 
 if __name__ == "__main__":
diff --git a/test/legacy_test/test_nn_functional_hot_op.py b/test/legacy_test/test_nn_functional_hot_op.py
@@ -21,6 +21,7 @@
 from paddle import base
 from paddle.base import core
 from paddle.nn import functional
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestOneHotOp(OpTest):
@@ -118,13 +119,59 @@ def test_check_output(self):
 
 
 class TestOneHotOpApi(unittest.TestCase):
+    @test_with_pir_api
     def test_api(self):
-        num_classes = 10
-        self._run(num_classes)
+        main = paddle.static.Program()
+        startup = paddle.static.Program()
+        with paddle.static.program_guard(main, startup):
+            num_classes = 10
+            label = paddle.static.data(
+                name="label", shape=[-1, 1], dtype="int64"
+            )
+            one_hot_label = functional.one_hot(x=label, num_classes=num_classes)
+
+            place = base.CPUPlace()
+            label_data = np.array(
+                [np.random.randint(0, 10 - 1) for i in range(6)]
+            ).reshape([6, 1])
+            label_data = label_data.astype('int64')
+
+            exe = base.Executor(place)
+            exe.run(startup)
+            ret = exe.run(
+                feed={
+                    'label': label_data,
+                },
+                fetch_list=[one_hot_label],
+                return_numpy=False,
+            )
 
+    @test_with_pir_api
     def test_api_with_depthTensor(self):
-        num_classes = paddle.assign(np.array([10], dtype=np.int32))
-        self._run(num_classes)
+        main = paddle.static.Program()
+        startup = paddle.static.Program()
+        with paddle.static.program_guard(main, startup):
+            num_classes = paddle.assign(np.array([10], dtype=np.int32))
+            label = paddle.static.data(
+                name="label", shape=[-1, 1], dtype="int64"
+            )
+            one_hot_label = functional.one_hot(x=label, num_classes=num_classes)
+
+            place = base.CPUPlace()
+            label_data = np.array(
+                [np.random.randint(0, 10 - 1) for i in range(6)]
+            ).reshape([6, 1])
+            label_data = label_data.astype('int64')
+
+            exe = base.Executor(place)
+            exe.run(startup)
+            ret = exe.run(
+                feed={
+                    'label': label_data,
+                },
+                fetch_list=[one_hot_label],
+                return_numpy=False,
+            )
 
     def test_api_with_dygraph(self):
         num_classes = 10
@@ -136,26 +183,6 @@ def test_api_with_dygraph(self):
                 x=base.dygraph.to_variable(label), num_classes=num_classes
             )
 
-    def _run(self, num_classes):
-        label = paddle.static.data(name="label", shape=[-1, 1], dtype="int64")
-        label.desc.set_need_check_feed(False)
-        one_hot_label = functional.one_hot(x=label, num_classes=num_classes)
-
-        place = base.CPUPlace()
-        label_data = np.array(
-            [np.random.randint(0, 10 - 1) for i in range(6)]
-        ).reshape([6, 1])
-
-        exe = base.Executor(place)
-        exe.run(base.default_startup_program())
-        ret = exe.run(
-            feed={
-                'label': label_data,
-            },
-            fetch_list=[one_hot_label],
-            return_numpy=False,
-        )
-
 
 class BadInputTestOnehotV2(unittest.TestCase):
     def test_error(self):
diff --git a/test/legacy_test/test_one_hot_v2_op.py b/test/legacy_test/test_one_hot_v2_op.py
@@ -20,6 +20,7 @@
 import paddle
 from paddle import base
 from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 
 def one_hot_wrapper(x, depth_tensor, **keargs):
@@ -128,13 +129,63 @@ def test_check_output(self):
 
 
 class TestOneHotOpApi(unittest.TestCase):
+    @test_with_pir_api
     def test_api(self):
-        depth = 10
-        self._run(depth)
+        main = paddle.static.Program()
+        startup = paddle.static.Program()
+        with paddle.static.program_guard(main, startup):
+            depth = 10
+            label = paddle.static.data(
+                name="label", shape=[-1, 1], dtype="int64"
+            )
+            one_hot_label = paddle.nn.functional.one_hot(
+                x=label, num_classes=depth
+            )
 
+            place = paddle.CPUPlace()
+            label_data = np.array(
+                [np.random.randint(0, 10 - 1) for i in range(6)]
+            ).reshape([6, 1])
+            label_data = label_data.astype('int64')
+
+            exe = base.Executor(place)
+            exe.run(startup)
+            ret = exe.run(
+                feed={
+                    'label': label_data,
+                },
+                fetch_list=[one_hot_label],
+                return_numpy=False,
+            )
+
+    @test_with_pir_api
     def test_api_with_depthTensor(self):
-        depth = paddle.assign(np.array([10], dtype=np.int32))
-        self._run(depth)
+        main = paddle.static.Program()
+        startup = paddle.static.Program()
+        with paddle.static.program_guard(main, startup):
+            depth = paddle.assign(np.array([10], dtype=np.int32))
+            label = paddle.static.data(
+                name="label", shape=[-1, 1], dtype="int64"
+            )
+            one_hot_label = paddle.nn.functional.one_hot(
+                x=label, num_classes=depth
+            )
+
+            place = paddle.CPUPlace()
+            label_data = np.array(
+                [np.random.randint(0, 10 - 1) for i in range(6)]
+            ).reshape([6, 1])
+            label_data = label_data.astype('int64')
+
+            exe = base.Executor(place)
+            exe.run(startup)
+            ret = exe.run(
+                feed={
+                    'label': label_data,
+                },
+                fetch_list=[one_hot_label],
+                return_numpy=False,
+            )
 
     def test_api_with_dygraph(self):
         depth = 10
@@ -149,30 +200,10 @@ def test_api_with_dygraph(self):
                 paddle.to_tensor(label), depth
             )
 
-    def _run(self, depth):
-        label = paddle.static.data(name="label", shape=[-1, 1], dtype="int64")
-        label.desc.set_need_check_feed(False)
-        one_hot_label = paddle.nn.functional.one_hot(x=label, num_classes=depth)
-
-        place = base.CPUPlace()
-        label_data = np.array(
-            [np.random.randint(0, 10 - 1) for i in range(6)]
-        ).reshape([6, 1])
-
-        exe = base.Executor(place)
-        exe.run(base.default_startup_program())
-        ret = exe.run(
-            feed={
-                'label': label_data,
-            },
-            fetch_list=[one_hot_label],
-            return_numpy=False,
-        )
-
 
 class BadInputTestOnehotV2(unittest.TestCase):
     def test_error(self):
-        with base.program_guard(base.Program()):
+        with paddle.static.program_guard(paddle.static.Program()):
 
             def test_bad_x():
                 label = paddle.static.data(
diff --git a/test/legacy_test/test_ops_nms.py b/test/legacy_test/test_ops_nms.py
@@ -20,6 +20,7 @@
 from test_nms_op import nms
 
 import paddle
+from paddle.pir_utils import test_with_pir_api
 
 
 def _find(condition):
@@ -138,56 +139,60 @@ def test_multiclass_nms_dynamic(self):
                     err_msg=f'paddle out: {out}\n py out: {out_py}\n',
                 )
 
+    @test_with_pir_api
     def test_multiclass_nms_static(self):
         for device in self.devices:
             for dtype in self.dtypes:
-                paddle.enable_static()
-                boxes, scores, category_idxs, categories = gen_args(
-                    self.num_boxes, dtype
-                )
-                boxes_static = paddle.static.data(
-                    shape=boxes.shape, dtype=boxes.dtype, name="boxes"
-                )
-                scores_static = paddle.static.data(
-                    shape=scores.shape, dtype=scores.dtype, name="scores"
-                )
-                category_idxs_static = paddle.static.data(
-                    shape=category_idxs.shape,
-                    dtype=category_idxs.dtype,
-                    name="category_idxs",
-                )
-                out = paddle.vision.ops.nms(
-                    boxes_static,
-                    self.threshold,
-                    scores_static,
-                    category_idxs_static,
-                    categories,
-                    self.topk,
-                )
-                place = paddle.CPUPlace()
-                if device == 'gpu':
-                    place = paddle.CUDAPlace(0)
-                exe = paddle.static.Executor(place)
-                out = exe.run(
-                    paddle.static.default_main_program(),
-                    feed={
-                        'boxes': boxes,
-                        'scores': scores,
-                        'category_idxs': category_idxs,
-                    },
-                    fetch_list=[out],
-                )
-                paddle.disable_static()
-                out_py = multiclass_nms(
-                    boxes, scores, category_idxs, self.threshold, self.topk
-                )
-                out = np.array(out)
-                out = np.squeeze(out)
-                np.testing.assert_array_equal(
-                    out,
-                    out_py,
-                    err_msg=f'paddle out: {out}\n py out: {out_py}\n',
-                )
+                with paddle.static.program_guard(
+                    paddle.static.Program(), paddle.static.Program()
+                ):
+                    paddle.enable_static()
+                    boxes, scores, category_idxs, categories = gen_args(
+                        self.num_boxes, dtype
+                    )
+                    boxes_static = paddle.static.data(
+                        shape=boxes.shape, dtype=boxes.dtype, name="boxes"
+                    )
+                    scores_static = paddle.static.data(
+                        shape=scores.shape, dtype=scores.dtype, name="scores"
+                    )
+                    category_idxs_static = paddle.static.data(
+                        shape=category_idxs.shape,
+                        dtype=category_idxs.dtype,
+                        name="category_idxs",
+                    )
+                    out = paddle.vision.ops.nms(
+                        boxes_static,
+                        self.threshold,
+                        scores_static,
+                        category_idxs_static,
+                        categories,
+                        self.topk,
+                    )
+                    place = paddle.CPUPlace()
+                    if device == 'gpu':
+                        place = paddle.CUDAPlace(0)
+                    exe = paddle.static.Executor(place)
+                    out = exe.run(
+                        paddle.static.default_main_program(),
+                        feed={
+                            'boxes': boxes,
+                            'scores': scores,
+                            'category_idxs': category_idxs,
+                        },
+                        fetch_list=[out],
+                    )
+                    paddle.disable_static()
+                    out_py = multiclass_nms(
+                        boxes, scores, category_idxs, self.threshold, self.topk
+                    )
+                    out = np.array(out)
+                    out = np.squeeze(out)
+                    np.testing.assert_array_equal(
+                        out,
+                        out_py,
+                        err_msg=f'paddle out: {out}\n py out: {out_py}\n',
+                    )
 
     def test_multiclass_nms_dynamic_to_static(self):
         for device in self.devices:
diff --git a/test/legacy_test/test_zero_dim_tensor.py b/test/legacy_test/test_zero_dim_tensor.py
@@ -5939,6 +5939,7 @@ def test_static_embedding(self):
         res = self.exe.run(prog, fetch_list=[emb])
         self.assertEqual(res[0].shape, (3,))
 
+    @test_with_pir_api
     def test_one_hot_label(self):
         label = paddle.full(shape=[], fill_value=2, dtype='int64')
         one_hot_label = paddle.nn.functional.one_hot(label, num_classes=4)
