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.common.api as me 19import mindspore.nn as nn 20from mindspore import Tensor, Parameter 21from mindspore import context 22from mindspore.ops import operations as P 23 24 25def test_get_parameter_layout(): 26 class Net(nn.Cell): 27 def __init__(self, strategy1, strategy2, weight): 28 super().__init__() 29 self.weight = Parameter(weight, "w1") 30 self.matmul = P.MatMul(transpose_a=False, transpose_b=True).shard(strategy1) 31 self.relu = P.ReLU().shard(strategy2) 32 33 def construct(self, x): 34 out = self.matmul(x, self.weight) 35 out = self.relu(out) 36 return out 37 38 context.reset_auto_parallel_context() 39 context.set_auto_parallel_context(device_num=8, global_rank=0) 40 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel") 41 strategy1 = ((2, 1), (4, 1)) 42 strategy2 = ((2, 4),) 43 context.set_context(mode=context.GRAPH_MODE) 44 45 x = Tensor(np.ones([32, 32]), dtype=ms.float32) 46 weight = Tensor(np.ones([64, 32]), dtype=ms.float32) 47 48 net = Net(strategy1, strategy2, weight) 49 net.set_auto_parallel() 50 net.set_train() 51 exe = me._cell_graph_executor 52 exe.compile(net, x, phase='train', auto_parallel_mode=True) 53 x_layout = ([2, 4], [1, -1], [16, 32], 0, True, '') # device_arrangement = [2, 4], tensor_map = [1, -1] 54 weight_layout = ([2, 4], [0, -1], [16, 32], 0, True, '') # device_arrangement = [2, 4], tensor_map = [0, -1] 55 expect_dict = {'x': x_layout, 'w1': weight_layout} 56 # to be resovled: static local variable count_p is used in step_parallel.cc, it needs to be reset between each ut 57 assert net.parameter_layout_dict == expect_dict 58 59 60if __name__ == '__main__': 61 test_get_parameter_layout() 62