• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2019 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15
16import numpy as np
17import pytest
18
19import mindspore
20import mindspore.context as context
21import mindspore.nn as nn
22from mindspore import Tensor, Parameter, ParameterTuple
23from mindspore.ops import operations as P
24from mindspore.ops import composite as C
25
26
27class NetIndexAdd(nn.Cell):
28    def __init__(self, x, axis):
29        super(NetIndexAdd, self).__init__()
30        self.input_x = Parameter(Tensor(x), name='x')
31        self.index_add = P.IndexAdd(axis)
32
33    def construct(self, idx, y):
34        z = self.index_add(self.input_x, idx, y)
35        return z
36
37
38@pytest.mark.level0
39@pytest.mark.platform_x86_gpu_training
40@pytest.mark.env_onecard
41def test_index_add():
42    x = np.arange(2 * 3 * 4 * 4).reshape(2, 3, 4, 4).astype(np.float32)
43    y0 = np.ones((1, 3, 4, 4), dtype=np.float32)
44    idx0 = np.array([1]).astype(np.int32)
45    axis0 = 0
46    expect = np.copy(x)
47    expect[idx0, :, :, :] = expect[idx0, :, :, :] + y0
48    context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
49    net = NetIndexAdd(x, axis0)
50    output = net(Tensor(idx0), Tensor(y0))
51    assert (output.asnumpy() == expect).all()
52    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
53    net = NetIndexAdd(x, axis0)
54    output = net(Tensor(idx0), Tensor(y0))
55    assert (output.asnumpy() == expect).all()
56
57    y1 = np.ndarray((2, 2, 4, 4)).astype(np.float32)
58    y1.fill(0.1)
59    idx1 = np.array([0, 2]).astype(np.int32)
60    axis1 = 1
61    expect = np.copy(x)
62    expect[:, idx1, :, :] = expect[:, idx1, :, :] + y1
63    context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
64    net = NetIndexAdd(x, axis1)
65    output = net(Tensor(idx1), Tensor(y1))
66    assert (output.asnumpy() == expect).all()
67    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
68    net = NetIndexAdd(x, axis1)
69    output = net(Tensor(idx1), Tensor(y1))
70    assert (output.asnumpy() == expect).all()
71
72    y2 = np.ones((2, 3, 2, 4)).astype(np.float32)
73    y2.fill(5.5)
74    idx2 = np.array([1, 3]).astype(np.int32)
75    axis2 = 2
76    expect = np.copy(x)
77    expect[:, :, idx2, :] = expect[:, :, idx2, :] + y2
78    context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
79    net = NetIndexAdd(x, axis2)
80    output = net(Tensor(idx2), Tensor(y2))
81    assert (output.asnumpy() == expect).all()
82    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
83    net = NetIndexAdd(x, axis2)
84    output = net(Tensor(idx2), Tensor(y2))
85    assert (output.asnumpy() == expect).all()
86
87    y3 = np.ones((2, 3, 4, 3)).astype(np.float32)
88    y3.fill(1000.00)
89    idx3 = np.array([0, 2, 3]).astype(np.int32)
90    axis3 = 3
91    expect = np.copy(x)
92    expect[:, :, :, idx3] = expect[:, :, :, idx3] + y3
93    context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
94    net = NetIndexAdd(x, axis3)
95    output = net(Tensor(idx3), Tensor(y3))
96    assert (output.asnumpy() == expect).all()
97    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
98    net = NetIndexAdd(x, axis3)
99    output = net(Tensor(idx3), Tensor(y3))
100    assert (output.asnumpy() == expect).all()
101
102
103@pytest.mark.level0
104@pytest.mark.platform_x86_gpu_training
105@pytest.mark.env_onecard
106def test_index_add_float16():
107    x = np.arange(2 * 3 * 4).reshape(2, 3, 4).astype(np.float16)
108    y = np.ones((2, 2, 4), dtype=np.float16)
109    idx = np.array([0, 2]).astype(np.int32)
110    axis = 1
111    expect = np.copy(x)
112    expect[:, idx, :] = expect[:, idx, :] + y
113    context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
114    net = NetIndexAdd(x, axis)
115    output = net(Tensor(idx), Tensor(y))
116    assert (output.asnumpy() == expect).all()
117    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
118    net = NetIndexAdd(x, axis)
119    output = net(Tensor(idx), Tensor(y))
120    assert (output.asnumpy() == expect).all()
121
122
123@pytest.mark.level0
124@pytest.mark.platform_x86_gpu_training
125@pytest.mark.env_onecard
126def test_index_add_int32():
127    x = np.arange(2 * 3 * 4).reshape(2, 3, 4).astype(np.int32)
128    y = np.ones((2, 2, 4), dtype=np.int32)
129    idx = np.array([0, 2]).astype(np.int32)
130    axis = 1
131    expect = np.copy(x)
132    expect[:, idx, :] = expect[:, idx, :] + y
133    context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
134    net = NetIndexAdd(x, axis)
135    output = net(Tensor(idx), Tensor(y))
136    assert (output.asnumpy() == expect).all()
137    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
138    net = NetIndexAdd(x, axis)
139    output = net(Tensor(idx), Tensor(y))
140    assert (output.asnumpy() == expect).all()
141
142
143@pytest.mark.level1
144@pytest.mark.platform_x86_gpu_training
145@pytest.mark.env_onecard
146def test_index_add_int8():
147    x = np.arange(2 * 3 * 4).reshape(2, 3, 4).astype(np.int8)
148    y = np.ones((2, 2, 4), dtype=np.int8)
149    idx = np.array([0, 2]).astype(np.int32)
150    axis = 1
151    expect = np.copy(x)
152    expect[:, idx, :] = expect[:, idx, :] + y
153    context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
154    net = NetIndexAdd(x, axis)
155    output = net(Tensor(idx), Tensor(y))
156    assert (output.asnumpy() == expect).all()
157    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
158    net = NetIndexAdd(x, axis)
159    output = net(Tensor(idx), Tensor(y))
160    assert (output.asnumpy() == expect).all()
161
162
163@pytest.mark.level1
164@pytest.mark.platform_x86_gpu_training
165@pytest.mark.env_onecard
166def test_index_add_uint8():
167    x = np.arange(2 * 3 * 4).reshape(2, 3, 4).astype(np.uint8)
168    y = np.ones((2, 2, 4), dtype=np.uint8)
169    idx = np.array([0, 2]).astype(np.int32)
170    axis = 1
171    expect = np.copy(x)
172    expect[:, idx, :] = expect[:, idx, :] + y
173    context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
174    net = NetIndexAdd(x, axis)
175    output = net(Tensor(idx), Tensor(y))
176    assert (output.asnumpy() == expect).all()
177    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
178    net = NetIndexAdd(x, axis)
179    output = net(Tensor(idx), Tensor(y))
180    assert (output.asnumpy() == expect).all()
181
182
183@pytest.mark.level0
184@pytest.mark.platform_x86_gpu_training
185@pytest.mark.env_onecard
186def test_index_add_float64():
187    x = np.arange(2 * 3 * 4).reshape(2, 3, 4).astype(np.float64)
188    y = np.ones((2, 2, 4), dtype=np.float64)
189    idx = np.array([0, 2]).astype(np.int32)
190    axis = 1
191    expect = np.copy(x)
192    expect[:, idx, :] = expect[:, idx, :] + y
193    context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
194    net = NetIndexAdd(x, axis)
195    output = net(Tensor(idx), Tensor(y))
196    assert (output.asnumpy() == expect).all()
197    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
198    net = NetIndexAdd(x, axis)
199    output = net(Tensor(idx), Tensor(y))
200    assert (output.asnumpy() == expect).all()
201
202
203@pytest.mark.level0
204@pytest.mark.platform_x86_gpu_training
205@pytest.mark.env_onecard
206def test_index_add_int16():
207    x = np.arange(2 * 3 * 4).reshape(2, 3, 4).astype(np.int16)
208    y = np.ones((2, 2, 4), dtype=np.int16)
209    idx = np.array([0, 2]).astype(np.int32)
210    axis = 1
211    expect = np.copy(x)
212    expect[:, idx, :] = expect[:, idx, :] + y
213    context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
214    net = NetIndexAdd(x, axis)
215    output = net(Tensor(idx), Tensor(y))
216    assert (output.asnumpy() == expect).all()
217    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
218    net = NetIndexAdd(x, axis)
219    output = net(Tensor(idx), Tensor(y))
220    assert (output.asnumpy() == expect).all()
221
222
223@pytest.mark.level0
224@pytest.mark.platform_x86_gpu_training
225@pytest.mark.env_onecard
226def test_index_add_invalid_inputs():
227    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
228    x = np.arange(2 * 3 * 4).reshape(2, 3, 4).astype(np.uint8)
229    y = np.ones((2, 2, 4), dtype=np.uint8)
230    with pytest.raises(TypeError):
231        #axis not int
232        net = NetIndexAdd(x, 1.0)
233
234        #x and y don't have the same type
235        y = np.ones((2, 2, 4), dtype=np.float32)
236        idx = np.array([0, 1]).astype(np.int32)
237        net = NetIndexAdd(x, 1)
238        _ = net(Tensor(idx), Tensor(y))
239
240    with pytest.raises(ValueError):
241        #index size not the same as len(y[axis])
242        idx = np.array([0]).astype(np.int32)
243        net = NetIndexAdd(x, 1)
244        _ = net(Tensor(idx), Tensor(y))
245
246        #x and y don't have same rank
247        y = np.ones((2, 2), dtype=np.uint8)
248        idx = np.array([0, 1]).astype(np.int32)
249        net = NetIndexAdd(x, 1)
250        _ = net(Tensor(idx), Tensor(y))
251
252        #x and y don't have same shape on dimensions other than axis-th dimension
253        y = np.ones((2, 2, 5), dtype=np.uint8)
254        idx = np.array([0, 1]).astype(np.int32)
255        net = NetIndexAdd(x, 1)
256        _ = net(Tensor(idx), Tensor(y))
257
258
259class IndexAddGradNet(nn.Cell):
260    def __init__(self, network):
261        super(IndexAddGradNet, self).__init__()
262        self.grad = C.GradOperation(get_all=True, sens_param=True, get_by_list=True)
263        self.network = network
264        self.params = ParameterTuple(network.trainable_params())
265
266    def construct(self, idx, y, dout):
267        out = self.grad(self.network, self.params)(idx, y, dout)
268        return out
269
270
271def index_add_grad_with_type(nptype):
272    x = np.arange(15).reshape(5, 3).astype(nptype)
273    net = NetIndexAdd(x, 1)
274    grad_net = IndexAddGradNet(net)
275    y = Tensor(np.arange(5).reshape(5, 1).astype(nptype))
276    dout = Tensor(np.array([[63., 64., 65.],
277                            [66., 67., 68.],
278                            [69., 70., 71.],
279                            [72., 73., 74.],
280                            [75., 76., 77.]]).astype(nptype))
281    index = Tensor(np.array([1]), dtype=mindspore.int32)
282    output = grad_net(index, y, dout)
283    ygrad = output[0][1]
284    xgrad = output[1][0]
285    expect_xgrad = np.array([[63., 64., 65.],
286                             [66., 67., 68.],
287                             [69., 70., 71.],
288                             [72., 73., 74.],
289                             [75., 76., 77.]]).astype(nptype)
290    expect_ygrad = np.array([[64.],
291                             [67.],
292                             [70.],
293                             [73.],
294                             [76.]]).astype(nptype)
295    np.testing.assert_array_equal(xgrad.asnumpy(), expect_xgrad)
296    np.testing.assert_array_equal(ygrad.asnumpy(), expect_ygrad)
297
298
299@pytest.mark.level0
300@pytest.mark.platform_x86_gpu_training
301@pytest.mark.env_onecard
302def test_index_add_grad_float64():
303    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
304    index_add_grad_with_type(np.float64)
305    context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
306    index_add_grad_with_type(np.float64)
307
308
309@pytest.mark.level0
310@pytest.mark.platform_x86_gpu_training
311@pytest.mark.env_onecard
312def test_index_add_grad_float32():
313    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
314    index_add_grad_with_type(np.float32)
315    context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
316    index_add_grad_with_type(np.float32)
317
318
319@pytest.mark.level1
320@pytest.mark.platform_x86_gpu_training
321@pytest.mark.env_onecard
322def test_index_add_grad_float16():
323    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
324    index_add_grad_with_type(np.float16)
325    context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
326    index_add_grad_with_type(np.float16)
327
328
329@pytest.mark.level1
330@pytest.mark.platform_x86_gpu_training
331@pytest.mark.env_onecard
332def test_index_add_grad_int32():
333    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
334    index_add_grad_with_type(np.int32)
335    context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
336    index_add_grad_with_type(np.int32)
337
338
339@pytest.mark.level0
340@pytest.mark.platform_x86_gpu_training
341@pytest.mark.env_onecard
342def test_index_add_grad_int16():
343    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
344    index_add_grad_with_type(np.int16)
345    context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
346    index_add_grad_with_type(np.int16)
347
348
349@pytest.mark.level1
350@pytest.mark.platform_x86_gpu_training
351@pytest.mark.env_onecard
352def test_index_add_grad_int8():
353    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
354    index_add_grad_with_type(np.int8)
355    context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
356    index_add_grad_with_type(np.int8)
357
358
359@pytest.mark.level1
360@pytest.mark.platform_x86_gpu_training
361@pytest.mark.env_onecard
362def test_index_add_grad_uint8():
363    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
364    index_add_grad_with_type(np.uint8)
365    context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
366    index_add_grad_with_type(np.uint8)
367