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, BatchNorm2d, BatchNorm1d 21from mindspore.ops import operations as P 22from tests.security_utils import security_off_wrap 23 24class Net(Cell): 25 def __init__(self, conv2d_weight, out_channel, kernel_size, pad_mode, stride, 26 strategy1=None, strategy2=None): 27 super().__init__() 28 self.conv2d = P.Conv2D(out_channel=out_channel, kernel_size=kernel_size, 29 pad_mode=pad_mode, stride=stride).shard(strategy1) 30 self.conv2d_weight = Parameter(conv2d_weight, "w1") 31 self.bn = BatchNorm2d(8) 32 self.bn.bn_train.shard(strategy2) 33 self.print = P.Print() 34 35 def construct(self, x, b): 36 out = self.conv2d(x, self.conv2d_weight) 37 self.print("output is", out) 38 out = self.bn(out) 39 return out 40 41 42_x = Tensor(np.ones([32, 16, 8, 8]), dtype=ms.float32) 43_w1 = Tensor(np.ones([8, 16, 2, 2]), dtype=ms.float32) 44_b = Tensor(np.ones([32, 16, 8, 8]), dtype=ms.float32) 45 46 47def compile_net(net): 48 optimizer = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9) 49 train_net = TrainOneStepCell(net, optimizer) 50 train_net.set_auto_parallel() 51 train_net.set_train() 52 _cell_graph_executor.compile(train_net, _x, _b) 53 context.reset_auto_parallel_context() 54 55 56@security_off_wrap 57def test_batchnorm_data_parallel(): 58 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=8, global_rank=0) 59 strategy1 = ((8, 1, 1, 1), (1, 1, 1, 1)) 60 strategy2 = ((8, 1, 1, 1), (1,), (1,), (1,), (1,)) 61 net = Net(_w1, out_channel=8, kernel_size=2, pad_mode="same", stride=1, strategy1=strategy1, strategy2=strategy2) 62 compile_net(net) 63 64 65@security_off_wrap 66def test_batchnorm_model_parallel1(): 67 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=8, global_rank=0) 68 strategy1 = ((2, 2, 1, 1), (2, 2, 1, 1)) 69 strategy2 = ((2, 1, 2, 2), (1,), (1,), (1,), (1,)) 70 net = Net(_w1, out_channel=8, kernel_size=2, pad_mode="same", stride=1, strategy1=strategy1, strategy2=strategy2) 71 compile_net(net) 72 73 74@security_off_wrap 75def test_batchnorm_model_parallel2(): 76 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=32, global_rank=0) 77 strategy1 = ((2, 2, 2, 2), (2, 2, 1, 1)) 78 strategy2 = ((1, 8, 1, 1), (8,), (8,), (8,), (8,)) 79 net = Net(_w1, out_channel=8, kernel_size=2, pad_mode="same", stride=2, strategy1=strategy1, strategy2=strategy2) 80 compile_net(net) 81 82 83class Net2(Cell): 84 def __init__(self, strategy1=None, strategy2=None): 85 super().__init__() 86 self.bn = BatchNorm1d(8) 87 self.bn.bn_train.shard(strategy1) 88 self.relu = P.ReLU().shard(strategy2) 89 90 def construct(self, x, b): 91 out = self.bn(x) 92 out = self.relu(out) 93 return out 94 95 96_x1 = Tensor(np.ones([32, 8]), dtype=ms.float32) 97_b1 = Tensor(np.ones([32, 8]), dtype=ms.float32) 98 99 100def compile_net2(net): 101 optimizer = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9) 102 train_net = TrainOneStepCell(net, optimizer) 103 train_net.set_auto_parallel() 104 train_net.set_train() 105 _cell_graph_executor.compile(train_net, _x1, _b1) 106 context.reset_auto_parallel_context() 107 108 109def test_batchnorm1d_data_parallel(): 110 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=8, global_rank=0) 111 strategy1 = ((8, 1), (1,), (1,), (1,), (1,)) 112 strategy2 = ((8, 1),) 113 net = Net2(strategy1=strategy1, strategy2=strategy2) 114 compile_net2(net) 115 116 117def test_batchnorm1d_model_parallel1(): 118 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=8, global_rank=0) 119 strategy1 = ((1, 8), (8,), (8,), (8,), (8,)) 120 strategy2 = ((1, 8),) 121 net = Net2(strategy1=strategy1, strategy2=strategy2) 122 compile_net2(net) 123 124 125def test_batchnorm1d_model_parallel2(): 126 context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=32, global_rank=0) 127 strategy1 = ((2, 4), (4,), (4,), (4,), (4,)) 128 strategy2 = ((2, 4),) 129 net = Net2(strategy1=strategy1, strategy2=strategy2) 130 compile_net2(net) 131