zzy34407230
/
mindspore2022
Not watched
1
0
Fork 0
Code
Releases 22 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: a64920c468
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'a64920c468'
${ noResults }
mindspore2022/tests/st/ops/gpu/test_index_add_op.py

367 lines
13 kB
Raw Blame History 

  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. 

  16. import numpy as np

  17. import pytest

  18. 

  19. import mindspore

  20. import mindspore.context as context

  21. import mindspore.nn as nn

  22. from mindspore import Tensor, Parameter, ParameterTuple

  23. from mindspore.ops import operations as P

  24. from mindspore.ops import composite as C

  25. 

  26. 

  27. class 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

  41. def 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

  106. def 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

  126. def 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

  146. def 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

  166. def 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

  186. def 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

  206. def 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

  226. def 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. 

  259. class 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. 

  271. def 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

  302. def 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

  312. def 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

  322. def 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

  332. def 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

  342. def 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

  352. def 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

  362. def 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)