1# Copyright 2021 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 import context, Tensor, Parameter 19from mindspore.common.api import _cell_graph_executor 20from mindspore.nn import Cell, TrainOneStepCell, Momentum 21from mindspore.ops import operations as P 22 23 24class Net(Cell): 25 def __init__(self, dim, index, strategy1=None, strategy2=None): 26 super().__init__() 27 self.gatherd = P.GatherD().shard(strategy1) 28 self.neg = P.Neg().shard(strategy2) 29 self.input = Parameter(index, "w1") 30 self.dim = dim 31 32 def construct(self, x, b): 33 out = self.gatherd(self.input, self.dim, x) 34 out = self.neg(out) 35 return out 36 37 38_x = Tensor(np.ones([16, 32, 64]), dtype=ms.int32) 39_w1 = Tensor(np.ones([16, 32, 64]), dtype=ms.float32) 40_b = Tensor(np.ones([16, 32, 64]), dtype=ms.float32) 41 42 43def compile_net(net): 44 optimizer = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9) 45 train_net = TrainOneStepCell(net, optimizer) 46 train_net.set_auto_parallel() 47 train_net.set_train() 48 _cell_graph_executor.compile(train_net, _x, _b) 49 context.reset_auto_parallel_context() 50 51 52def test_gathernd_dim0(): 53 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0) 54 strategy1 = ((1, 2, 8), (1, 2, 8)) 55 strategy2 = ((1, 2, 8),) 56 net = Net(0, _w1, strategy1, strategy2) 57 compile_net(net) 58 59 60def test_gathernd_dim2(): 61 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0) 62 strategy1 = ((2, 8, 1), (2, 8, 1)) 63 strategy2 = ((2, 8, 1),) 64 net = Net(2, _w1, strategy1, strategy2) 65 compile_net(net) 66 67 68def test_gathernd_dim2_default_batch_parallel(): 69 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0) 70 strategy1 = None 71 strategy2 = ((2, 8, 1),) 72 net = Net(2, _w1, strategy1, strategy2) 73 compile_net(net) 74 75 76def test_gathernd_auto_parallel(): 77 context.set_auto_parallel_context(parallel_mode="auto_parallel", device_num=16, global_rank=0) 78 net = Net(1, _w1) 79 compile_net(net) 80 81 82def test_gathernd_repeat_calc(): 83 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0) 84 strategy1 = ((1, 2, 4), (1, 2, 4)) 85 strategy2 = ((1, 2, 4),) 86 net = Net(0, _w1, strategy1, strategy2) 87 compile_net(net) 88