Huawei_Technology
/
mindspore
Not watched
1
0
Fork 0
Code
Releases 13 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: bc20de741a
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'bc20de741a'
${ noResults }
mindspore/tests/ut/python/parallel/test_arithmetic.py

528 lines
18 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. import numpy as np

  16. 

  17. import mindspore as ms

  18. import mindspore.nn as nn

  19. from mindspore import Parameter, Tensor, context

  20. from mindspore.common.api import _executor

  21. from mindspore.ops import composite as C

  22. from mindspore.ops import operations as P

  23. from tests.ut.python.ops.test_math_ops import VirtualLoss

  24. 

  25. 

  26. class NetWithLoss(nn.Cell):

  27.     def __init__(self, network):

  28.         super(NetWithLoss, self).__init__()

  29.         self.loss = VirtualLoss()

  30.         self.network = network

  31. 

  32.     def construct(self, x, y, b):

  33.         predict = self.network(x, y, b)

  34.         return self.loss(predict)

  35. 

  36. 

  37. class GradWrap(nn.Cell):

  38.     def __init__(self, network):

  39.         super(GradWrap, self).__init__()

  40.         self.network = network

  41. 

  42.     def construct(self, x, y, b):

  43.         return C.grad_all(self.network)(x, y, b)

  44. 

  45. 

  46. def compile_net(net, x, y, b):

  47.     net.set_auto_parallel()

  48.     _executor.compile(net, x, y, b)

  49. 

  50. 

  51. def test_matmul_sub():

  52.     class Net(nn.Cell):

  53.         def __init__(self, strategy1, strategy2):

  54.             super().__init__()

  55.             self.matmul = P.MatMul().set_strategy(strategy1)

  56.             self.sub = P.Sub().set_strategy(strategy2)

  57. 

  58.         def construct(self, x, y, b):

  59.             out = self.matmul(x, y)

  60.             out = self.sub(out, b)

  61.             return out

  62. 

  63.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  64.     strategy1 = ((2, 2), (2, 2))

  65.     strategy2 = ((4, 2), (4, 2))

  66.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  67.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  68. 

  69.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  70.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  71.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  72.     compile_net(net, x, y, b)

  73. 

  74. 

  75. def test_matmul_add():

  76.     class Net(nn.Cell):

  77.         def __init__(self, strategy1, strategy2):

  78.             super().__init__()

  79.             self.matmul = P.MatMul().set_strategy(strategy1)

  80.             self.add = P.TensorAdd().set_strategy(strategy2)

  81. 

  82.         def construct(self, x, y, b):

  83.             out = self.matmul(x, y)

  84.             out = self.add(out, b)

  85.             return out

  86. 

  87.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  88.     strategy1 = ((2, 2), (2, 2))

  89.     strategy2 = ((4, 2), (4, 2))

  90.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  91.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  92. 

  93.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  94.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  95.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  96.     compile_net(net, x, y, b)

  97. 

  98. 

  99. def test_matmul_mul():

  100.     class Net(nn.Cell):

  101.         def __init__(self, strategy1, strategy2):

  102.             super().__init__()

  103.             self.matmul = P.MatMul().set_strategy(strategy1)

  104.             self.mul = P.Mul().set_strategy(strategy2)

  105. 

  106.         def construct(self, x, y, b):

  107.             out = self.matmul(x, y)

  108.             out = self.mul(out, b)

  109.             return out

  110. 

  111.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  112.     strategy1 = ((2, 2), (2, 2))

  113.     strategy2 = ((4, 2), (4, 2))

  114.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  115.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  116. 

  117.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  118.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  119.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  120.     compile_net(net, x, y, b)

  121. 

  122. 

  123. def test_matmul_div():

  124.     class Net(nn.Cell):

  125.         def __init__(self, strategy1, strategy2):

  126.             super().__init__()

  127.             self.matmul = P.MatMul().set_strategy(strategy1)

  128.             self.div = P.Div().set_strategy(strategy2)

  129. 

  130.         def construct(self, x, y, b):

  131.             out = self.matmul(x, y)

  132.             out = self.div(out, b)

  133.             return out

  134. 

  135.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  136.     strategy1 = ((2, 2), (2, 2))

  137.     strategy2 = ((4, 2), (4, 2))

  138.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  139.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  140. 

  141.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  142.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  143.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  144.     compile_net(net, x, y, b)

  145. 

  146. 

  147. def test_matmul_greater():

  148.     class Net(nn.Cell):

  149.         def __init__(self, strategy1, strategy2):

  150.             super().__init__()

  151.             self.matmul = P.MatMul().set_strategy(strategy1)

  152.             self.greater = P.Greater().set_strategy(strategy2)

  153. 

  154.         def construct(self, x, y, b):

  155.             out = self.matmul(x, y)

  156.             out = self.greater(out, b)

  157.             return out

  158. 

  159.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  160.     strategy1 = ((2, 2), (2, 2))

  161.     strategy2 = ((4, 2), (4, 2))

  162.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  163.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  164. 

  165.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  166.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  167.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  168.     compile_net(net, x, y, b)

  169. 

  170. 

  171. def test_matmul_add_broadcast():

  172.     class Net(nn.Cell):

  173.         def __init__(self, strategy1, strategy2):

  174.             super().__init__()

  175.             self.matmul = P.MatMul().set_strategy(strategy1)

  176.             self.add = P.TensorAdd().set_strategy(strategy2)

  177. 

  178.         def construct(self, x, y, b):

  179.             out = self.matmul(x, y)

  180.             out = self.add(out, b)

  181.             return out

  182. 

  183.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  184.     strategy1 = ((2, 2), (2, 2))

  185.     strategy2 = ((4, 2), (2,))

  186.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  187.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  188. 

  189.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  190.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  191.     b = Tensor(np.ones([64]), dtype=ms.float32)

  192.     compile_net(net, x, y, b)

  193. 

  194. 

  195. def test_matmul_add_broadcast2():

  196.     class Net(nn.Cell):

  197.         def __init__(self, strategy1, strategy2):

  198.             super().__init__()

  199.             self.matmul = P.MatMul().set_strategy(strategy1)

  200.             self.add = P.TensorAdd().set_strategy(strategy2)

  201. 

  202.         def construct(self, x, y, b):

  203.             out = self.matmul(x, y)

  204.             out = self.add(out, b)

  205.             return out

  206. 

  207.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  208.     strategy1 = ((2, 4), (4, 1))

  209.     strategy2 = ((4, 1), (1, 2))

  210.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  211.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  212. 

  213.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  214.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  215.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  216.     compile_net(net, x, y, b)

  217. 

  218. 

  219. def test_matmul_sub_broadcast():

  220.     class Net(nn.Cell):

  221.         def __init__(self, strategy1, strategy2):

  222.             super().__init__()

  223.             self.matmul = P.MatMul().set_strategy(strategy1)

  224.             self.sub = P.Sub().set_strategy(strategy2)

  225. 

  226.         def construct(self, x, y, b):

  227.             out = self.matmul(x, y)

  228.             out = self.sub(out, b)

  229.             return out

  230. 

  231.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  232.     strategy1 = ((2, 2), (2, 2))

  233.     strategy2 = ((4, 2), (2,))

  234.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  235.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  236. 

  237.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  238.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  239.     b = Tensor(np.ones([64]), dtype=ms.float32)

  240.     compile_net(net, x, y, b)

  241. 

  242. 

  243. def test_matmul_sub_broadcast2():

  244.     class Net(nn.Cell):

  245.         def __init__(self, strategy1, strategy2):

  246.             super().__init__()

  247.             self.matmul = P.MatMul().set_strategy(strategy1)

  248.             self.sub = P.Sub().set_strategy(strategy2)

  249. 

  250.         def construct(self, x, y, b):

  251.             out = self.matmul(x, y)

  252.             out = self.sub(out, b)

  253.             return out

  254. 

  255.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  256.     strategy1 = ((2, 4), (4, 1))

  257.     strategy2 = ((4, 1), (1, 2))

  258.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  259.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  260. 

  261.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  262.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  263.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  264.     compile_net(net, x, y, b)

  265. 

  266. 

  267. def test_matmul_mul_broadcast():

  268.     class Net(nn.Cell):

  269.         def __init__(self, strategy1, strategy2):

  270.             super().__init__()

  271.             self.matmul = P.MatMul().set_strategy(strategy1)

  272.             self.mul = P.Mul().set_strategy(strategy2)

  273. 

  274.         def construct(self, x, y, b):

  275.             out = self.matmul(x, y)

  276.             out = self.mul(out, b)

  277.             return out

  278. 

  279.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  280.     strategy1 = ((2, 2), (2, 2))

  281.     strategy2 = ((4, 2), (2,))

  282.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  283.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  284. 

  285.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  286.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  287.     b = Tensor(np.ones([64]), dtype=ms.float32)

  288.     compile_net(net, x, y, b)

  289. 

  290. 

  291. def test_matmul_mul_broadcast2():

  292.     class Net(nn.Cell):

  293.         def __init__(self, strategy1, strategy2):

  294.             super().__init__()

  295.             self.matmul = P.MatMul().set_strategy(strategy1)

  296.             self.mul = P.Mul().set_strategy(strategy2)

  297. 

  298.         def construct(self, x, y, b):

  299.             out = self.matmul(x, y)

  300.             out = self.mul(out, b)

  301.             return out

  302. 

  303.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  304.     strategy1 = ((2, 4), (4, 1))

  305.     strategy2 = ((4, 1), (1, 2))

  306.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  307.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  308. 

  309.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  310.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  311.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  312.     compile_net(net, x, y, b)

  313. 

  314. 

  315. def test_matmul_div_broadcast():

  316.     class Net(nn.Cell):

  317.         def __init__(self, strategy1, strategy2):

  318.             super().__init__()

  319.             self.matmul = P.MatMul().set_strategy(strategy1)

  320.             self.div = P.Div().set_strategy(strategy2)

  321. 

  322.         def construct(self, x, y, b):

  323.             out = self.matmul(x, y)

  324.             out = self.div(out, b)

  325.             return out

  326. 

  327.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  328.     strategy1 = ((2, 2), (2, 2))

  329.     strategy2 = ((4, 2), (2,))

  330.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  331.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  332. 

  333.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  334.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  335.     b = Tensor(np.ones([64]), dtype=ms.float32)

  336.     compile_net(net, x, y, b)

  337. 

  338. 

  339. def test_matmul_div_broadcast2():

  340.     class Net(nn.Cell):

  341.         def __init__(self, strategy1, strategy2):

  342.             super().__init__()

  343.             self.matmul = P.MatMul().set_strategy(strategy1)

  344.             self.div = P.Div().set_strategy(strategy2)

  345. 

  346.         def construct(self, x, y, b):

  347.             out = self.matmul(x, y)

  348.             out = self.div(out, b)

  349.             return out

  350. 

  351.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  352.     strategy1 = ((2, 4), (4, 1))

  353.     strategy2 = ((4, 1), (1, 2))

  354.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  355.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  356. 

  357.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  358.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  359.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  360.     compile_net(net, x, y, b)

  361. 

  362. 

  363. def test_matmul_greater_broadcast():

  364.     class Net(nn.Cell):

  365.         def __init__(self, strategy1, strategy2):

  366.             super().__init__()

  367.             self.matmul = P.MatMul().set_strategy(strategy1)

  368.             self.greater = P.Greater().set_strategy(strategy2)

  369. 

  370.         def construct(self, x, y, b):

  371.             out = self.matmul(x, y)

  372.             out = self.greater(out, b)

  373.             return out

  374. 

  375.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  376.     strategy1 = ((2, 2), (2, 2))

  377.     strategy2 = ((4, 2), (2,))

  378.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  379.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  380. 

  381.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  382.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  383.     b = Tensor(np.ones([64]), dtype=ms.float32)

  384.     compile_net(net, x, y, b)

  385. 

  386. 

  387. def test_matmul_greater_broadcast2():

  388.     class Net(nn.Cell):

  389.         def __init__(self, strategy1, strategy2):

  390.             super().__init__()

  391.             self.matmul = P.MatMul().set_strategy(strategy1)

  392.             self.greater = P.Greater().set_strategy(strategy2)

  393. 

  394.         def construct(self, x, y, b):

  395.             out = self.matmul(x, y)

  396.             out = self.greater(out, b)

  397.             return out

  398. 

  399.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  400.     strategy1 = ((2, 4), (4, 1))

  401.     strategy2 = ((4, 1), (1, 2))

  402.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  403.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  404. 

  405.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  406.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  407.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  408.     compile_net(net, x, y, b)

  409. 

  410. 

  411. def test_matmul_floordiv():

  412.     class Net(nn.Cell):

  413.         def __init__(self, strategy1, strategy2):

  414.             super().__init__()

  415.             self.matmul = P.MatMul().set_strategy(strategy1)

  416.             self.floordiv = P.FloorDiv().set_strategy(strategy2)

  417. 

  418.         def construct(self, x, y, b):

  419.             out = self.matmul(x, y)

  420.             out = self.floordiv(out, b)

  421.             return out

  422. 

  423.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  424.     strategy1 = ((2, 2), (2, 2))

  425.     strategy2 = ((4, 2), (4, 2))

  426.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  427.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  428. 

  429.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  430.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  431.     b = Tensor(np.ones([64, 64]), dtype=ms.float32)

  432.     compile_net(net, x, y, b)

  433. 

  434. 

  435. def test_matmul_floordiv_broadcast():

  436.     class Net(nn.Cell):

  437.         def __init__(self, strategy1, strategy2):

  438.             super().__init__()

  439.             self.matmul = P.MatMul().set_strategy(strategy1)

  440.             self.floordiv = P.FloorDiv().set_strategy(strategy2)

  441. 

  442.         def construct(self, x, y, b):

  443.             out = self.matmul(x, y)

  444.             out = self.floordiv(out, b)

  445.             return out

  446. 

  447.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  448.     strategy1 = ((2, 2), (2, 2))

  449.     strategy2 = ((4, 2), (2,))

  450.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  451.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  452. 

  453.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  454.     y = Tensor(np.ones([32, 64]), dtype=ms.float32)

  455.     b = Tensor(np.ones([64]), dtype=ms.float32)

  456.     compile_net(net, x, y, b)

  457. 

  458. 

  459. def test_matmul_floordiv_broadcast2():

  460.     class Net(nn.Cell):

  461.         def __init__(self, strategy1, strategy2):

  462.             super().__init__()

  463.             self.matmul = P.MatMul().set_strategy(strategy1)

  464.             self.floordiv = P.FloorDiv().set_strategy(strategy2)

  465. 

  466.         def construct(self, x, y, b):

  467.             out = self.matmul(x, y)

  468.             out = self.floordiv(out, b)

  469.             return out

  470. 

  471.     context.set_auto_parallel_context(device_num=8, global_rank=0)

  472.     strategy1 = ((2, 4), (4, 1))

  473.     strategy2 = ((4, 1), (1, 2))

  474.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2)))

  475.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  476. 

  477.     x = Tensor(np.ones([64, 32]), dtype=ms.float32)

  478.     y = Tensor(np.ones([32, 1]), dtype=ms.float32)

  479.     b = Tensor(np.ones([1, 64]), dtype=ms.float32)

  480.     compile_net(net, x, y, b)

  481. 

  482. 

  483. def test_assign_sub():

  484.     class Net(nn.Cell):

  485.         def __init__(self):

  486.             super().__init__()

  487.             self.assign_sub = P.AssignSub()

  488.             self.mul = P.Mul()

  489.             self.mul_weight = Parameter(Tensor(np.full([128, 32],

  490.                                                        0.5, dtype=np.float32)),

  491.                                         name="mul_weight")

  492.             self.assignsub_weight = Parameter(Tensor(np.full([128, 32],

  493.                                                              1.1, dtype=np.float32)),

  494.                                               name="assignsub_weight")

  495. 

  496.         def construct(self, x):

  497.             out = self.mul(x, self.mul_weight)

  498.             out = self.assign_sub(self.assignsub_weight, out)

  499.             return out

  500. 

  501.     class SubNetWithLoss(nn.Cell):

  502.         def __init__(self, network):

  503.             super(SubNetWithLoss, self).__init__()

  504.             self.loss = VirtualLoss()

  505.             self.network = network

  506. 

  507.         def construct(self, x):

  508.             predict = self.network(x,)

  509.             return self.loss(predict)

  510. 

  511.     class SubGradWrap(nn.Cell):

  512.         def __init__(self, network):

  513.             super(SubGradWrap, self).__init__()

  514.             self.network = network

  515. 

  516.         def construct(self, x):

  517.             return C.grad_all(self.network)(x)

  518. 

  519.     def compile_sub_net(net, x):

  520.         net.set_auto_parallel()

  521.         _executor.compile(net, x)

  522. 

  523.     context.set_auto_parallel_context(device_num=64, global_rank=15)

  524.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  525.     net = SubGradWrap(SubNetWithLoss(Net()))

  526.     x = Tensor(np.ones([128, 32]), dtype=ms.float32)

  527.     compile_sub_net(net, x)