diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/fused_ada_factor_cpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/cpu/fused_ada_factor_cpu_kernel.h index c51770378a0..c9430aade41 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/fused_ada_factor_cpu_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/fused_ada_factor_cpu_kernel.h @@ -50,8 +50,8 @@ class FusedAdaFactorCPUKernel : public CPUKernel { bool enable_weight_decay_{false}; bool need_factor_{false}; size_t elem_num_{0}; - size_t last_row_dim_size_{0}; - size_t last_col_dim_size_{0}; + size_t last_row_dim_size_{1}; + size_t last_col_dim_size_{1}; TypeId param_dtype_{kTypeUnknown}; enum InputEnum { diff --git a/mindspore/python/mindspore/nn/optim/adafactor.py b/mindspore/python/mindspore/nn/optim/adafactor.py index 7644ed51f66..ba6197a8334 100644 --- a/mindspore/python/mindspore/nn/optim/adafactor.py +++ b/mindspore/python/mindspore/nn/optim/adafactor.py @@ -13,6 +13,7 @@ # limitations under the License. # ============================================================================ """adafactor""" +from mindspore import context from mindspore.common import dtype as mstype from mindspore.log import logging from mindspore.common.initializer import initializer @@ -27,22 +28,6 @@ from mindspore.nn.optim.optimizer import opt_init_args_register from .optimizer import Optimizer -def _get_lr(step, rms, learning_rate, relative_step, warmup_init, scale_parameter, eps): - """update optimizer learning rete""" - rel_step_sz = learning_rate - if relative_step: - if warmup_init: - min_step = 1e-6 * step * 1.0 - else: - min_step = 1e-2 * 1.0 - - rel_step_sz = P.Minimum()(min_step, 1.0 / P.Sqrt()(step * 1.0)) - param_scale = 1.0 - if scale_parameter: - param_scale = P.Maximum()(eps[1], rms) - return rel_step_sz * param_scale * F.ones_like(rms) - - def _rms(update_tensor): """calculate rms""" return F.sqrt(P.ReduceMean(False)(F.square(update_tensor))) @@ -59,18 +44,14 @@ def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col): return P.Mul()(r_factor, c_factor) -_adam_opt = C.MultitypeFuncGraph("adam_opt") +_adafactor_opt = C.MultitypeFuncGraph("adafactor_opt") -@_adam_opt.register("Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Bool", "Bool", - "Bool", "Bool", "Bool", "Bool", "Bool", "Tensor", "Tensor", "Tensor", - "Tensor", "Tensor", "Tensor", "Tensor", "Tensor") -def _run_opt_with_one_number(eps, clip_threshold, decay_rate, beta1, - weight_decay, scale_lr, scale_parameter, relative_step, - warmup_init, compression, use_first_moment, weight_decay_flag, - learning_rate, step, grad, param, - exp_avg, exp_avg_sq_row, - exp_avg_sq_col, exp_avg_sq): +@_adafactor_opt.register("Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Bool", "Bool", "Bool", "Bool", "Tensor", + "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor") +def _run_opt_with_one_number(eps, clip_threshold, beta1, beta2t, weight_decay, scale_parameter, + compression, use_first_moment, weight_decay_flag, learning_rate, + grad, param, exp_avg, exp_avg_sq_row, exp_avg_sq_col, exp_avg_sq): """Apply ada factor optimizer to the weight parameter using Tensor.""" success = True grad_dtype = F.dtype(grad) @@ -84,38 +65,24 @@ def _run_opt_with_one_number(eps, clip_threshold, decay_rate, beta1, factored = len(grad_shape) >= 2 - # State Initialization - exp_avg_update = exp_avg - exp_avg_sq_update = exp_avg_sq - exp_avg_sq_row_update = exp_avg_sq_row - exp_avg_sq_col_update = exp_avg_sq_col - - if use_first_moment: - if compression: - exp_avg_update = F.cast(exp_avg, mstype.float16) - - if factored: - exp_avg_sq_row_update = F.cast(exp_avg_sq_row, grad_dtype) - exp_avg_sq_col_update = F.cast(exp_avg_sq_col, grad_dtype) - else: - exp_avg_sq_update = F.cast(exp_avg_sq, grad_dtype) - - if scale_lr: + if scale_parameter: rms = _rms(p_data_fp32) - learning_rate_update = _get_lr(step, rms, learning_rate, relative_step, warmup_init, scale_parameter, eps) + param_scale = P.Maximum()(eps[1], rms) + learning_rate_update = learning_rate * param_scale * F.ones_like(rms) learning_rate_update = F.assign(learning_rate, F.cast(learning_rate_update, F.dtype(learning_rate))) else: - learning_rate_update = learning_rate * 1.0 + learning_rate_update = learning_rate - beta2t = 1.0 - P.Pow()(step, decay_rate) update = (grad ** 2) + eps[0] if factored: + exp_avg_sq_row_update = F.cast(exp_avg_sq_row, grad_dtype) exp_avg_sq_row_update = P.Mul()(exp_avg_sq_row_update, beta2t) update_mean = P.ReduceMean()(update, -1) * (1.0 - beta2t) exp_avg_sq_row_update = P.Add()(exp_avg_sq_row_update, update_mean) exp_avg_sq_row_update = F.assign(exp_avg_sq_row, F.cast(exp_avg_sq_row_update, F.dtype(exp_avg_sq_row))) + exp_avg_sq_col_update = F.cast(exp_avg_sq_col, grad_dtype) exp_avg_sq_col_update = P.Mul()(exp_avg_sq_col_update, beta2t) update_mean = P.ReduceMean()(update, -2) * (1.0 - beta2t) exp_avg_sq_col_update = P.Add()(exp_avg_sq_col_update, update_mean) @@ -124,6 +91,7 @@ def _run_opt_with_one_number(eps, clip_threshold, decay_rate, beta1, update = _approx_sq_grad(exp_avg_sq_row_update, exp_avg_sq_col_update) update = P.Mul()(update, grad) else: + exp_avg_sq_update = F.cast(exp_avg_sq, grad_dtype) update = update * (1.0 - beta2t) exp_avg_sq_update = P.Add()(P.Mul()(exp_avg_sq_update, beta2t), update) exp_avg_sq_update = F.assign(exp_avg_sq, F.cast(exp_avg_sq_update, F.dtype(exp_avg_sq))) @@ -135,8 +103,9 @@ def _run_opt_with_one_number(eps, clip_threshold, decay_rate, beta1, update = P.Mul()(P.Div()(update, update_coff), learning_rate_update) if use_first_moment: + exp_avg_update = exp_avg if compression: - exp_avg_update = F.cast(exp_avg_update, grad_dtype) + exp_avg_update = F.cast(exp_avg, grad_dtype) exp_avg_update = P.Add()(P.Mul()(exp_avg_update, beta1), update * (1 - beta1)) update = F.assign(exp_avg, F.cast(exp_avg_update, F.dtype(exp_avg))) @@ -144,18 +113,27 @@ def _run_opt_with_one_number(eps, clip_threshold, decay_rate, beta1, p_data_fp32_coff = p_data_fp32 * -weight_decay * learning_rate_update p_data_fp32 = P.Add()(p_data_fp32, p_data_fp32_coff) p_data_fp32 = P.Sub()(p_data_fp32, update) - P.Assign()(param, F.cast(p_data_fp32, F.dtype(param))) - return success + return F.depend(success, P.Assign()(param, F.cast(p_data_fp32, F.dtype(param)))) -def trans_to_tensor(paras, is_tuple=False, fp32=True): - if paras is None or isinstance(paras, bool): - return paras +@_adafactor_opt.register("Function", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", + "Tensor", "Tensor", "Tensor", "Tensor") +def _run_fused_ada_factor(fused_ada_factor, eps, clip_threshold, beta1, beta2t, weight_decay, learning_rate, + grad, param, exp_avg, exp_avg_sq_row, exp_avg_sq_col, exp_avg_sq): + success = True + ret = fused_ada_factor(eps, clip_threshold, beta1, beta2t, weight_decay, learning_rate, + grad, param, exp_avg, exp_avg_sq_row, exp_avg_sq_col, exp_avg_sq) + return F.depend(success, ret) + + +def trans_to_tensor(param, is_tuple=False, fp32=True): + if param is None or isinstance(param, bool): + return param data_type = mstype.float32 if fp32 else mstype.float16 if is_tuple: - new_paras = [Tensor(ele, data_type) for ele in paras] - return tuple(new_paras) - return Tensor(paras, data_type) + new_param = [Tensor(ele, data_type) for ele in param] + return tuple(new_param) + return Tensor(param, data_type) class AdaFactor(Optimizer): @@ -344,9 +322,17 @@ class AdaFactor(Optimizer): self.relative_step = relative_step self.warmup_init = warmup_init self.compression = compression - + if not self.scale_lr: + self.scale_parameter = False self.init_ada_factor_state(beta1) self.step = Parameter(initializer(0, [1], mstype.float32), name='afactor_step') + self.fused_ada_factor = P.FusedAdaFactor(enable_scale_parameter=self.scale_parameter, + enable_first_moment=self.use_first_moment, + enable_weight_decay=self.weight_decay_flag) + if context.get_context("device_target") == "CPU": + self.use_fused_ada_factor = True + else: + self.use_fused_ada_factor = False print("AdaFactor init completed", self.learning_rate) def init_ada_factor_state(self, beta1): @@ -361,35 +347,31 @@ class AdaFactor(Optimizer): self.exp_avg_sq = [] self.exp_avg_sq_col = [] self.exp_avg_sq_row = [] - for paras in self.parameters: - paras_dtype = paras.dtype - paras_shape = paras.shape - paras_name = paras.name - if len(paras_shape) > 1: - self.exp_avg_sq_row.append(Parameter(initializer(0, shape=paras_shape[:-1], dtype=paras_dtype), - name="exp_avg_sq_row_{}".format(paras_name))) - self.exp_avg_sq_col.append(Parameter(initializer(0, shape=paras_shape[:-2] + paras_shape[-1:], - dtype=paras_dtype), - name="exp_avg_sq_col_{}".format(paras_name))) - if self.compression: - self.exp_avg_sq.append(Parameter(initializer(0, shape=(1,), dtype=mstype.float16), - name="exp_avg_sq_{}".format(paras_name))) - else: - self.exp_avg_sq.append(Parameter(initializer(0, shape=(1,), dtype=paras_dtype), - name="exp_avg_sq_{}".format(paras_name))) + for param in self.parameters: + param_dtype = param.dtype + param_shape = param.shape + param_name = param.name + if len(param_shape) > 1: + self.exp_avg_sq_row.append(Parameter(initializer(0, shape=param_shape[:-1], dtype=param_dtype), + name="exp_avg_sq_row_{}".format(param_name))) + self.exp_avg_sq_col.append(Parameter(initializer(0, shape=param_shape[:-2] + param_shape[-1:], + dtype=param_dtype), + name="exp_avg_sq_col_{}".format(param_name))) + self.exp_avg_sq.append(Parameter(initializer(0, shape=(1,), dtype=param_dtype), + name="exp_avg_sq_{}".format(param_name))) else: - self.exp_avg_sq_row.append(Parameter(initializer(0, shape=(1,), dtype=paras_dtype), - name="exp_avg_sq_row_{}".format(paras_name))) - self.exp_avg_sq_col.append(Parameter(initializer(0, shape=(1,), dtype=paras_dtype), - name="exp_avg_sq_col_{}".format(paras_name))) + self.exp_avg_sq_row.append(Parameter(initializer(0, shape=(1,), dtype=param_dtype), + name="exp_avg_sq_row_{}".format(param_name))) + self.exp_avg_sq_col.append(Parameter(initializer(0, shape=(1,), dtype=param_dtype), + name="exp_avg_sq_col_{}".format(param_name))) if self.compression: - self.exp_avg_sq.append(Parameter(initializer(0, shape=paras_shape, dtype=mstype.float16), - name="exp_avg_sq_{}".format(paras_name))) + self.exp_avg_sq.append(Parameter(initializer(0, shape=param_shape, dtype=mstype.float16), + name="exp_avg_sq_{}".format(param_name))) else: - self.exp_avg_sq.append(Parameter(initializer(0, shape=paras_shape, dtype=paras_dtype), - name="exp_avg_sq_{}".format(paras_name))) + self.exp_avg_sq.append(Parameter(initializer(0, shape=param_shape, dtype=param_dtype), + name="exp_avg_sq_{}".format(param_name))) self.exp_avg_sq_row = ParameterTuple(self.exp_avg_sq_row) self.exp_avg_sq_col = ParameterTuple(self.exp_avg_sq_col) @@ -406,13 +388,25 @@ class AdaFactor(Optimizer): def construct(self, gradients): lr = self.get_lr() step = F.assign_add(self.step, 1) - success = self.hyper_map(F.partial(_adam_opt, self.eps, self.clip_threshold, self.decay_rate, - self.beta1, self.weight_decay, self.scale_lr, - self.scale_parameter, self.relative_step, - self.warmup_init, self.compression, self.use_first_moment, - self.weight_decay_flag, lr, step), - gradients, self.parameters, self.exp_avg, self.exp_avg_sq_row, - self.exp_avg_sq_col, self.exp_avg_sq) + if self.scale_lr and self.relative_step: + if self.warmup_init: + min_step = 1e-6 * step + else: + min_step = 1e-2 + lr = P.Minimum()(min_step, 1.0 / P.Sqrt()(step * 1.0)) + beta2t = 1.0 - P.Pow()(step, self.decay_rate) + + if self.use_fused_ada_factor: + success = self.hyper_map(F.partial(_adafactor_opt, self.fused_ada_factor, self.eps, self.clip_threshold, + self.beta1, beta2t, self.weight_decay, lr), + gradients, self.parameters, self.exp_avg, self.exp_avg_sq_row, + self.exp_avg_sq_col, self.exp_avg_sq) + else: + success = self.hyper_map(F.partial(_adafactor_opt, self.eps, self.clip_threshold, self.beta1, beta2t, + self.weight_decay, self.scale_parameter, self.compression, + self.use_first_moment, self.weight_decay_flag, lr), + gradients, self.parameters, self.exp_avg, self.exp_avg_sq_row, + self.exp_avg_sq_col, self.exp_avg_sq) return success @@ -423,3 +417,8 @@ class AdaFactor(Optimizer): optimizer operation. """ self._set_base_target(value) + if value == 'CPU': + self.fused_ada_factor.add_prim_attr("primitive_target", "CPU") + self.use_fused_ada_factor = True + else: + self.use_fused_ada_factor = False diff --git a/mindspore/python/mindspore/ops/operations/__init__.py b/mindspore/python/mindspore/ops/operations/__init__.py index 0868345149d..b60a1477f02 100644 --- a/mindspore/python/mindspore/ops/operations/__init__.py +++ b/mindspore/python/mindspore/ops/operations/__init__.py @@ -44,7 +44,7 @@ from .debug_ops import (ImageSummary, InsertGradientOf, HookBackward, ScalarSumm from .control_ops import GeSwitch, Merge from .inner_ops import (ScalarCast, Randperm, NoRepeatNGram, LambApplyOptimizerAssign, LambApplyWeightAssign, MakeRefKey, - FusedWeightScaleApplyMomentum, FusedCastAdamWeightDecay) + FusedWeightScaleApplyMomentum, FusedCastAdamWeightDecay, FusedAdaFactor) from .math_ops import (Abs, ACos, Asin, Asinh, AddN, AccumulateNV2, AssignAdd, AssignSub, Atan2, BatchMatMul, BitwiseAnd, BitwiseOr, Ger, @@ -177,6 +177,7 @@ __all__ = [ 'FusedSparseAdam', 'FusedSparseLazyAdam', 'AdamNoUpdateParam', + 'FusedAdaFactor', 'Softplus', 'Softmax', 'Softsign', diff --git a/mindspore/python/mindspore/ops/operations/inner_ops.py b/mindspore/python/mindspore/ops/operations/inner_ops.py index 0637ce0c069..2699948039e 100755 --- a/mindspore/python/mindspore/ops/operations/inner_ops.py +++ b/mindspore/python/mindspore/ops/operations/inner_ops.py @@ -254,6 +254,7 @@ class LambApplyOptimizerAssign(PrimitiveWithInfer): Supported Platforms: ``Ascend`` """ + @prim_attr_register def __init__(self): """Initialize LambApplyOptimizerAssign""" @@ -316,6 +317,7 @@ class LambApplyWeightAssign(PrimitiveWithInfer): Supported Platforms: ``Ascend`` """ + @prim_attr_register def __init__(self): """Initialize LambApplyWeightAssign""" @@ -558,3 +560,132 @@ class FusedCastAdamWeightDecay(PrimitiveWithInfer): "decay": decay_dtype} validator.check_scalar_or_tensor_types_same(args, [mstype.float32], self.name, True) return var_dtype, m_dtype, v_dtype + + +class FusedAdaFactor(PrimitiveWithInfer): + r""" + Updates gradients by the Adaptive Learning Rates with Sublinear Memory Cost (Adafactor) algorithm. + + The Adafactor algorithm is proposed in `Adafactor: Adafactor: Adaptive Learning Rates with Sublinear Memory + Cost `_. + + .. warning:: + This is an experimental prototype that is subject to change and/or deletion. + + Adafactor for weight vector are as follows, + + .. math:: + \begin{array}{l} \\ + \alpha_{t}=\max \left(\epsilon_{2}, \operatorname{RMS}\left(X_{t-1}\right)\right) \rho_{t} \\ + G_{t}=\nabla f_{t}\left(X_{t-1}\right) \\ + \hat{V}_{t}=\hat{\beta}_{2} \hat{V}_{t-1}+\left(1-\hat{\beta}_{2_{t}}\right)\left(G_{t}^{2}+ \\ + \epsilon_{1} 1_{n}\right) \\ + U_{t}=G_{t} / \sqrt{\hat{V}_{t}} \\ + \hat{U}_{t}=U_{t} / \max \left(1, \operatorname{RMS}\left(U_{t}\right) / d\right) \\ + X_{t}=X_{t-1}-\alpha_{t} \hat{U}_{t} + \end{array} + + Adafactor for weight matrices are as follows, + + .. math:: + \begin{array}{l} \\ + \alpha_{t}=\max \left(\epsilon_{2}, \operatorname{RMS}\left(X_{t-1}\right)\right) \rho_{t} \\ + G_{t}=\nabla f_{t}\left(X_{t-1}\right) \\ + R_{t}=\hat{\beta}_{2 t} R_{t-1}+\left(1-\hat{\beta}_{2 t}\right)\left(G_{t}^{2}+ \\ + \epsilon_{1} 1_{n} 1_{m}^{\top}\right) 1_{m} \\ + C_{t}=\hat{\beta}_{2 t} C_{t-1}+\left(1-\hat{\beta}_{2 t}\right) 1_{n}^{\top}\left(G_{t}^{2}+ \\ + \epsilon_{1} 1_{n} 1_{m}^{\top}\right) \\ + \hat{V}_{t}=R_{t} C_{t} / 1_{n}^{\top} R_{t} \\ + U_{t}=G_{t} / \sqrt{\hat{V}_{t}} \\ + \hat{U}_{t}=U_{t} / \max \left(1, \operatorname{RMS}\left(U_{t}\right) / d\right) \\ + X_{t}=X_{t-1}-\alpha_{t} U_{t} + \end{array} + + Where RMS is: + + .. math:: + \operatorname{RMS}\left(U_{t}\right)=\operatorname{RMS}_{x \in X}\left(u_{x t}\right)= \\ + \sqrt{\operatorname{Mean}_{x \in X}\left(\frac{\left(g_{x t}\right)^{2}}{\hat{v}_{x t}}\right)} + + :math:`x` is each individual parameter, + :math:`t` is assumed to be the current number of steps, + :math:`a_{t}` is the learning rate, + :math:`f(X)` is the loss function, + :math:`\epsilon1` and :math:`\epsilon2` is a small positive number to prevent errors, + :math:`d` is the clipping threshold, + :math:`\beta_{2}` is the moment decay, + :math:`\rho` is the relative step size, + :math:`R` is the running averages of the row sums of the squared gradient, + :math:`C` is the running averages of the column sums of the squared gradient. + + Args: + enable_weight_decay (bool): If True, enable weight decay. default: False + enable_first_moment (bool): If True, enable first moment. default: False + enable_scale_parameter (bool): If True, enable scale learning rate using parameter. default: False + + Inputs: + - **epsilon** (Tensor) - input epsilon pair. + - **clip_threshold** (float) - The threshold of root mean square of final gradient update. + - **beta1** (float) - The exponential decay rate for the 1nd moment estimations. + - **beta2** (float) - The exponential decay rate for the 2nd moment estimations. + - **weight_decay** (float) - The weight decay value, must be a scalar tensor with float data type. + - **learning_rate** (float) - The learning rate value. + - **gradient** (Tensor) - Gradient. + - **param** (Tensor) - Weights to be updated. + - **exp_avg** (Tensor) - The exponential moving average of 1st moment optimizer state. + - **exp_avg_sq_row** (Tensor) - The exponential moving average of square of gradient square row factor. + - **exp_avg_sq_col** (Tensor) - The exponential moving average of square of gradient square col factor. + - **exp_avg_sq** (Tensor) - The exponential moving average of square of gradient square. + + Outputs: + - **dummy_param** (Tensor) - The same shape and data type as `param`. + + Supported Platforms: + ``CPU`` + + Examples: + >>> import numpy as np + >>> import mindspore.context as context + >>> import mindspore.nn as nn + >>> import mindspore.ops as ops + >>> from mindspore import Tensor, Parameter + >>> from mindspore import dtype as mstype + >>> param_shape = [2, 3, 2] + >>> class Net(nn.Cell): + ... def __init__(self): + ... super(Net, self).__init__() + ... self.opt = ops.FusedAdaFactor() + ... self.param = Parameter(Tensor(np.ones(param_shape), mstype.float32), name="param") + ... self.exp_avg = Parameter(Tensor(np.zeros(param_shape), mstype.float32), name="exp_avg") + ... self.exp_avg_sq = Parameter(Tensor(np.zeros(param_shape), mstype.float32), name="exp_avg_sq") + ... self.exp_avg_sq_row = Parameter(Tensor(np.zeros([2, 3]), mstype.float32), name="exp_avg_sq_row") + ... self.exp_avg_sq_col = Parameter(Tensor(np.zeros([2, 2]), mstype.float32), name="exp_avg_sq_col") + ... + ... def construct(self, epsilon, clip_threshold, beta1, beta2, weight_decay, lr, grad): + ... out = self.opt(epsilon, clip_threshold, beta1, beta2, weight_decay, lr, grad, self.param, + ... self.exp_avg, self.exp_avg_sq_row, self.exp_avg_sq_col, self.exp_avg_sq) + ... return out + >>> context.set_context(mode=context.GRAPH_MODE, device_target="CPU") + >>> net = Net() + >>> gradient = Tensor(np.ones(param_shape), mstype.float32) + >>> net((1e-30, 1e-3), 1.0, 0.9, 0.8, 1e-2, 0.03, gradient) + >>> print(net.param.asnumpy()) + """ + + @prim_attr_register + def __init__(self, enable_scale_parameter=False, enable_first_moment=False, enable_weight_decay=False): + self.add_prim_attr('side_effect_mem', True) + validator.check_value_type("enable_scale_parameter", enable_scale_parameter, [bool], self.name) + validator.check_value_type("enable_first_moment", enable_first_moment, [bool], self.name) + validator.check_value_type("enable_weight_decay", enable_weight_decay, [bool], self.name) + + def infer_shape(self, epsilon_shape, clip_threshold_shape, beta1_shape, beta2t_shape, weight_decay_shape, + learning_rate_shape, grad_shape, param_shape, exp_avg_shape, exp_avg_sq_row_shape, + exp_avg_sq_col_shape, exp_avg_sq_shape): + validator.check("grad_shape", grad_shape, "param_shape", param_shape, Rel.EQ, self.name) + return param_shape + + def infer_dtype(self, epsilon_type, clip_threshold_type, beta1_type, beta2t_type, weight_decay_type, + learning_rate_type, grad_type, param_type, exp_avg_type, exp_avg_sq_row_type, + exp_avg_sq_col_type, exp_avg_sq_type): + return param_type diff --git a/tests/st/ops/cpu/test_fused_ada_factor_op.py b/tests/st/ops/cpu/test_fused_ada_factor_op.py new file mode 100644 index 00000000000..7f8ef4297d1 --- /dev/null +++ b/tests/st/ops/cpu/test_fused_ada_factor_op.py @@ -0,0 +1,56 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0(the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http: // www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# == == == == == == == == == == == == == == == == == == == == == == == == == == == == == == == == == == == == == == +import pytest +import numpy as np +import mindspore.context as context +import mindspore.nn as nn +import mindspore.ops as ops +from mindspore import Tensor, Parameter +from mindspore import dtype as mstype + +param_shape = [2, 3, 2] + + +class Net(nn.Cell): + def __init__(self): + super(Net, self).__init__() + self.opt = ops.FusedAdaFactor() + self.param = Parameter(Tensor(np.ones(param_shape), mstype.float32), name="param") + self.exp_avg = Parameter(Tensor(np.zeros(param_shape), mstype.float32), name="exp_avg") + self.exp_avg_sq = Parameter(Tensor(np.zeros(param_shape), mstype.float32), name="exp_avg_sq") + self.exp_avg_sq_row = Parameter(Tensor(np.zeros([2, 3]), mstype.float32), name="exp_avg_sq_row") + self.exp_avg_sq_col = Parameter(Tensor(np.zeros([2, 2]), mstype.float32), name="exp_avg_sq_col") + + def construct(self, epsilon, clip_threshold, beta1, beta2, weight_decay, lr, grad): + out = self.opt(epsilon, clip_threshold, beta1, beta2, weight_decay, lr, grad, self.param, self.exp_avg, + self.exp_avg_sq_row, self.exp_avg_sq_col, self.exp_avg_sq) + return out + + +@pytest.mark.level0 +@pytest.mark.platform_x86_cpu +@pytest.mark.env_onecard +def test_adafactor(): + ''' + Feature: AdaFactor + Description: Test AdaFactor + Expectation: Run success + ''' + context.set_context(mode=context.GRAPH_MODE, device_target="CPU") + net = Net() + gradient = Tensor(np.ones(param_shape), mstype.float32) + net((1e-30, 1e-3), 1.0, 0.9, 0.8, 1e-2, 0.03, gradient) + diff = net.param.asnumpy() - np.ones(param_shape) * 0.97 + assert np.all(diff < 1e-3) diff --git a/tests/st/optimizer/test_fused_adafactor_cpu.py b/tests/st/optimizer/test_fused_adafactor_cpu.py new file mode 100644 index 00000000000..3148b9c1db2 --- /dev/null +++ b/tests/st/optimizer/test_fused_adafactor_cpu.py @@ -0,0 +1,51 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +import numpy as np +import pytest + +import mindspore.context as context +import mindspore.nn as nn +from mindspore import Tensor +from mindspore.nn import TrainOneStepCell, WithLossCell +from tests.st.networks.models.lenet import LeNet + +context.set_context(mode=context.GRAPH_MODE, device_target="CPU") + + +@pytest.mark.level0 +@pytest.mark.platform_x86_cpu +@pytest.mark.env_onecard +def test_lenet(): + ''' + Feature: AdaFactor + Description: Test AdaFactor + Expectation: Run lenet success + ''' + data = Tensor(np.ones([32, 3, 32, 32]).astype(np.float32) * 0.01) + label = Tensor(np.ones([32]).astype(np.int32)) + net = LeNet() + net.batch_size = 32 + learning_rate = 0.01 + optimizer = nn.AdaFactor(filter(lambda x: x.requires_grad, net.get_parameters()), learning_rate, + scale_parameter=False, relative_step=False, beta1=0) + criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean') + net_with_criterion = WithLossCell(net, criterion) + train_network = TrainOneStepCell(net_with_criterion, optimizer) # optimizer + train_network.set_train() + loss = [] + for _ in range(10): + res = train_network(data, label) + loss.append(res.asnumpy()) + assert np.all(loss[-1] < 0.1)