test_gather_nd_op.py 10.9 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
#   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
16

17 18
import numpy as np
from op_test import OpTest
19

20
import paddle
21
import paddle.fluid as fluid
22 23 24


class TestGatherNdOpWithEmptyIndex(OpTest):
H
hong 已提交
25
    # Index has empty element, which means copy entire tensor
26 27 28

    def setUp(self):
        self.op_type = "gather_nd"
29
        self.prim_op_type = "prim"
H
hong 已提交
30
        self.python_api = paddle.gather_nd
31
        xnp = np.random.random((5, 20)).astype("float64")
32 33 34 35 36 37
        self.inputs = {'X': xnp, 'Index': np.array([[], []]).astype("int32")}
        self.outputs = {
            'Out': np.vstack((xnp[np.newaxis, :], xnp[np.newaxis, :]))
        }

    def test_check_output(self):
38
        self.check_output(check_eager=False)
39 40

    def test_check_grad(self):
41
        self.check_grad(['X'], 'Out', check_eager=False, check_prim=True)
42 43


44 45 46
class TestGatherNdOpWithIndex1(OpTest):
    def setUp(self):
        self.op_type = "gather_nd"
47
        self.prim_op_type = "prim"
H
hong 已提交
48
        self.python_api = paddle.gather_nd
49 50 51 52 53
        xnp = np.random.random((5, 20)).astype("float64")
        self.inputs = {'X': xnp, 'Index': np.array([1]).astype("int32")}
        self.outputs = {'Out': self.inputs["X"][self.inputs["Index"]]}

    def test_check_output(self):
54
        self.check_output(check_eager=False)
55 56

    def test_check_grad(self):
W
wangxiaoning 已提交
57
        self.check_grad(['X'], 'Out', check_eager=False, check_prim=True)
58 59


60
class TestGatherNdOpWithLowIndex(OpTest):
61
    # Index has low rank, X has high rank
62 63 64

    def setUp(self):
        self.op_type = "gather_nd"
65
        self.prim_op_type = "prim"
H
hong 已提交
66
        self.python_api = paddle.gather_nd
67
        self.enable_cinn = False
Z
zhupengyang 已提交
68
        xnp = np.random.uniform(0, 100, (10, 10)).astype("float64")
69 70 71 72
        index = np.array([[1], [2]]).astype("int64")

        self.inputs = {'X': xnp, 'Index': index}

73 74 75
        self.outputs = {
            'Out': xnp[tuple(index.T)]
        }  # [[14, 25, 1], [76, 22, 3]]
76 77

    def test_check_output(self):
78
        self.check_output(check_eager=False)
79 80

    def test_check_grad(self):
81
        self.check_grad(['X'], 'Out', check_eager=False, check_prim=True)
82 83


84
class TestGatherNdOpIndex1(OpTest):
85
    # Index has low rank, X has high rank
86 87 88

    def setUp(self):
        self.op_type = "gather_nd"
89
        self.prim_op_type = "prim"
H
hong 已提交
90
        self.python_api = paddle.gather_nd
91
        self.init_input()
92

93
        self.inputs = {'X': self.xnp, 'Index': self.index}
94

95 96 97 98 99 100
        self.outputs = {'Out': self.xnp[tuple(self.index.T)]}
        self.enable_cinn = False

    def init_input(self):
        self.xnp = np.random.uniform(0, 100, (10, 10)).astype("float64")
        self.index = np.array([1, 2]).astype("int32")
101 102

    def test_check_output(self):
103
        self.check_output(check_eager=False)
104 105

    def test_check_grad(self):
106 107 108 109 110 111 112
        self.check_grad(['X'], 'Out', check_eager=False, check_prim=True)


class TestGatherNdOpIndex1FP16(TestGatherNdOpIndex1):
    def init_input(self):
        self.xnp = np.random.uniform(0, 100, (10, 10)).astype("float16")
        self.index = np.array([1, 2]).astype("int32")
113 114


115
class TestGatherNdOpWithSameIndexAsX(OpTest):
116
    # Index has same rank as X's rank
117 118
    def setUp(self):
        self.op_type = "gather_nd"
119
        self.prim_op_type = "prim"
H
hong 已提交
120
        self.python_api = paddle.gather_nd
121
        self.enable_cinn = False
Z
zhupengyang 已提交
122
        xnp = np.random.uniform(0, 100, (10, 10)).astype("float64")
123 124 125
        index = np.array([[1, 1], [2, 1]]).astype("int64")

        self.inputs = {'X': xnp, 'Index': index}
126
        self.outputs = {'Out': xnp[tuple(index.T)]}  # [25, 22]
127 128

    def test_check_output(self):
129
        self.check_output(check_eager=False)
130 131

    def test_check_grad(self):
132
        self.check_grad(['X'], 'Out', check_eager=False, check_prim=True)
133 134 135


class TestGatherNdOpWithHighRankSame(OpTest):
136
    # Both Index and X have high rank, and Rank(Index) = Rank(X)
137 138 139

    def setUp(self):
        self.op_type = "gather_nd"
140
        self.prim_op_type = "prim"
H
hong 已提交
141
        self.python_api = paddle.gather_nd
S
ShenLiang 已提交
142
        shape = (5, 2, 3, 1, 10)
143
        xnp = np.random.rand(*shape).astype("float64")
S
ShenLiang 已提交
144
        index = np.vstack([np.random.randint(0, s, size=2) for s in shape]).T
145 146 147 148 149

        self.inputs = {'X': xnp, 'Index': index.astype("int32")}
        self.outputs = {'Out': xnp[tuple(index.T)]}

    def test_check_output(self):
150
        self.check_output(check_eager=False)
151 152

    def test_check_grad(self):
153
        self.check_grad(['X'], 'Out', check_eager=False, check_prim=True)
154 155 156


class TestGatherNdOpWithHighRankDiff(OpTest):
157
    # Both Index and X have high rank, and Rank(Index) < Rank(X)
158 159 160

    def setUp(self):
        self.op_type = "gather_nd"
161
        self.prim_op_type = "prim"
H
hong 已提交
162
        self.python_api = paddle.gather_nd
S
ShenLiang 已提交
163
        shape = (2, 3, 4, 1, 10)
164
        xnp = np.random.rand(*shape).astype("float64")
S
ShenLiang 已提交
165 166
        index = np.vstack([np.random.randint(0, s, size=200) for s in shape]).T
        index_re = index.reshape([20, 5, 2, 5])
167 168

        self.inputs = {'X': xnp, 'Index': index_re.astype("int32")}
S
ShenLiang 已提交
169
        self.outputs = {'Out': xnp[tuple(index.T)].reshape([20, 5, 2])}
170 171

    def test_check_output(self):
172
        self.check_output(check_eager=False)
173 174

    def test_check_grad(self):
175
        self.check_grad(['X'], 'Out', check_eager=False, check_prim=True)
176 177


178
# Test Python API
179
class TestGatherNdOpAPI(unittest.TestCase):
180
    def test_case1(self):
G
GGBond8488 已提交
181 182 183 184 185
        x1 = paddle.static.data(
            name='x1', shape=[-1, 30, 40, 50, 60], dtype='float32'
        )
        index1 = paddle.static.data(
            name='index1', shape=[-1, 2, 4], dtype='int32'
186
        )
187
        output1 = paddle.gather_nd(x1, index1)
188 189

    def test_case2(self):
G
GGBond8488 已提交
190 191 192 193 194 195
        x2 = paddle.static.data(
            name='x2', shape=[-1, 30, 40, 50], dtype='float32'
        )
        index2 = paddle.static.data(
            name='index2', shape=[-1, 2, 2], dtype='int64'
        )
196
        output2 = paddle.gather_nd(x2, index2)
197 198

    def test_case3(self):
G
GGBond8488 已提交
199 200 201 202
        x3 = paddle.static.data(name='x3', shape=[-1, 3, 4, 5], dtype='float32')
        index3 = paddle.static.data(
            name='index3', shape=[-1, 2, 1], dtype='int32'
        )
203
        output3 = paddle.gather_nd(x3, index3, name="gather_nd_layer")
204 205


206
# Test Raise Index Error
207
class TestGatherNdOpRaise(unittest.TestCase):
208 209 210
    def test_check_raise(self):
        def check_raise_is_test():
            try:
G
GGBond8488 已提交
211 212
                x = paddle.static.data(
                    name='x', shape=[-1, 3, 4, 5], dtype='float32'
213
                )
G
GGBond8488 已提交
214 215
                index = paddle.static.data(
                    name='index', shape=[-1, 2, 10], dtype='int32'
216
                )
217
                output = paddle.gather_nd(x, index)
218
            except Exception as e:
219
                t = "Input(Index).shape[-1] should be no greater than Input(X).rank"
220 221 222 223 224 225
                if t in str(e):
                    raise IndexError

        self.assertRaises(IndexError, check_raise_is_test)


226 227
class TestGatherNdError(unittest.TestCase):
    def test_error(self):
228 229 230
        with paddle.static.program_guard(
            paddle.static.Program(), paddle.static.Program()
        ):
231 232

            shape = [8, 9, 6]
233 234
            x = paddle.fluid.data(shape=shape, dtype='float32', name='x')
            index = paddle.fluid.data(shape=shape, dtype='bool', name='index')
235 236 237
            index_float = paddle.fluid.data(
                shape=shape, dtype='float32', name='index_float'
            )
238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259
            np_x = np.random.random(shape).astype('float32')
            np_index = np.array(np.random.randint(2, size=shape, dtype=bool))

            def test_x_type():
                paddle.gather_nd(np_x, index)

            self.assertRaises(TypeError, test_x_type)

            def test_index_type():
                paddle.gather_nd(x, np_index)

            self.assertRaises(TypeError, test_index_type)

            def test_index_dtype():
                paddle.gather_nd(x, index_float)

            self.assertRaises(TypeError, test_index_dtype)


class TestGatherNdAPI2(unittest.TestCase):
    def test_static(self):
        with fluid.program_guard(fluid.Program(), fluid.Program()):
G
GGBond8488 已提交
260 261 262 263
            data1 = paddle.static.data('data1', shape=[-1, 2], dtype='float64')
            data1.desc.set_need_check_feed(False)
            index = paddle.static.data('index', shape=[-1, 1], dtype='int32')
            index.desc.set_need_check_feed(False)
264 265 266 267
            out = paddle.gather_nd(data1, index)
            place = fluid.CPUPlace()
            exe = fluid.Executor(place)
            input = np.array([[1, 2], [3, 4], [5, 6]])
G
GGBond8488 已提交
268
            index_1 = np.array([[1]]).astype('int32')
269 270 271
            (result,) = exe.run(
                feed={"data1": input, "index": index_1}, fetch_list=[out]
            )
272
            expected_output = np.array([[3, 4]])
273
        np.testing.assert_allclose(result, expected_output, rtol=1e-05)
274

张春乔 已提交
275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307
    def test_static_fp16_with_gpu(self):
        if paddle.fluid.core.is_compiled_with_cuda():
            place = paddle.CUDAPlace(0)
            with paddle.static.program_guard(
                paddle.static.Program(), paddle.static.Program()
            ):
                input = np.array(
                    [[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]],
                    dtype='float16',
                )
                index = np.array([[0, 1]], dtype='int32')
                res_np = np.array([[3, 4]], dtype='float16')

                x = paddle.static.data(
                    name="x", shape=[2, 3, 2], dtype="float16"
                )
                x.desc.set_need_check_feed(False)
                idx = paddle.static.data(
                    name="index", shape=[1, 2], dtype="int32"
                )
                idx.desc.set_need_check_feed(False)

                y = paddle.gather_nd(x, idx)

                exe = paddle.static.Executor(place)
                res = exe.run(
                    paddle.static.default_main_program(),
                    feed={"x": input, "index": index},
                    fetch_list=[y],
                )

                np.testing.assert_allclose(res[0], res_np, rtol=1e-05)

308 309 310 311 312 313
    def test_imperative(self):
        paddle.disable_static()
        input_1 = np.array([[1, 2], [3, 4], [5, 6]])
        index_1 = np.array([[1]])
        input = fluid.dygraph.to_variable(input_1)
        index = fluid.dygraph.to_variable(index_1)
314
        output = paddle.gather(input, index)
315
        output_np = output.numpy()
316 317
        expected_output = np.array([[3, 4]])
        np.testing.assert_allclose(output_np, expected_output, rtol=1e-05)
318 319 320
        paddle.enable_static()


321
if __name__ == "__main__":
H
hong 已提交
322
    paddle.enable_static()
323
    unittest.main()