• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020 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
18from mindspore.common import dtype as mstype
19from mindspore import context, Tensor, Parameter
20from mindspore.nn import Cell, Momentum
21from mindspore.ops import operations as P
22from mindspore.train import Model
23from tests.dataset_mock import MindData
24
25
26class Dataset(MindData):
27    def __init__(self, predict, label, length=3):
28        super(Dataset, self).__init__(size=length)
29        self.predict = predict
30        self.label = label
31        self.index = 0
32        self.length = length
33
34    def __iter__(self):
35        return self
36
37    def __next__(self):
38        if self.index >= self.length:
39            raise StopIteration
40        self.index += 1
41        return self.predict, self.label
42
43    def reset(self):
44        self.index = 0
45
46
47class Net(Cell):
48    def __init__(self, weight, start, limit, delta, strategy1=None, strategy2=None, strategy3=None):
49        super().__init__()
50        self.mul = P.Mul().shard(strategy1)
51        if isinstance(start, float):
52            self.type = mstype.float32
53        else:
54            self.type = mstype.int32
55        self.start = Tensor(start, self.type)
56        self.limit = Tensor(limit, self.type)
57        self.delta = Tensor(delta, self.type)
58        self.range = P.Range()
59        self.range.shard(strategy2)
60        self.mul2 = P.Mul().shard(strategy3)
61        self.weight = Parameter(weight, "w")
62
63    def construct(self, x, b):
64        r_out = self.range(self.start, self.limit, self.delta)
65        out = self.mul(x, self.weight)
66        out = self.mul2(out, r_out)
67        return out
68
69
70dev_num = 4
71_x = Tensor(np.ones([64 // dev_num, 8]), dtype=ms.float32)
72_b = Tensor(np.ones([8]), dtype=ms.float32)
73_w1 = Tensor(np.ones([64, 8]), dtype=ms.float32)
74
75
76def compile_net(net):
77    learning_rate = 0.1
78    momentum = 0.9
79    epoch_size = 2
80    dataset = Dataset(_x, _b)
81    opt = Momentum(net.trainable_params(), learning_rate, momentum)
82    model = Model(net, optimizer=opt)
83    model.train(epoch_size, dataset, dataset_sink_mode=False)
84    context.reset_auto_parallel_context()
85
86
87def test_range():
88    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=dev_num, global_rank=2)
89    strategy1 = ((2, 2), (2, 2))
90    strategy2 = ((2,),)
91    strategy3 = ((2, 2), (2,))
92    net = Net(_w1, 0, 8, 1, strategy1, strategy2, strategy3)
93    compile_net(net)
94
95
96def test_range2():
97    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=dev_num, global_rank=0)
98    strategy1 = ((4, 1), (4, 1))
99    strategy2 = ((1,),)
100    strategy3 = ((4, 1), (1,))
101    net = Net(_w1, 0.0, 4.0, 0.5, strategy1, strategy2, strategy3)
102    compile_net(net)
103
104
105def test_range3():
106    context.set_auto_parallel_context(parallel_mode="auto_parallel", device_num=dev_num, global_rank=2)
107    net = Net(_w1, 0.0, 4.0, 0.5)
108    compile_net(net)
109