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