Huawei_Technology
/
mindspore
Not watched
1
0
Fork 0
Code
Releases 13 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: ef13a4b6fb
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'ef13a4b6fb'
${ noResults }
mindspore/example/resnet50_imagenet2012_THOR/model/thor.py

200 lines
10 kB
Raw Blame History 

  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. """momentum"""

  16. import mindspore.common.dtype as mstype

  17. from mindspore.common.initializer import initializer

  18. from mindspore.common.parameter import Parameter

  19. from mindspore.common.parameter import ParameterTuple

  20. from mindspore.common.tensor import Tensor

  21. from mindspore.nn.optim.optimizer import Optimizer

  22. from mindspore.ops import functional as F, composite as C, operations as P

  23. from mindspore.parallel._utils import _get_device_num, _get_mirror_mean

  24. from model.grad_reducer_thor import DistributedGradReducerThor

  25. 

  26. momentum_opt = C.MultitypeFuncGraph("momentum_opt")

  27. 

  28. 

  29. @momentum_opt.register("Function", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")

  30. def _tensor_run_opt_ext(opt, learning_rate, momentum, gradient, weight, moment):

  31.     """Apply momentum optimizer to the weight parameter using Tensor."""

  32.     success = True

  33.     success = F.depend(success, opt(weight, moment, learning_rate, gradient, momentum))

  34.     return success

  35. 

  36. 

  37. op_add = P.AddN()

  38. apply_decay = C.MultitypeFuncGraph("apply_decay")

  39. 

  40. 

  41. @apply_decay.register("Number", "Bool", "Tensor", "Tensor")

  42. def _tensor_apply_decay(weight_decay, if_apply, weight, gradient):

  43.     """Get grad with weight_decay."""

  44.     if if_apply:

  45.         return op_add((weight * weight_decay, gradient))

  46.     return gradient

  47. 

  48. 

  49. class THOR(Optimizer):

  50.     """THOR"""

  51.     def __init__(self, params, learning_rate, momentum, matrix_A, matrix_G, A_inv_max, G_inv_max, weight_decay=0.0,

  52.                  loss_scale=1.0,

  53.                  decay_filter=lambda x: x.name not in []):

  54.         super(THOR, self).__init__(learning_rate, params, weight_decay, loss_scale)

  55.         if isinstance(momentum, float) and momentum < 0.0:

  56.             raise ValueError("momentum should be at least 0.0, but got momentum {}".format(momentum))

  57.         self.momentum = Parameter(Tensor(momentum, mstype.float32), name="momentum")

  58.         self.params = self.parameters

  59.         self.moments = self.params.clone(prefix="moments", init='zeros')

  60.         self.hyper_map = C.HyperMap()

  61.         self.opt = P.ApplyMomentum()

  62.         self.matrix_A = ParameterTuple(matrix_A)

  63.         self.matrix_G = ParameterTuple(matrix_G)

  64.         self.A_inv_max = ParameterTuple(A_inv_max)

  65.         self.G_inv_max = ParameterTuple(G_inv_max)

  66.         self.cube_matmul_left = P.CusMatMulCubeFraczLeftCast()

  67.         self.cube_matmul_left_fc = P.CusMatMulCubeDenseLeft()

  68.         self.cube_matmul_right_fc = P.CusMatMulCubeDenseRight()

  69.         self.cube_matmul_right_mul = P.CusMatMulCubeFraczRightMul()

  70.         self.transpose = P.Transpose()

  71.         self.shape = P.Shape()

  72.         self.reshape = P.Reshape()

  73.         self.mul = P.Mul()

  74.         self.weight_idx = []

  75.         for i in range(len(self.params)):

  76.             if "conv" in self.params[i].name or "end_point" in self.params[i].name:

  77.                 self.weight_idx.append(i)

  78.         self.weight_idx.append(len(self.params))

  79.         self.feature_map = [1.0 / 12544, 1.0 / 3136, 1.0 / 3136, 1.0 / 3136, 1.0 / 3136, 1.0 / 3136, 1.0 / 3136,

  80.                             1.0 / 3136, 1.0 / 3136, 1.0 / 3136, 1.0 / 3136, 1.0 / 3136,

  81.                             1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784,

  82.                             1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784,

  83.                             1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196,

  84.                             1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196,

  85.                             1.0 / 196, 1.0 / 196, 1.0 / 196,

  86.                             1.0 / 49, 1.0 / 49, 1.0 / 49, 1.0 / 49, 1.0 / 49, 1.0 / 49, 1.0 / 49, 1.0 / 49, 1.0 / 49,

  87.                             1.0]

  88.         mean = _get_mirror_mean()

  89.         degree = _get_device_num()

  90.         self.grad_reducer_Amax = DistributedGradReducerThor(self.parameters, 2, mean, degree)

  91.         self.grad_reducer_Gmax = DistributedGradReducerThor(self.parameters, 5, mean, degree)

  92.         self.grad_reducer_A = DistributedGradReducerThor(self.parameters, 3, mean, degree)

  93.         self.grad_reducer_G = DistributedGradReducerThor(self.parameters, 4, mean, degree)

  94.         self.matrix_A_inv = ()

  95.         self.matrix_G_inv = ()

  96.         self.matrix_max_inv = ()

  97. 

  98.         for i in range(54):

  99.             self.matrix_max_inv = self.matrix_max_inv + (

  100.                 Parameter(initializer(1, [1], mstype.float32), name="matrix_max" + str(i), requires_grad=False),)

  101.         self.log = P.Log()

  102.         self.exp = P.Exp()

  103.         self.sqrt = P.Sqrt()

  104.         self.matrix_max_inv = ParameterTuple(self.matrix_max_inv)

  105.         self.assign = P.Assign()

  106.         self.cast = P.Cast()

  107.         self.thor = True

  108.         self.weight_decay = weight_decay * loss_scale

  109.         self.decay_flags = tuple(decay_filter(x) for x in self.parameters)

  110. 

  111.     def construct(self, gradients):

  112.         params = self.params

  113.         moments = self.moments

  114.         if self.thor:

  115.             matrix_A_allreduce = ()

  116.             matrix_G_allreduce = ()

  117.             matrix_A_max_allreduce = ()

  118.             matrix_G_max_allreduce = ()

  119.             for i in range(54):

  120.                 g = gradients[i * 3]

  121.                 matrix_A = self.matrix_A[i]

  122.                 matrix_G = self.matrix_G[i]

  123.                 A_max = self.A_inv_max[i]

  124.                 G_max = self.G_inv_max[i]

  125.                 matrix_A = F.depend(matrix_A, g)

  126.                 matrix_G = F.depend(matrix_G, g)

  127.                 A_max = F.depend(A_max, g)

  128.                 G_max = F.depend(G_max, g)

  129.                 matrix_A_allreduce = matrix_A_allreduce + (matrix_A,)

  130.                 matrix_G_allreduce = matrix_G_allreduce + (matrix_G,)

  131.                 matrix_A_max_allreduce = matrix_A_max_allreduce + (A_max,)

  132.                 matrix_G_max_allreduce = matrix_G_max_allreduce + (G_max,)

  133.             matrix_A_allreduce = self.grad_reducer_A(matrix_A_allreduce)

  134.             matrix_G_allreduce = self.grad_reducer_G(matrix_G_allreduce)

  135.             matrix_A_max_allreduce = self.grad_reducer_Amax(matrix_A_max_allreduce)

  136.             matrix_G_max_allreduce = self.grad_reducer_Gmax(matrix_G_max_allreduce)

  137.             new_grads = ()

  138.             for i in range(54):

  139.                 g = gradients[i * 3]

  140.                 temp_a = matrix_A_allreduce[i]

  141.                 temp_g = matrix_G_allreduce[i]

  142.                 temp_a = self.cast(temp_a, mstype.float32)

  143.                 temp_g = self.cast(temp_g, mstype.float32)

  144.                 matrix_A_inv_max = self.log(matrix_A_max_allreduce[i])

  145.                 matrix_A_inv_max = self.mul(matrix_A_inv_max, -1)

  146.                 matrix_A_inv_max = self.exp(matrix_A_inv_max)

  147.                 temp_a = self.mul(temp_a, matrix_A_inv_max)

  148.                 matrix_G_inv_max = self.log(matrix_G_max_allreduce[i])

  149.                 matrix_G_inv_max = self.mul(matrix_G_inv_max, -1)

  150.                 matrix_G_inv_max = self.exp(matrix_G_inv_max)

  151.                 temp_g = self.mul(temp_g, matrix_G_inv_max)

  152.                 temp_max = self.mul(matrix_A_max_allreduce[i], matrix_G_max_allreduce[i])

  153.                 temp_max = self.mul(temp_max, self.feature_map[i])

  154.                 temp_a = self.cast(temp_a, mstype.float16)

  155.                 temp_g = self.cast(temp_g, mstype.float16)

  156.                 if i == 53:

  157.                     g = self.cube_matmul_left_fc(temp_g, g)

  158.                     g = self.cube_matmul_right_fc(g, temp_a, temp_max)

  159.                 else:

  160.                     g = self.cube_matmul_left(temp_g, g)

  161.                     g = self.cube_matmul_right_mul(g, temp_a, temp_max)

  162.                 fake_A = self.assign(self.matrix_A[i], temp_a)

  163.                 fake_G = self.assign(self.matrix_G[i], temp_g)

  164.                 fake_max = self.assign(self.matrix_max_inv[i], temp_max)

  165.                 g = F.depend(g, fake_A)

  166.                 g = F.depend(g, fake_G)

  167.                 g = F.depend(g, fake_max)

  168.                 if i == 53:

  169.                     new_grads = new_grads + (g,)

  170.                 else:

  171.                     new_grads = new_grads + (g, gradients[i * 3 + 1], gradients[i * 3 + 2])

  172.             gradients = new_grads

  173.         else:

  174.             new_grads = ()

  175.             for i in range(54):

  176.                 g = gradients[i * 3]

  177.                 matrix_A = self.matrix_A[i]

  178.                 matrix_G = self.matrix_G[i]

  179.                 matrix_max = self.matrix_max_inv[i]

  180.                 matrix_A = F.depend(matrix_A, g)

  181.                 matrix_G = F.depend(matrix_G, g)

  182.                 matrix_max = F.depend(matrix_max, g)

  183.                 if i == 53:

  184.                     g = self.cube_matmul_left_fc(matrix_G, g)

  185.                     g = self.cube_matmul_right_fc(g, matrix_A, matrix_max)

  186.                     new_grads = new_grads + (g,)

  187.                 else:

  188.                     g = self.cube_matmul_left(matrix_G, g)

  189.                     g = self.cube_matmul_right_mul(g, matrix_A, matrix_max)

  190.                     new_grads = new_grads + (g, gradients[i * 3 + 1], gradients[i * 3 + 2])

  191.             gradients = new_grads

  192. 

  193.         if self.weight_decay > 0:

  194.             gradients = self.hyper_map(F.partial(apply_decay, self.weight_decay), self.decay_flags,

  195.                                        params, gradients)

  196.         gradients = self.scale_grad(gradients)

  197.         lr = self.get_lr()

  198.         success = self.hyper_map(F.partial(momentum_opt, self.opt, lr, self.momentum), gradients, params, moments)

  199.         return success