• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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.context as context
18import mindspore.nn as nn
19from mindspore import Tensor, Model, ms_function
20from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits
21from mindspore.ops import operations as P
22
23context.set_context(device_target="Ascend")
24
25input_channel = 2048
26output_channel = 512
27num_class = 10
28batch_size = 32
29
30
31class MsWrapper(nn.Cell):
32    def __init__(self, network):
33        super(MsWrapper, self).__init__(auto_prefix=False)
34        self._network = network
35
36    @ms_function
37    def construct(self, *args):
38        return self._network(*args)
39
40
41def me_train_tensor(net, input_np, label_np, epoch_size=2):
42    loss = SoftmaxCrossEntropyWithLogits(sparse=True)
43    opt = nn.Momentum(Tensor(np.array([0.1])), Tensor(np.array([0.9])),
44                      filter(lambda x: x.requires_grad, net.get_parameters()))
45    context.set_context(mode=context.GRAPH_MODE)
46    Model(net, loss, opt)
47    _network = nn.WithLossCell(net, loss)
48    _train_net = MsWrapper(nn.TrainOneStepCell(_network, opt))
49    _train_net.set_train()
50    for epoch in range(0, epoch_size):
51        print(f"epoch %d" % (epoch))
52        output = _train_net(Tensor(input_np), Tensor(label_np))
53        print(output.asnumpy())
54
55
56def test_conv_bn_add_relu_fusion():
57    class Net(nn.Cell):
58        def __init__(self):
59            super(Net, self).__init__()
60            self.conv = nn.Conv2d(input_channel, output_channel,
61                                  kernel_size=1, stride=1, padding=0, has_bias=False, pad_mode="same")
62            self.conv1 = nn.Conv2d(input_channel, output_channel,
63                                   kernel_size=1, stride=1, padding=0, has_bias=False, pad_mode="same")
64            self.bn = nn.BatchNorm2d(output_channel, momentum=0.1, eps=0.0001)
65            self.add = P.Add()
66            self.relu = P.ReLU()
67            self.mean = P.ReduceMean(keep_dims=True)
68            self.reshape = P.Reshape()
69            self.dense = nn.Dense(output_channel, num_class)
70
71        def construct(self, input_x):
72            output = self.conv(input_x)
73            output = self.bn(output)
74            output = self.add(output, self.conv1(input_x))
75            output = self.relu(output)
76            output = self.mean(output, (-2, -1))
77            output = self.reshape(output, (batch_size, output_channel))
78            output = self.dense(output)
79            return output
80
81    net = Net()
82    input_np = np.ones([batch_size, input_channel, 7, 7]).astype(np.float32) * 0.01
83    label_np = np.ones([batch_size]).astype(np.int32)
84    me_train_tensor(net, input_np, label_np)
85
86
87def test_conv_bn_relu_fusion():
88    class Net(nn.Cell):
89        def __init__(self):
90            super(Net, self).__init__()
91            self.conv = nn.Conv2d(input_channel, output_channel,
92                                  kernel_size=1, stride=1, padding=0, has_bias=False, pad_mode="same")
93            self.bn = nn.BatchNorm2d(output_channel, momentum=0.1, eps=0.0001)
94            self.relu = P.ReLU()
95            self.mean = P.ReduceMean(keep_dims=True)
96            self.reshape = P.Reshape()
97            self.dense = nn.Dense(output_channel, num_class)
98
99        def construct(self, input_x):
100            output = self.conv(input_x)
101            output = self.bn(output)
102            output = self.relu(output)
103            output = self.mean(output, (-2, -1))
104            output = self.reshape(output, (batch_size, output_channel))
105            output = self.dense(output)
106            return output
107
108    net = Net()
109    input_np = np.ones([batch_size, input_channel, 7, 7]).astype(np.float32) * 0.01
110    label_np = np.ones([batch_size]).astype(np.int32)
111    me_train_tensor(net, input_np, label_np)
112
113
114def test_conv_bn_fusion():
115    class Net(nn.Cell):
116        def __init__(self):
117            super(Net, self).__init__()
118            self.conv = nn.Conv2d(input_channel, output_channel,
119                                  kernel_size=1, stride=1, padding=0, has_bias=False, pad_mode="same")
120            self.bn = nn.BatchNorm2d(output_channel, momentum=0.1, eps=0.0001)
121            self.mean = P.ReduceMean(keep_dims=True)
122            self.reshape = P.Reshape()
123            self.dense = nn.Dense(output_channel, num_class)
124
125        def construct(self, input_x):
126            output = self.conv(input_x)
127            output = self.bn(output)
128            output = self.mean(output, (-2, -1))
129            output = self.reshape(output, (batch_size, output_channel))
130            output = self.dense(output)
131            return output
132
133    net = Net()
134    input_np = np.ones([batch_size, input_channel, 7, 7]).astype(np.float32) * 0.01
135    label_np = np.ones([batch_size]).astype(np.int32)
136    me_train_tensor(net, input_np, label_np)
137