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: 9cee0d2143
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from '9cee0d2143'
${ noResults }
mindspore2022/tests/ut/python/parallel/test_reduce_method_info.py

586 lines
21 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 Tensor

  20. from mindspore import context

  21. from mindspore.common.api import _executor

  22. from mindspore.ops import composite as C

  23. from mindspore.ops import operations as P

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

  25. 

  26. 

  27. grad_all = C.GradOperation(get_all=True)

  28. 

  29. 

  30. class NetWithLossNoBias(nn.Cell):

  31.     def __init__(self, network):

  32.         super(NetWithLossNoBias, self).__init__()

  33.         self.loss = VirtualLoss()

  34.         self.network = network

  35. 

  36.     def construct(self, x, y):

  37.         predict = self.network(x, y)

  38.         return self.loss(predict)

  39. 

  40. 

  41. class NetWithLoss(nn.Cell):

  42.     def __init__(self, network):

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

  44.         self.loss = VirtualLoss()

  45.         self.network = network

  46. 

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

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

  49.         return self.loss(predict)

  50. 

  51. 

  52. class GradWrapNoBias(nn.Cell):

  53.     def __init__(self, network):

  54.         super(GradWrapNoBias, self).__init__()

  55.         self.network = network

  56. 

  57.     def construct(self, x, y):

  58.         return grad_all(self.network)(x, y)

  59. 

  60. 

  61. class GradWrap(nn.Cell):

  62.     def __init__(self, network):

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

  64.         self.network = network

  65. 

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

  67.         return grad_all(self.network)(x, y, b)

  68. 

  69. 

  70. def compile_net_no_bias(net, x, y):

  71.     net.set_auto_parallel()

  72.     _executor.compile(net, x, y)

  73. 

  74. 

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

  76.     net.set_auto_parallel()

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

  78. 

  79. 

  80. # model_parallel test

  81. def test_sum_mul():

  82.     class Net(nn.Cell):

  83.         def __init__(self, strategy1, strategy2, strategy3):

  84.             super().__init__()

  85.             self.mul1 = P.Mul().shard(strategy1)

  86.             self.reduce_sum = P.ReduceSum(keep_dims=False).shard(strategy2)

  87.             self.mul2 = P.Mul().shard(strategy3)

  88. 

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

  90.             out = self.mul1(x, y)

  91.             out = self.reduce_sum(out, (1,))

  92.             out = self.mul2(out, b)

  93.             return out

  94. 

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

  96.     strategy1 = ((1, 1, 8), (1, 1, 8))

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

  98.     strategy3 = ((2, 4), (2, 4))

  99.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2, strategy3)))

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

  101. 

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

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

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

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

  106. 

  107. 

  108. def test_sum_mul2():

  109.     class Net(nn.Cell):

  110.         def __init__(self, strategy1, strategy2, strategy3):

  111.             super().__init__()

  112.             self.mul1 = P.Mul().shard(strategy1)

  113.             self.reduce_sum = P.ReduceSum(keep_dims=False).shard(strategy2)

  114.             self.mul2 = P.Mul().shard(strategy3)

  115. 

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

  117.             out = self.mul1(x, y)

  118.             out = self.reduce_sum(out, (0, 1))

  119.             out = self.mul2(out, b)

  120.             return out

  121. 

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

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

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

  125.     strategy3 = ((2, 4), (2, 4))

  126.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2, strategy3)))

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

  128. 

  129.     x = Tensor(np.ones([128, 128, 64, 64]), dtype=ms.float32)

  130.     y = Tensor(np.ones([128, 128, 64, 64]), dtype=ms.float32)

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

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

  133. 

  134. 

  135. def test_sum_mul3():

  136.     class Net(nn.Cell):

  137.         def __init__(self, strategy1, strategy2, strategy3):

  138.             super().__init__()

  139.             self.mul1 = P.Mul().shard(strategy1)

  140.             self.reduce_sum = P.ReduceSum(keep_dims=False).shard(strategy2)

  141.             self.mul2 = P.Mul().shard(strategy3)

  142. 

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

  144.             out = self.mul1(x, y)

  145.             out = self.reduce_sum(out, -1)

  146.             out = self.mul2(out, b)

  147.             return out

  148. 

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

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

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

  152.     strategy3 = ((2, 4), (2, 4))

  153.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2, strategy3)))

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

  155. 

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

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

  158.     b = Tensor(np.ones([128, 32]), dtype=ms.float32)

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

  160. 

  161. 

  162. def test_sum_mul4():

  163.     class Net(nn.Cell):

  164.         def __init__(self, strategy1, strategy2, strategy3):

  165.             super().__init__()

  166.             self.mul1 = P.Mul().shard(strategy1)

  167.             self.reduce_sum = P.ReduceSum(keep_dims=True).shard(strategy2)

  168.             self.mul2 = P.Mul().shard(strategy3)

  169. 

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

  171.             out = self.mul1(x, y)

  172.             out = self.reduce_sum(out, -1)

  173.             out = self.mul2(out, b)

  174.             return out

  175. 

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

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

  178.     strategy2 = ((2, 2, 2),)

  179.     strategy3 = ((4, 2, 1), (4, 2, 1))

  180.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2, strategy3)))

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

  182. 

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

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

  185.     b = Tensor(np.ones([128, 32, 1]), dtype=ms.float32)

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

  187. 

  188. 

  189. def test_sum_mul5():

  190.     class Net(nn.Cell):

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

  192.             super().__init__()

  193.             self.mul1 = P.Mul().shard(strategy1)

  194.             self.reduce_sum = P.ReduceSum(keep_dims=True).shard(strategy2)

  195. 

  196.         def construct(self, x, y):

  197.             out = self.mul1(x, y)

  198.             out = self.reduce_sum(out, 0)

  199.             return out

  200. 

  201.     context.set_auto_parallel_context(device_num=64, global_rank=0)

  202.     strategy1 = ((1, 8, 8), (1, 8, 8))

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

  204.     net = GradWrapNoBias(NetWithLossNoBias(Net(strategy1, strategy2)))

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

  206. 

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

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

  209.     compile_net_no_bias(net, x, y)

  210. 

  211. 

  212. def test_sum_mul6():

  213.     class Net(nn.Cell):

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

  215.             super().__init__()

  216.             self.mul1 = P.Mul().shard(strategy1)

  217.             self.reduce_sum = P.ReduceSum(keep_dims=True).shard(strategy2)

  218. 

  219.         def construct(self, x, y):

  220.             out = self.mul1(x, y)

  221.             out = self.reduce_sum(out, 1)

  222.             return out

  223. 

  224.     context.set_auto_parallel_context(device_num=64, global_rank=0)

  225.     strategy1 = ((1, 8, 8), (1, 8, 8))

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

  227.     net = GradWrapNoBias(NetWithLossNoBias(Net(strategy1, strategy2)))

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

  229. 

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

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

  232.     compile_net_no_bias(net, x, y)

  233. 

  234. 

  235. def test_sum_mul7():

  236.     class Net(nn.Cell):

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

  238.             super().__init__()

  239.             self.mul1 = P.Mul().shard(strategy1)

  240.             self.reduce_sum = P.ReduceSum(keep_dims=True).shard(strategy2)

  241. 

  242.         def construct(self, x, y):

  243.             out = self.mul1(x, y)

  244.             out = self.reduce_sum(out, (0, 1))

  245.             return out

  246. 

  247.     context.set_auto_parallel_context(device_num=64, global_rank=0)

  248.     strategy1 = ((1, 8, 8), (1, 8, 8))

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

  250.     net = GradWrapNoBias(NetWithLossNoBias(Net(strategy1, strategy2)))

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

  252. 

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

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

  255.     compile_net_no_bias(net, x, y)

  256. 

  257. 

  258. def test_max_mul():

  259.     class Net(nn.Cell):

  260.         def __init__(self, strategy1, strategy2, strategy3):

  261.             super().__init__()

  262.             self.mul1 = P.Mul().shard(strategy1)

  263.             self.reduce_max = P.ReduceMax(keep_dims=False).shard(strategy2)

  264.             self.mul2 = P.Mul().shard(strategy3)

  265. 

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

  267.             out = self.mul1(x, y)

  268.             out = self.reduce_max(out, -1)

  269.             out = self.mul2(out, b)

  270.             return out

  271. 

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

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

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

  275.     strategy3 = ((2, 4), (2, 4))

  276.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2, strategy3)))

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

  278. 

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

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

  281.     b = Tensor(np.ones([128, 32]), dtype=ms.float32)

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

  283. 

  284. 

  285. def test_min_mul():

  286.     class Net(nn.Cell):

  287.         def __init__(self, strategy1, strategy2, strategy3):

  288.             super().__init__()

  289.             self.mul1 = P.Mul().shard(strategy1)

  290.             self.reduce_min = P.ReduceMin(keep_dims=False).shard(strategy2)

  291.             self.mul2 = P.Mul().shard(strategy3)

  292. 

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

  294.             out = self.mul1(x, y)

  295.             out = self.reduce_min(out, 0)

  296.             out = self.mul2(out, b)

  297.             return out

  298. 

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

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

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

  302.     strategy3 = ((2, 4), (2, 4))

  303.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2, strategy3)))

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

  305. 

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

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

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

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

  310. 

  311. 

  312. def test_reduce_mean_mul_float32():

  313.     class Net(nn.Cell):

  314.         def __init__(self, strategy1, strategy2, strategy3):

  315.             super().__init__()

  316.             self.mul1 = P.Mul().shard(strategy1)

  317.             self.reduce_mean = P.ReduceMean(keep_dims=False).shard(strategy2)

  318.             self.mul2 = P.Mul().shard(strategy3)

  319. 

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

  321.             out = self.mul1(x, y)

  322.             out = self.reduce_mean(out, 0)

  323.             out = self.mul2(out, b)

  324.             return out

  325. 

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

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

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

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

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

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

  332. 

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

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

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

  336. 

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

  338. 

  339. 

  340. class ArgMaxWithValueNet(nn.Cell):

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

  342.         super().__init__()

  343.         self.mul1 = P.Mul().shard(strategy1)

  344.         self.arg_max_with_value = P.ArgMaxWithValue(keep_dims=False, axis=-1).shard(strategy2)

  345.         self.mul2 = P.Mul().shard(strategy3)

  346. 

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

  348.         out = self.mul1(x, y)

  349.         _, out = self.arg_max_with_value(out)

  350.         out = self.mul2(out, b)

  351.         return out

  352. 

  353. 

  354. class ArgMinWithValueNet(nn.Cell):

  355.     def __init__(self, strategy1, strategy2, strategy3):

  356.         super().__init__()

  357.         self.mul1 = P.Mul().shard(strategy1)

  358.         self.arg_min_with_value = P.ArgMinWithValue(keep_dims=False, axis=-1).shard(strategy2)

  359.         self.mul2 = P.Mul().shard(strategy3)

  360. 

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

  362.         out = self.mul1(x, y)

  363.         _, out = self.arg_min_with_value(out)

  364.         out = self.mul2(out, b)

  365.         return out

  366. 

  367. 

  368. def gen_inputs_and_compile_net(net):

  369.     x = Tensor(np.ones([128, 64, 64]), dtype=ms.float32)

  370.     y = Tensor(np.ones([128, 64, 64]), dtype=ms.float32)

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

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

  373. 

  374. 

  375. def gen_inputs_and_compile_net_no_bias(net):

  376.     x = Tensor(np.ones([128, 64, 64]), dtype=ms.float32)

  377.     y = Tensor(np.ones([128, 64, 64]), dtype=ms.float32)

  378.     compile_net_no_bias(net, x, y)

  379. 

  380. 

  381. def tobefixed_test_arg_max_with_value_mul_semi_axis_parallel():

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

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

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

  385.     strategy3 = ((2, 4), (2, 4))

  386.     net = GradWrap(NetWithLoss(ArgMaxWithValueNet(strategy1, strategy2, strategy3)))

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

  388.     gen_inputs_and_compile_net(net)

  389. 

  390. 

  391. def test_arg_max_with_value_mul_semi():

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

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

  394.     strategy2 = ((4, 1, 1),)

  395.     strategy3 = ((2, 4), (2, 4))

  396.     net = GradWrap(NetWithLoss(ArgMaxWithValueNet(strategy1, strategy2, strategy3)))

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

  398.     gen_inputs_and_compile_net(net)

  399. 

  400. 

  401. def test_arg_max_with_value_mul_auto():

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

  403.     strategy1 = None

  404.     strategy2 = None

  405.     strategy3 = None

  406.     net = GradWrap(NetWithLoss(ArgMaxWithValueNet(strategy1, strategy2, strategy3)))

  407.     context.set_auto_parallel_context(parallel_mode="auto_parallel")

  408.     gen_inputs_and_compile_net(net)

  409. 

  410. 

  411. def test_arg_min_with_value_mul_semi_axis_parallel():

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

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

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

  415.     strategy3 = ((2, 4), (2, 4))

  416.     net = GradWrap(NetWithLoss(ArgMinWithValueNet(strategy1, strategy2, strategy3)))

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

  418.     gen_inputs_and_compile_net(net)

  419. 

  420. 

  421. def test_arg_min_with_value_mul_semi():

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

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

  424.     strategy2 = ((4, 1, 1),)

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

  426.     net = GradWrap(NetWithLoss(ArgMinWithValueNet(strategy1, strategy2, strategy3)))

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

  428.     gen_inputs_and_compile_net(net)

  429. 

  430. 

  431. def test_arg_min_with_value_mul_auto():

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

  433.     strategy1 = None

  434.     strategy2 = None

  435.     strategy3 = None

  436.     net = GradWrap(NetWithLoss(ArgMinWithValueNet(strategy1, strategy2, strategy3)))

  437.     context.set_auto_parallel_context(parallel_mode="auto_parallel")

  438.     gen_inputs_and_compile_net(net)

  439. 

  440. 

  441. class ArgMinWithValueNet2(nn.Cell):

  442.     def __init__(self, strategy1, strategy2, strategy3):

  443.         super().__init__()

  444.         self.mul1 = P.Mul().shard(strategy1)

  445.         self.arg_min_with_value = P.ArgMinWithValue(keep_dims=True, axis=-1).shard(strategy2)

  446.         self.relu = P.ReLU().shard(strategy3)

  447. 

  448.     def construct(self, x, y):

  449.         out = self.mul1(x, y)

  450.         _, out = self.arg_min_with_value(out)

  451.         out = self.relu(out)

  452.         return out

  453. 

  454. 

  455. def tobefixed_test_arg_min_with_value_mul_semi_axis_parallel2():

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

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

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

  459.     strategy3 = ((2, 4, 1),)

  460.     net = GradWrapNoBias(NetWithLossNoBias(ArgMinWithValueNet2(strategy1, strategy2, strategy3)))

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

  462.     gen_inputs_and_compile_net_no_bias(net)

  463. 

  464. 

  465. def test_arg_min_with_value_mul_semi2():

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

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

  468.     strategy2 = ((4, 1, 1),)

  469.     strategy3 = ((2, 4, 1),)

  470.     net = GradWrapNoBias(NetWithLossNoBias(ArgMinWithValueNet2(strategy1, strategy2, strategy3)))

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

  472.     gen_inputs_and_compile_net_no_bias(net)

  473. 

  474. 

  475. def test_arg_min_with_value_mul_auto2():

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

  477.     strategy1 = None

  478.     strategy2 = None

  479.     strategy3 = None

  480.     net = GradWrapNoBias(NetWithLossNoBias(ArgMinWithValueNet2(strategy1, strategy2, strategy3)))

  481.     context.set_auto_parallel_context(parallel_mode="auto_parallel")

  482.     gen_inputs_and_compile_net_no_bias(net)

  483. 

  484. 

  485. def test_cross_batch():

  486.     class Net(nn.Cell):

  487.         def __init__(self, strategy1, strategy2, strategy3):

  488.             super().__init__()

  489.             self.mul1 = P.Mul().shard(strategy1)

  490.             self.reduce_sum = P.ReduceSum(keep_dims=False).shard(strategy2)

  491.             self.reduce_mean = P.ReduceMean(keep_dims=False).shard(strategy3).add_prim_attr("cross_batch", True)

  492. 

  493.         def construct(self, x, y):

  494.             out = self.mul1(x, y)

  495.             out = self.reduce_sum(out, -1)

  496.             out = self.reduce_mean(out, 0)

  497.             return out

  498. 

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

  500.     strategy1 = ((4, 2), (4, 2))

  501.     strategy2 = ((2, 1),)

  502.     strategy3 = ((8,),)

  503.     net = GradWrapNoBias(NetWithLossNoBias(Net(strategy1, strategy2, strategy3)))

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

  505. 

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

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

  508.     compile_net_no_bias(net, x, y)

  509. 

  510. 

  511. def test_cross_batch2():

  512.     class Net(nn.Cell):

  513.         def __init__(self, strategy1, strategy2, strategy3):

  514.             super().__init__()

  515.             self.mul1 = P.Mul().shard(strategy1)

  516.             self.reduce_mean = P.ReduceMean(keep_dims=False).shard(strategy2)

  517.             self.reduce_sum = P.ReduceSum(keep_dims=False).shard(strategy3).add_prim_attr("cross_batch", True)

  518. 

  519.         def construct(self, x, y):

  520.             out = self.mul1(x, y)

  521.             out = self.reduce_mean(out, -1)

  522.             out = self.reduce_sum(out, 0)

  523.             return out

  524. 

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

  526.     strategy1 = ((4, 2), (4, 2))

  527.     strategy2 = ((2, 1),)

  528.     strategy3 = ((8,),)

  529.     net = GradWrapNoBias(NetWithLossNoBias(Net(strategy1, strategy2, strategy3)))

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

  531. 

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

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

  534.     compile_net_no_bias(net, x, y)

  535. 

  536. 

  537. def test_cross_batch_auto():

  538.     class Net(nn.Cell):

  539.         def __init__(self):

  540.             super().__init__()

  541.             self.mul1 = P.Mul()

  542.             self.reduce_mean = P.ReduceMean(keep_dims=False)

  543.             self.reduce_sum = P.ReduceSum(keep_dims=False).add_prim_attr("cross_batch", True)

  544. 

  545.         def construct(self, x, y):

  546.             out = self.mul1(x, y)

  547.             out = self.reduce_mean(out, -1)

  548.             out = self.reduce_sum(out, 0)

  549.             return out

  550. 

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

  552.     net = GradWrapNoBias(NetWithLossNoBias(Net()))

  553.     context.set_auto_parallel_context(parallel_mode="auto_parallel")

  554. 

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

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

  557.     compile_net_no_bias(net, x, y)

  558. 

  559. 

  560. def test_max_empty_tuple():

  561.     class Net(nn.Cell):

  562.         def __init__(self, strategy1, strategy2, strategy3):

  563.             super().__init__()

  564.             self.mul = P.Mul().shard(strategy1)

  565.             self.reduce_max = P.ReduceMax(keep_dims=False).shard(strategy2)

  566.             self.add = P.TensorAdd().shard(strategy3)

  567. 

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

  569.             out = self.mul(x, y)

  570.             out = self.reduce_max(out)

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

  572.             return out

  573. 

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

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

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

  577.     strategy3 = ((), (1, 1))

  578.     net = GradWrap(NetWithLoss(Net(strategy1, strategy2, strategy3)))

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

  580. 

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

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

  583.     b = Tensor(np.ones([128, 32]), dtype=ms.float32)

  584. 

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