• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2019 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import numpy as np
16
17import mindspore as ms
18import mindspore.nn as nn
19from mindspore import Tensor
20from mindspore import context
21from mindspore.common import dtype as mstype
22from mindspore.common.api import _cell_graph_executor
23from mindspore.ops import composite as C
24from mindspore.ops import operations as P
25from mindspore.parallel._utils import _reset_op_id as reset_op_id
26from tests.ut.python.ops.test_math_ops import VirtualLoss
27
28
29grad_all = C.GradOperation(get_all=True)
30
31
32class NetWithLoss(nn.Cell):
33    def __init__(self, network):
34        super(NetWithLoss, self).__init__()
35        self.loss = VirtualLoss()
36        self.network = network
37
38    def construct(self, x, y, z, w):
39        predict = self.network(x, y, z, w)
40        return self.loss(predict)
41
42
43class GradWrap(nn.Cell):
44    def __init__(self, network):
45        super(GradWrap, self).__init__()
46        self.network = network
47
48    def construct(self, x, y, z, w):
49        return grad_all(self.network)(x, y, z, w)
50
51    # model_parallel test
52
53
54def test_double_star_graph():
55    class Net(nn.Cell):
56        def __init__(self):
57            super().__init__()
58            self.matmul1 = P.MatMul()
59            self.matmul2 = P.MatMul()
60            self.matmul3 = P.MatMul()
61            self.cast1 = P.Cast()
62            self.cast2 = P.Cast()
63
64        def construct(self, x, y, z, w):
65            m1_result = self.matmul1(x, y)
66            m2_result = self.matmul2(z, w)
67            m3_result = self.matmul3(self.cast1(m2_result, mstype.float16), self.cast2(m1_result, mstype.float16))
68
69            return m3_result
70
71    size = 8
72    context.set_auto_parallel_context(device_num=size, global_rank=0)
73
74    x = Tensor(np.ones([32, 8]), dtype=ms.float32)
75    y = Tensor(np.ones([8, 16]), dtype=ms.float32)
76    z = Tensor(np.ones([8, 16]), dtype=ms.float32)
77    w = Tensor(np.ones([16, 32]), dtype=ms.float32)
78
79    net = NetWithLoss(Net())
80    context.set_auto_parallel_context(parallel_mode="auto_parallel")
81    net.set_auto_parallel()
82    reset_op_id()
83
84    net.set_train()
85    _cell_graph_executor.compile(net, x, y, z, w, phase='train')
86    strategies = _cell_graph_executor._get_shard_strategy(net)
87    expected_strategies = {'Default/network-Net/Cast-op1': [[8, 1]],
88                           'Default/network-Net/Cast-op3': [[1, 8]],
89                           'Default/network-Net/MatMul-op2': [[8, 1], [1, 1]],
90                           'Default/network-Net/MatMul-op4': [[1, 1], [1, 8]],
91                           'Default/network-Net/MatMul-op0': [[1, 8], [8, 1]]}
92    assert strategies == expected_strategies
93