diff --git a/mindspore/nn/optim/adam.py b/mindspore/nn/optim/adam.py index 9ae14312472..055eaae7c65 100755 --- a/mindspore/nn/optim/adam.py +++ b/mindspore/nn/optim/adam.py @@ -31,8 +31,8 @@ _learning_rate_update_func = ['linear', 'cos', 'sin'] adam_opt = C.MultitypeFuncGraph("adam_opt") -@adam_opt.register("Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor") -def _update_run_op(beta1, beta2, eps, lr, weight_decay_tensor, param, m, v, gradient): +@adam_opt.register("Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Bool") +def _update_run_op(beta1, beta2, eps, lr, weight_decay_tensor, param, m, v, gradient, decay_flag): """ Update parameters. @@ -67,7 +67,8 @@ def _update_run_op(beta1, beta2, eps, lr, weight_decay_tensor, param, m, v, grad next_v = op_mul(beta2, v) + op_mul(op_cast(F.tuple_to_array((1.0,)), mstype.float32) - beta2, op_square(gradient)) update = next_m / (op_sqrt(next_v) + eps) - update = update + op_mul(weight_decay_tensor, param) + if decay_flag: + update = update + op_mul(weight_decay_tensor, param) update_with_lr = op_mul(lr, update) next_param = param - op_reshape(update_with_lr, op_shape(param)) @@ -90,6 +91,17 @@ def _check_param_value(beta1, beta2, eps, weight_decay, prim_name): validator.check_number_range("weight_decay", weight_decay, 0.0, float("inf"), Rel.INC_LEFT, prim_name) +def _check_learning_rate_value(learning_rate, end_learning_rate, decay_steps, power, prim_name): + """Check the type of inputs.""" + validator.check_float_positive('learning_rate', learning_rate, prim_name) + validator.check_float_legal_value('learning_rate', learning_rate, prim_name) + validator.check_float_positive('end_learning_rate', end_learning_rate, prim_name) + validator.check_float_legal_value('end_learning_rate', end_learning_rate, prim_name) + validator.check_float_positive('power', power, prim_name) + validator.check_float_legal_value('power', power, prim_name) + validator.check_integer('decay_steps', decay_steps, 0, Rel.GT, prim_name) + + @adam_opt.register("Function", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Number", "Tensor", "Tensor", "Tensor", "Tensor") def _run_opt_with_one_number(opt, lr, beta1_power, beta2_power, beta1, beta2, eps, gradient, params, moment1, @@ -126,8 +138,13 @@ class Adam(Optimizer): Args: params (list[Parameter]): A list of parameter, which will be updated. The element in `params` should be class mindspore.Parameter. - learning_rate (Union[float, Tensor, Iterable]): The Learning rate. - Iterable type is used for the dynamic learning rate. + learning_rate (Union[float, Tensor, Iterable]): A value for the learning rate. When the learning_rate is + Iterable or a Tensor and the dims of the Tensor is 1, + use dynamic learning rate, then the i-th step will + take the i-th value as the learning rate. + When the learning_rate is float or learning_rate is a Tensor + but the dims of the Tensor is 0, use fixed learning rate. + Other cases are not supported. Default: 1e-3. beta1 (float): The exponential decay rate for the 1st moment estimates. Should be in range (0.0, 1.0). beta2 (float): The exponential decay rate for the 2nd moment estimates. Should be in range (0.0, 1.0). eps (float): Term added to the denominator to improve numerical stability. Should be greater than 0. @@ -140,6 +157,8 @@ class Adam(Optimizer): weight_decay (float): Weight decay (L2 penalty). Default: 0.0. loss_scale (float): A floating point value for the loss scale. Default: 1.0. Should be equal to or greater than 1. + decay_filter (Function): A function to determine whether to apply weight decay on parameters. Default: + lambda x: 'LayerNorm' not in x.name and 'bias' not in x.name. Inputs: - **gradients** (tuple[Tensor]) - The gradients of `params`, the shape is the same as `params`. @@ -207,7 +226,13 @@ class AdamWeightDecay(Optimizer): Args: params (list[Parameter]): A list of parameter, which will be updated. The element in `params` should be class mindspore.Parameter. - learning_rate (float): A floating point value for the learning rate. Default: 1e-3. + learning_rate (Union[float, Tensor, Iterable]): A value for the learning rate. When the learning_rate is + Iterable or a Tensor and the dims of the Tensor is 1, + use dynamic learning rate, then the i-th step will + take the i-th value as the learning rate. + When the learning_rate is float or learning_rate is a Tensor + but the dims of the Tensor is 0, use fixed learning rate. + Other cases are not supported. Default: 1e-3. beta1 (float): The exponential decay rate for the 1st moment estimates. Default: 0.9. Should be in range (0.0, 1.0). beta2 (float): The exponential decay rate for the 2nd moment estimates. Default: 0.999. @@ -215,6 +240,8 @@ class AdamWeightDecay(Optimizer): eps (float): Term added to the denominator to improve numerical stability. Default: 1e-6. Should be greater than 0. weight_decay (float): Weight decay (L2 penalty). Default: 0.0. + decay_filter (Function): A function to determine whether to apply weight decay on parameters. Default: + lambda x: 'LayerNorm' not in x.name and 'bias' not in x.name. Inputs: - **gradients** (tuple[Tensor]) - The gradients of `params`, the shape is the same as `params`. @@ -228,10 +255,10 @@ class AdamWeightDecay(Optimizer): >>> optim = nn.AdamWeightDecay(params=net.trainable_params()) >>> model = Model(net, loss_fn=loss, optimizer=optim, metrics=None) """ - def __init__(self, params, learning_rate=1e-3, beta1=0.9, beta2=0.999, eps=1e-6, weight_decay=0.0): + def __init__(self, params, learning_rate=1e-3, beta1=0.9, beta2=0.999, eps=1e-6, weight_decay=0.0, + decay_filter=lambda x: 'beta' not in x.name and 'gamma' not in x.name): super(AdamWeightDecay, self).__init__(learning_rate, params) _check_param_value(beta1, beta2, eps, weight_decay, self.cls_name) - self.lr = Tensor(np.array([learning_rate]).astype(np.float32)) self.beta1 = Tensor(np.array([beta1]).astype(np.float32)) self.beta2 = Tensor(np.array([beta2]).astype(np.float32)) self.eps = Tensor(np.array([eps]).astype(np.float32)) @@ -240,13 +267,15 @@ class AdamWeightDecay(Optimizer): self.params = self.parameters self.moments1 = self.params.clone(prefix="adam_m", init='zeros') self.moments2 = self.params.clone(prefix="adam_v", init='zeros') + self.decay_flag = tuple(decay_filter(x) for x in self.params) self.hyper_map = C.HyperMap() def construct(self, gradients): - updated_velocity = self.hyper_map(F.partial(adam_opt, self.beta1, self.beta2, self.eps, self.lr, + lr = self.get_lr() + updated_velocity = self.hyper_map(F.partial(adam_opt, self.beta1, self.beta2, self.eps, lr, self.weight_decay_tensor), - self.params, self.moments1, self.moments2, gradients) + self.params, self.moments1, self.moments2, gradients, self.decay_flag) return updated_velocity @@ -269,6 +298,8 @@ class AdamWeightDecayDynamicLR(Optimizer): eps (float): Term added to the denominator to improve numerical stability. Default: 1e-6. Should be greater than 0. weight_decay (float): Weight decay (L2 penalty). Default: 0.0. + decay_filter (Function): A function to determine whether to apply weight decay on parameters. Default: + lambda x: 'LayerNorm' not in x.name and 'bias' not in x.name. Inputs: - **gradients** (tuple[Tensor]) - The gradients of `params`, the shape is the same as `params`. @@ -291,10 +322,11 @@ class AdamWeightDecayDynamicLR(Optimizer): beta1=0.9, beta2=0.999, eps=1e-6, - weight_decay=0.0): + weight_decay=0.0, + decay_filter=lambda x: 'beta' not in x.name and 'gamma' not in x.name): super(AdamWeightDecayDynamicLR, self).__init__(learning_rate, params) _check_param_value(beta1, beta2, eps, weight_decay, self.cls_name) - + _check_learning_rate_value(learning_rate, end_learning_rate, decay_steps, power, self.cls_name) # turn them to scalar when me support scalar/tensor mix operations self.global_step = Parameter(initializer(0, [1]), name="global_step") self.decay_steps = Tensor(np.array([decay_steps]).astype(np.float32)) @@ -308,7 +340,7 @@ class AdamWeightDecayDynamicLR(Optimizer): self.params = self.parameters self.moments1 = self.params.clone(prefix="adam_m", init='zeros') self.moments2 = self.params.clone(prefix="adam_v", init='zeros') - + self.decay_flag = tuple(decay_filter(x) for x in self.params) self.hyper_map = C.HyperMap() self.min = P.Minimum() self.pow = P.Pow() @@ -320,7 +352,7 @@ class AdamWeightDecayDynamicLR(Optimizer): lr = self.diff_learning_rate * self.pow(self.one - p, self.power) + self.end_learning_rate updated_velocity = self.hyper_map(F.partial(adam_opt, self.beta1, self.beta2, self.eps, lr, self.weight_decay_tensor), - self.params, self.moments1, self.moments2, gradients) + self.params, self.moments1, self.moments2, gradients, self.decay_flag) added_global_step = self.global_step + self.one F.control_depend(lr, added_global_step) diff --git a/mindspore/nn/optim/lamb.py b/mindspore/nn/optim/lamb.py index afcbf8cda47..01ec9844530 100755 --- a/mindspore/nn/optim/lamb.py +++ b/mindspore/nn/optim/lamb.py @@ -112,16 +112,18 @@ def _check_param_value(decay_steps, warmup_steps, start_learning_rate, end_learning_rate, power, beta1, beta2, eps, weight_decay, prim_name): """Check the type of inputs.""" - validator.check_value_type("decay_steps", decay_steps, [int], prim_name) - validator.check_value_type("warmup_steps", warmup_steps, [int], prim_name) - validator.check_value_type("start_learning_rate", start_learning_rate, [float], prim_name) - validator.check_value_type("end_learning_rate", end_learning_rate, [float], prim_name) - validator.check_value_type("power", power, [float], prim_name) + validator.check_float_positive('start_learning_rate', start_learning_rate, prim_name) + validator.check_float_legal_value('start_learning_rate', start_learning_rate, prim_name) + validator.check_float_positive('end_learning_rate', end_learning_rate, prim_name) + validator.check_float_legal_value('end_learning_rate', end_learning_rate, prim_name) + validator.check_float_positive('power', power, prim_name) + validator.check_float_legal_value('power', power, prim_name) + validator.check_integer('decay_steps', decay_steps, 0, Rel.GT, prim_name) + validator.check_integer('warmup_steps', decay_steps, 0, Rel.GT, prim_name) validator.check_value_type("beta1", beta1, [float], prim_name) validator.check_value_type("beta2", beta2, [float], prim_name) validator.check_value_type("eps", eps, [float], prim_name) validator.check_value_type("weight_dacay", weight_decay, [float], prim_name) - validator.check_number_range("decay_steps", decay_steps, 1, float("inf"), Rel.INC_LEFT, prim_name) validator.check_number_range("beta1", beta1, 0.0, 1.0, Rel.INC_NEITHER, prim_name) validator.check_number_range("beta2", beta2, 0.0, 1.0, Rel.INC_NEITHER, prim_name) validator.check_number_range("eps", eps, 0.0, float("inf"), Rel.INC_NEITHER, prim_name) diff --git a/mindspore/nn/optim/sgd.py b/mindspore/nn/optim/sgd.py index cda5aa904a9..bf2ed21d50f 100755 --- a/mindspore/nn/optim/sgd.py +++ b/mindspore/nn/optim/sgd.py @@ -42,7 +42,13 @@ class SGD(Optimizer): Args: params (list[Parameter]): A list of parameter, which will be updated. The element in `params` should be class mindspore.Parameter. - learning_rate (float): A floating point value for the learning rate. Default: 0.1. + learning_rate (Union[float, Tensor, Iterable]): A value for the learning rate. When the learning_rate is + Iterable or a Tensor and the dims of the Tensor is 1, + use dynamic learning rate, then the i-th step will + take the i-th value as the learning rate. + When the learning_rate is float or learning_rate is a Tensor + but the dims of the Tensor is 0, use fixed learning rate. + Other cases are not supported. Default: 0.1. momentum (float): A floating point value the momentum. Default: 0. dampening (float): A floating point value of dampening for momentum. Default: 0. weight_decay (float): Weight decay (L2 penalty). Default: 0.