• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020 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# ============================================================================
15""" test model train """
16import mindspore.nn as nn
17from mindspore import Tensor, Model
18from mindspore.common import dtype as mstype
19from mindspore.common.parameter import ParameterTuple, Parameter
20from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits
21from mindspore.nn.optim import Momentum
22from mindspore.ops import composite as C
23from mindspore.ops import operations as P
24
25
26def get_reordered_parameters(parameters):
27    """get_reordered_parameters"""
28    # put the bias parameter to the end
29    non_bias_param = []
30    bias_param = []
31    for item in parameters:
32        if item.name.find("bias") >= 0:
33            bias_param.append(item)
34        else:
35            non_bias_param.append(item)
36    reordered_params = tuple(non_bias_param + bias_param)
37    return len(non_bias_param), len(reordered_params), reordered_params
38
39
40def get_net_trainable_reordered_params(net):
41    params = net.trainable_params()
42    return get_reordered_parameters(params)
43
44
45class TrainOneStepWithLarsCell(nn.Cell):
46    """TrainOneStepWithLarsCell definition"""
47
48    def __init__(self, network, optimizer, sens=1.0):
49        super(TrainOneStepWithLarsCell, self).__init__(auto_prefix=False)
50        self.network = network
51        self.slice_index, self.params_len, weights = get_net_trainable_reordered_params(self.network)
52        self.weights = ParameterTuple(weights)
53        self.optimizer = optimizer
54        self.grad = C.GradOperation(get_by_list=True,
55                                    sens_param=True)
56        self.sens = Parameter(Tensor([sens], mstype.float32), name='sens', requires_grad=False)
57        self.weight_decay = 1.0
58        self.lars = P.Lars(epsilon=1.0, hyperpara=1.0)
59
60    def construct(self, data, label):
61        weights = self.weights
62        loss = self.network(data, label)
63        grads = self.grad(self.network, weights)(data, label, self.sens)
64        non_bias_weights = weights[0: self.slice_index]
65        non_bias_grads = grads[0: self.slice_index]
66        bias_grads = grads[self.slice_index: self.params_len]
67        lars_grads = self.lars(non_bias_weights, non_bias_grads, self.weight_decay)
68        new_grads = lars_grads + bias_grads
69        self.optimizer(new_grads)
70        return loss
71
72
73# fn is a function use i as input
74def lr_gen(fn, epoch_size):
75    for i in range(epoch_size):
76        yield fn(i)
77
78
79def me_train_tensor(net, input_np, label_np, epoch_size=2):
80    """me_train_tensor"""
81    loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
82    # reorder the net parameters , leave the parameters that need to be passed into lars to the end part
83
84    opt = Momentum(get_net_trainable_reordered_params(net)[2], lr_gen(lambda i: 0.1, epoch_size), 0.9, 0.01, 1024)
85    Model(net, loss, opt)
86    _network = nn.WithLossCell(net, loss)
87    TrainOneStepWithLarsCell(_network, opt)
88    data = Tensor(input_np)
89    label = Tensor(label_np)
90    net(data, label)
91