parallel_ut_refactoring1

tests/ut/python/parallel/test_arithmetic.py

@@ -53,6 +53,11 @@ def compile_net(net, x, y, b):
 
 
 def test_matmul_sub():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-sub net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -77,6 +82,11 @@ def test_matmul_sub():
 
 
 def test_matmul_add():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-add net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -101,6 +111,11 @@ def test_matmul_add():
 
 
 def test_matmul_mul():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-mul net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -124,6 +139,11 @@ def test_matmul_mul():
     compile_net(net, x, y, b)
 
 def test_matmul_mod():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-mod net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -147,6 +167,11 @@ def test_matmul_mod():
     compile_net(net, x, y, b)
 
 def test_matmul_floormod():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-floormod net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -171,6 +196,11 @@ def test_matmul_floormod():
 
 
 def test_matmul_atan2():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-atan2 net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -195,6 +225,11 @@ def test_matmul_atan2():
 
 
 def test_matmul_divNoNan():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-divNoNan net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -219,6 +254,11 @@ def test_matmul_divNoNan():
 
 
 def test_matmul_logicaland():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-logical_and net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -248,6 +288,11 @@ def test_matmul_logicaland():
 
 
 def test_matmul_logicalor():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-logical_or net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -277,6 +322,11 @@ def test_matmul_logicalor():
 
 
 def test_matmul_div():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-div net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -301,6 +351,11 @@ def test_matmul_div():
 
 
 def test_matmul_add_broadcast():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-add broadcast net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -325,6 +380,11 @@ def test_matmul_add_broadcast():
 
 
 def test_matmul_add_broadcast2():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-add broadcast net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -349,6 +409,11 @@ def test_matmul_add_broadcast2():
 
 
 def test_matmul_sub_broadcast():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-sub broadcast net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -373,6 +438,11 @@ def test_matmul_sub_broadcast():
 
 
 def test_matmul_sub_broadcast2():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-sub broadcast net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -397,6 +467,11 @@ def test_matmul_sub_broadcast2():
 
 
 def test_matmul_mul_broadcast():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-mul broadcast net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -421,6 +496,11 @@ def test_matmul_mul_broadcast():
 
 
 def test_matmul_mul_broadcast2():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-mul broadcast net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -445,6 +525,11 @@ def test_matmul_mul_broadcast2():
 
 
 def test_matmul_div_broadcast():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-div broadcast net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -469,6 +554,11 @@ def test_matmul_div_broadcast():
 
 
 def test_matmul_div_broadcast2():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-div broadcast net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -493,6 +583,11 @@ def test_matmul_div_broadcast2():
 
 
 def test_matmul_greater_broadcast():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-greater broadcast net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -517,6 +612,11 @@ def test_matmul_greater_broadcast():
 
 
 def test_matmul_greater_broadcast2():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-greater broadcast net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -541,6 +641,11 @@ def test_matmul_greater_broadcast2():
 
 
 def test_matmul_floordiv():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-floordiv net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -565,6 +670,11 @@ def test_matmul_floordiv():
 
 
 def test_matmul_floordiv_broadcast():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-floordiv broadcast net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -589,6 +699,11 @@ def test_matmul_floordiv_broadcast():
 
 
 def test_matmul_floordiv_broadcast2():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: matmul-floordiv broadcast net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
@@ -613,6 +728,11 @@ def test_matmul_floordiv_broadcast2():
 
 
 def test_assign_sub():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: mul-assign_sub net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self):
             super().__init__()
@@ -661,6 +781,11 @@ def test_assign_sub():
 
 
 def test_assign_add():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: mul-assign_add net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self):
             super().__init__()
@@ -709,6 +834,11 @@ def test_assign_add():
 
 
 def test_assign():
+    """
+    Feature: distribute operator sub in auto parallel.
+    Description: mul-assign_sub net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class Net(nn.Cell):
         def __init__(self):
             super().__init__()

tests/ut/python/parallel/test_get_next.py

@@ -23,10 +23,19 @@ from mindspore.ops import composite as C
 from mindspore.ops import operations as P
 
 context.set_context(mode=context.GRAPH_MODE)
-
-
 grad_by_list = C.GradOperation(get_by_list=True)
 
+class Net(nn.Cell):
+    def __init__(self, channel=1, w=0.25, strategy1=None, strategy2=None):
+        super().__init__()
+        self.norm = P.L2Normalize().shard(strategy1)
+        self.prelu = P.PReLU().shard(strategy2)
+        self.w = Parameter(initializer(w, [channel,]), name='w')
+
+    def construct(self, data):
+        x = self.norm(data)
+        x = self.prelu(x, self.w)
+        return x
 
 class NetWithLoss(nn.Cell):
     def __init__(self, network, types, shapes, output_num, strategy3=None, strategy4=None, axis=-1):
@@ -59,37 +68,21 @@ def compile_net(net):
     net.set_auto_parallel()
     _cell_graph_executor.compile(net)
 
-
 def test_get_next_single():
-    class Net(nn.Cell):
-        def __init__(self, channel=1, w=0.25):
-            super().__init__()
-            self.norm = P.L2Normalize(axis=1)
-            self.prelu = P.PReLU()
-            self.w = Parameter(initializer(w, [channel,]), name='w')
-
-        def construct(self, data):
-            x = self.norm(data)
-            x = self.prelu(x, self.w)
-            return x
-
+    """
+    Feature: test get next ops
+    Description: standalone, getnext-norm-prelu-loss.
+    Expectation: compile well done.
+    """
     net = GradWrap(NetWithLoss(Net(), [ms.float32, ms.int32], [[32, 64], [32]], 2))
     _cell_graph_executor.compile(net)
 
-
 def test_get_next_semi_auto_parallel():
-    class Net(nn.Cell):
-        def __init__(self, channel=1, w=0.25, strategy1=None, strategy2=None):
-            super().__init__()
-            self.norm = P.L2Normalize().shard(strategy1)
-            self.prelu = P.PReLU().shard(strategy2)
-            self.w = Parameter(initializer(w, [channel,]), name='w')
-
-        def construct(self, data):
-            x = self.norm(data)
-            x = self.prelu(x, self.w)
-            return x
-
+    """
+    Feature: test get next ops
+    Description: semi-auto parallel, getnext-norm-prelu-loss.
+    Expectation: compile well done.
+    """
     context.set_auto_parallel_context(device_num=4, global_rank=0)
     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
     network = Net(strategy1=((1, 4),), strategy2=((4, 1), (1,)))
@@ -102,18 +95,11 @@ def test_get_next_semi_auto_parallel():
 
 
 def test_get_next_semi_auto_parallel1():
-    class Net(nn.Cell):
-        def __init__(self, channel=1, w=0.25, strategy1=None, strategy2=None):
-            super().__init__()
-            self.norm = P.L2Normalize().shard(strategy1)
-            self.prelu = P.PReLU().shard(strategy2)
-            self.w = Parameter(initializer(w, [channel,]), name='w')
-
-        def construct(self, data):
-            x = self.norm(data)
-            x = self.prelu(x, self.w)
-            return x
-
+    """
+    Feature: test get next ops
+    Description: semi-auto parallel, getnext-norm-prelu-loss.
+    Expectation: compile well done.
+    """
     context.set_auto_parallel_context(device_num=4, global_rank=0)
     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
     network = Net(strategy1=((1, 4),), strategy2=((4, 1), (1,)))
@@ -126,18 +112,11 @@ def test_get_next_semi_auto_parallel1():
 
 
 def test_get_next_auto_parallel():
-    class Net(nn.Cell):
-        def __init__(self, channel=1, w=0.25, strategy1=None, strategy2=None):
-            super().__init__()
-            self.norm = P.L2Normalize().shard(strategy1)
-            self.prelu = P.PReLU().shard(strategy2)
-            self.w = Parameter(initializer(w, [channel,]), name='w')
-
-        def construct(self, data):
-            x = self.norm(data)
-            x = self.prelu(x, self.w)
-            return x
-
+    """
+    Feature: test get next ops
+    Description: auto parallel, getnext-norm-prelu-loss.
+    Expectation: compile well done.
+    """
     context.set_auto_parallel_context(device_num=4, global_rank=0)
     context.set_auto_parallel_context(parallel_mode="auto_parallel")
     network = Net()
@@ -147,7 +126,12 @@ def test_get_next_auto_parallel():
 
 
 def test_only_one_get_next():
-    class Net(nn.Cell):
+    """
+    Feature: test get next ops
+    Description: semi-auto parallel, only getnext.
+    Expectation: compile well done.
+    """
+    class Net1(nn.Cell):
         def __init__(self):
             super().__init__()
             self.get_next = P.GetNext([ms.float32, ms.int32], [[32, 64], [32]], 2, "")
@@ -157,6 +141,6 @@ def test_only_one_get_next():
 
     context.set_auto_parallel_context(device_num=4, global_rank=0)
     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
-    net = Net()
+    net = Net1()
     net.set_train()
     compile_net(net)

tests/ut/python/parallel/test_parallel_mirror_group.py

@@ -65,7 +65,7 @@ def auto_parallel_compile_net(strategy1=None, strategy2=None, strategy3=None):
 def test_mirror_group():
     """
     Feature: save and load mirror group
-    Description: semi-auto, disable parallel optimizer.
+    Description: semi-auto parallel, disable parallel optimizer.
     Expectation: group info list match expectation value.
     """
     os.environ['GROUP_INFO_FILE'] = "./test_mirror_group.pb"
@@ -77,10 +77,25 @@ def test_mirror_group():
     context.reset_auto_parallel_context()
     del os.environ['GROUP_INFO_FILE']
 
+def test_mirror_group_auto_parallel():
+    """
+    Feature: save and load mirror group
+    Description: auto parallel, disable parallel optimizer.
+    Expectation: group info list match expectation value.
+    """
+    os.environ['GROUP_INFO_FILE'] = "./test_mirror_group_auto_parallel.pb"
+    context.set_auto_parallel_context(parallel_mode="auto_parallel",
+                                      device_num=32, enable_parallel_optimizer=False)
+    auto_parallel_compile_net(((8, 1), (1, 4)), ((32, 1), (1, 1)), ((8, 4), (4, 1)))
+    group_info_list = restore_group_info_list("./test_mirror_group_auto_parallel.pb")
+    assert group_info_list == [0, 4, 8, 12, 16, 20, 24, 28]
+    context.reset_auto_parallel_context()
+    del os.environ['GROUP_INFO_FILE']
+
 def test_data_parallel_group():
     """
     Feature: save and load mirror group
-    Description: data-parallel, disable parallel optimizer.
+    Description: data parallel , disable parallel optimizer.
     Expectation: group info list match expectation value.
     """
     os.environ['GROUP_INFO_FILE'] = "./test_data_parallel_group.pb"
@@ -96,7 +111,7 @@ def test_data_parallel_group():
 def test_mirror_group_parallel_optimizer():
     """
     Feature: save and load mirror group
-    Description: semi-auto, enable parallel optimizer.
+    Description: semi-auto parallel, enable parallel optimizer.
     Expectation: group info list match expectation value.
     """
     os.environ['GROUP_INFO_FILE'] = "./test_mirror_group_parallel_optimizer.pb"
@@ -111,7 +126,7 @@ def test_mirror_group_parallel_optimizer():
 def test_mirror_group_parallel_optimizer_not_full_shard():
     """
     Feature: save and load mirror group
-    Description: semi-auto, enable parallel optimizer but not fully shard.
+    Description: semi-auto parallel, enable parallel optimizer but not fully shard.
     Expectation: group info list match expectation value.
     """
     os.environ['GROUP_INFO_FILE'] = "./test_mirror_group_parallel_optimizer_not_full_shard.pb"
@@ -126,7 +141,7 @@ def test_mirror_group_parallel_optimizer_not_full_shard():
 def test_pipeline_split_stage0_mirror_group():
     """
     Feature: save and load mirror group
-    Description: semi-auto, pipeline parallel.
+    Description: semi-auto parallel, pipeline parallel.
     Expectation: group info list match expectation value.
     """
     import mindspore as ms
@@ -147,3 +162,30 @@ def test_pipeline_split_stage0_mirror_group():
     model.train(2, dataset, dataset_sink_mode=False)
     group_info_list = restore_group_info_list("./test_pipeline_split_stage0_mirror_group.pb")
     assert group_info_list == [0, 8, 16, 24]
+    del os.environ['GROUP_INFO_FILE']
+
+def test_pipeline_split_stage1_mirror_group():
+    """
+    Feature: save and load mirror group
+    Description: semi-auto parallel, pipeline parallel.
+    Expectation: group info list match expectation value.
+    """
+    import mindspore as ms
+    from mindspore import Model
+    from .test_pipeline_split import PipelineCell, PipelineSplit, DatasetLenet
+    os.environ['GROUP_INFO_FILE'] = "./test_pipeline_split_stage1_mirror_group.pb"
+    context.set_auto_parallel_context(device_num=64, global_rank=63, pipeline_stages=2)
+    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
+    data = Tensor(np.ones([32, 64]), dtype=ms.float32)
+    label = Tensor(np.ones([64, 64]), dtype=ms.float32)
+    strategy1 = ((4, 1), (1, 8))
+    strategy2 = ((4, 1), (1, 1))
+    net = PipelineCell(PipelineSplit(strategy1, strategy2), 4)
+    params = net.network.cell.block[1].trainable_params()
+    dataset = DatasetLenet(data, label, 3)
+    optimizer = nn.Lamb(params, learning_rate=0.01)
+    model = Model(net, optimizer=optimizer)
+    model.train(2, dataset, dataset_sink_mode=False)
+    group_info_list = restore_group_info_list("./test_pipeline_split_stage1_mirror_group.pb")
+    assert group_info_list == [39, 47, 55, 63]
+    del os.environ['GROUP_INFO_FILE']

tests/ut/python/parallel/test_prelu.py

@@ -26,7 +26,6 @@ from tests.ut.python.ops.test_math_ops import VirtualLoss
 
 grad_all = C.GradOperation(get_all=True)
 
-
 class NetWithLoss(nn.Cell):
     def __init__(self, network):
         super(NetWithLoss, self).__init__()
@@ -37,7 +36,6 @@ class NetWithLoss(nn.Cell):
         predict = self.network(x, y)
         return self.loss(predict)
 
-
 class GradWrap(nn.Cell):
     def __init__(self, network):
         super(GradWrap, self).__init__()
@@ -46,88 +44,75 @@ class GradWrap(nn.Cell):
     def construct(self, x, y):
         return grad_all(self.network)(x, y)
 
+class Net(nn.Cell):
+    def __init__(self, strategy=None):
+        super().__init__()
+        self.prelu = P.PReLU().shard(strategy)
+
+    def construct(self, x, y):
+        out = self.prelu(x, y)
+        return out
 
 def compile_net(net, x, y):
     net.set_auto_parallel()
     net.set_train()
     _cell_graph_executor.compile(net, x, y)
 
-
-def test_prelu_single_success1():
-    class Net(nn.Cell):
-        def __init__(self):
-            super().__init__()
-            self.prelu = P.PReLU()
-
-        def construct(self, x, y):
-            out = self.prelu(x, y)
-            return out
-
-    context.reset_auto_parallel_context()
-    net = GradWrap(NetWithLoss(Net()))
-    x = Tensor(np.random.rand(1, 33, 4, 4), ms.float32)
-    w = Tensor(np.random.rand(33), ms.float32)
+def common_train_compile(input1_shape, input2_shape, strategy=None):
+    x = Tensor(np.random.rand(*input1_shape), dtype=ms.float32)
+    w = Tensor(np.random.rand(*input2_shape), dtype=ms.float32)
+    net = GradWrap(NetWithLoss(Net(strategy)))
     compile_net(net, x, w)
-
-
-def test_prelu_single_success2():
-    class Net(nn.Cell):
-        def __init__(self):
-            super().__init__()
-            self.prelu = P.PReLU()
-
-        def construct(self, x, y):
-            out = self.prelu(x, y)
-            return out
-
     context.reset_auto_parallel_context()
-    net = GradWrap(NetWithLoss(Net()))
-    x = Tensor(np.random.rand(1, 33, 4, 4), ms.float32)
-    w = Tensor([0.1], ms.float32)
-    compile_net(net, x, w)
-
 
 def test_prelu_parallel_success1():
-    class Net(nn.Cell):
-        def __init__(self, strategy):
-            super().__init__()
-            self.prelu = P.PReLU().shard(strategy)
-
-        def construct(self, x, y):
-            out = self.prelu(x, y)
-            return out
-
-    context.reset_auto_parallel_context()
+    """
+    Feature: distribute operator prelu in auto parallel.
+    Description: prelu net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     context.set_auto_parallel_context(device_num=8, global_rank=0)
     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
-    strategy = ((1, 1, 1, 1), (1,))
-    x = Tensor(np.random.rand(4, 4, 32, 64), dtype=ms.float32)
-    w = Tensor(np.random.rand(4), dtype=ms.float32)
-    net = GradWrap(NetWithLoss(Net(strategy)))
-    compile_net(net, x, w)
+    common_train_compile((4, 4, 32, 64), (4,), ((1, 1, 1, 1), (1,)))
 
 
 def test_prelu_parallel_success2():
-    class Net(nn.Cell):
-        def __init__(self, strategy):
-            super().__init__()
-            self.prelu = P.PReLU().shard(strategy)
-
-        def construct(self, x, y):
-            out = self.prelu(x, y)
-            return out
+    """
+    Feature: distribute operator prelu in auto parallel.
+    Description: prelu net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
+    context.set_auto_parallel_context(device_num=64, global_rank=0)
+    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
+    common_train_compile((4, 4, 32, 64), (4,), ((2, 1, 4, 8), (1,)))
 
-    context.reset_auto_parallel_context()
+def test_prelu_parallel_success3():
+    """
+    Feature: distribute operator prelu in auto parallel.
+    Description: prelu net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     context.set_auto_parallel_context(device_num=64, global_rank=0)
     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
-    strategy = ((2, 1, 4, 8), (1,))
-    x = Tensor(np.random.rand(4, 4, 32, 64), dtype=ms.float32)
-    w = Tensor(np.random.rand(4), dtype=ms.float32)
-    net = GradWrap(NetWithLoss(Net(strategy)))
-    compile_net(net, x, w)
+    common_train_compile((4, 16, 32, 64), (16,), ((2, 4, 4, 2), (4,)))
 
 
-def test_prelu_parallel_success3():
+def test_prelu_parallel_success4():
+    """
+    Feature: distribute operator prelu in auto parallel.
+    Description: prelu net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
+    context.set_auto_parallel_context(device_num=64, global_rank=0)
+    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
+    common_train_compile((4, 16, 32, 64), (1,), ((2, 4, 4, 2), (1,)))
+
+def test_matmul_prelu_parallel_success():
+    """
+    Feature: distribute operator prelu in auto parallel.
+    Description: matmul-prelu net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     class NetWithLoss3(nn.Cell):
         def __init__(self, network):
             super(NetWithLoss3, self).__init__()
@@ -146,7 +131,7 @@ def test_prelu_parallel_success3():
         def construct(self, x, y, w):
             return grad_all(self.network)(x, y, w)
 
-    class Net(nn.Cell):
+    class Net3(nn.Cell):
         def __init__(self, strategy1, strategy2):
             super().__init__()
             self.matmul = P.MatMul().shard(strategy1)
@@ -165,47 +150,7 @@ def test_prelu_parallel_success3():
     x = Tensor(np.random.rand(128, 64), dtype=ms.float32)
     y = Tensor(np.random.rand(64, 16), dtype=ms.float32)
     w = Tensor(np.random.rand(16), dtype=ms.float32)
-    net = GradWrap3(NetWithLoss3(Net(strategy1, strategy2)))
+    net = GradWrap3(NetWithLoss3(Net3(strategy1, strategy2)))
     net.set_auto_parallel()
     net.set_train()
     _cell_graph_executor.compile(net, x, y, w)
-
-
-def test_prelu_parallel_success4():
-    class Net(nn.Cell):
-        def __init__(self, strategy):
-            super().__init__()
-            self.prelu = P.PReLU().shard(strategy)
-
-        def construct(self, x, y):
-            out = self.prelu(x, y)
-            return out
-
-    context.reset_auto_parallel_context()
-    context.set_auto_parallel_context(device_num=64, global_rank=0)
-    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
-    strategy = ((2, 4, 4, 2), (4,))
-    x = Tensor(np.random.rand(4, 16, 32, 64), dtype=ms.float32)
-    w = Tensor(np.random.rand(16), dtype=ms.float32)
-    net = GradWrap(NetWithLoss(Net(strategy)))
-    compile_net(net, x, w)
-
-
-def test_prelu_parallel_success5():
-    class Net(nn.Cell):
-        def __init__(self, strategy):
-            super().__init__()
-            self.prelu = P.PReLU().shard(strategy)
-
-        def construct(self, x, y):
-            out = self.prelu(x, y)
-            return out
-
-    context.reset_auto_parallel_context()
-    context.set_auto_parallel_context(device_num=64, global_rank=0)
-    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
-    strategy = ((2, 4, 4, 2), (1,))
-    x = Tensor(np.random.rand(4, 16, 32, 64), dtype=ms.float32)
-    w = Tensor(np.random.rand(1), dtype=ms.float32)
-    net = GradWrap(NetWithLoss(Net(strategy)))
-    compile_net(net, x, w)

tests/ut/python/parallel/test_prelu_cell.py

@@ -18,19 +18,15 @@ import mindspore as ms
 import mindspore.nn as nn
 from mindspore import Tensor
 from mindspore import context
-from mindspore.common.initializer import initializer
-from mindspore.common.parameter import Parameter
 from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits
 from mindspore.nn.optim.momentum import Momentum
-from mindspore.ops import functional as F
-from mindspore.ops import operations as P
 from mindspore.train import Model
 from mindspore.context import ParallelMode
+from mindspore.nn import PReLU
 from tests.dataset_mock import MindData
 
 context.set_context(mode=context.GRAPH_MODE)
 
-
 class Dataset(MindData):
     def __init__(self, predict, label, length=3, input_num=2):
         super(Dataset, self).__init__(size=length)
@@ -55,33 +51,6 @@ class Dataset(MindData):
         self.index = 0
 
 
-class PReLU(nn.Cell):
-    def __init__(self, channel=1, w=0.25):
-        super(PReLU, self).__init__()
-        if isinstance(w, (np.float32, float)):
-            tmp = np.empty((channel,), dtype=np.float32)
-            tmp.fill(w)
-            w = Tensor(tmp)
-        elif isinstance(w, list):
-            w = Tensor(w)
-
-        if not isinstance(w, Tensor):
-            raise TypeError("w only support np.float32, float or Tensor type.")
-
-        self.w = Parameter(initializer(w, [channel,]), name='a')
-        self.prelu = P.PReLU()
-        self.relu = P.ReLU().shard(((1,),))
-        self.sub = P.Sub().shard(((1,), (1,)))
-        self.assign_sub = P.AssignSub().shard(((1,), (1,)))
-
-    def construct(self, x):
-        u = self.relu(self.w)
-        tmp = self.sub(self.w, u)
-        x = F.depend(x, self.assign_sub(self.w, tmp))
-        v = self.prelu(x, u)
-        return v
-
-
 class PReLUNet(nn.Cell):
     def __init__(self):
         super(PReLUNet, self).__init__()
@@ -91,11 +60,6 @@ class PReLUNet(nn.Cell):
         x = self.prelu(x)
         return x
 
-
-def prelu_net():
-    return PReLUNet()
-
-
 def reshape_common(parallel_mode):
     learning_rate = 0.1
     momentum = 0.9
@@ -106,7 +70,7 @@ def reshape_common(parallel_mode):
     predict = Tensor(np.ones([32, 256]), dtype=ms.float32)
     label = Tensor(np.ones([32]), dtype=ms.int32)
     dataset = Dataset(predict, label, 2)
-    net = prelu_net()
+    net = PReLUNet()
 
     loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
     opt = Momentum(net.trainable_params(), learning_rate, momentum)
@@ -115,4 +79,9 @@ def reshape_common(parallel_mode):
 
 
 def test_prelu_cell():
+    """
+    Feature: distribute operator prelu in auto parallel.
+    Description: prelu net with strategy in semi auto parallel.
+    Expectation: compile done without error.
+    """
     reshape_common(ParallelMode.SEMI_AUTO_PARALLEL)