• 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 as ms
18import mindspore.nn as nn
19from mindspore import Tensor
20from mindspore import context
21from mindspore.common.initializer import initializer
22from mindspore.common.parameter import Parameter
23from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits
24from mindspore.nn.optim.momentum import Momentum
25from mindspore.ops import functional as F
26from mindspore.ops import operations as P
27from mindspore.train import Model
28from mindspore.context import ParallelMode
29from tests.dataset_mock import MindData
30
31context.set_context(mode=context.GRAPH_MODE)
32
33
34class Dataset(MindData):
35    def __init__(self, predict, label, length=3, input_num=2):
36        super(Dataset, self).__init__(size=length)
37        self.predict = predict
38        self.label = label
39        self.index = 0
40        self.length = length
41        self.input_num = input_num
42
43    def __iter__(self):
44        return self
45
46    def __next__(self):
47        if self.index >= self.length:
48            raise StopIteration
49        self.index += 1
50        if self.input_num == 2:
51            return (self.predict, self.label)
52        return (self.predict,)
53
54    def reset(self):
55        self.index = 0
56
57
58class PReLU(nn.Cell):
59    def __init__(self, channel=1, w=0.25):
60        super(PReLU, self).__init__()
61        if isinstance(w, (np.float32, float)):
62            tmp = np.empty((channel,), dtype=np.float32)
63            tmp.fill(w)
64            w = Tensor(tmp)
65        elif isinstance(w, list):
66            w = Tensor(w)
67
68        if not isinstance(w, Tensor):
69            raise TypeError("w only support np.float32, float or Tensor type.")
70
71        self.w = Parameter(initializer(w, [channel,]), name='a')
72        self.prelu = P.PReLU()
73        self.relu = P.ReLU().shard(((1,),))
74        self.sub = P.Sub().shard(((1,), (1,)))
75        self.assign_sub = P.AssignSub().shard(((1,), (1,)))
76
77    def construct(self, x):
78        u = self.relu(self.w)
79        tmp = self.sub(self.w, u)
80        x = F.depend(x, self.assign_sub(self.w, tmp))
81        v = self.prelu(x, u)
82        return v
83
84
85class PReLUNet(nn.Cell):
86    def __init__(self):
87        super(PReLUNet, self).__init__()
88        self.prelu = PReLU(channel=256)
89
90    def construct(self, x):
91        x = self.prelu(x)
92        return x
93
94
95def prelu_net():
96    return PReLUNet()
97
98
99def reshape_common(parallel_mode):
100    learning_rate = 0.1
101    momentum = 0.9
102    epoch_size = 2
103
104    context.reset_auto_parallel_context()
105    context.set_auto_parallel_context(parallel_mode=parallel_mode, device_num=8)
106    predict = Tensor(np.ones([32, 256]), dtype=ms.float32)
107    label = Tensor(np.ones([32]), dtype=ms.int32)
108    dataset = Dataset(predict, label, 2)
109    net = prelu_net()
110
111    loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
112    opt = Momentum(net.trainable_params(), learning_rate, momentum)
113    model = Model(net, loss, opt)
114    model.train(epoch_size, dataset, dataset_sink_mode=False)
115
116
117def test_prelu_cell():
118    reshape_common(ParallelMode.SEMI_AUTO_PARALLEL)
119