diff --git a/imperative/python/megengine/optimizer/sgd.py b/imperative/python/megengine/optimizer/sgd.py index fe5efda23121dc0a5dc6db3d67ede3d8fcb5580b..a4d122820e69007b4457158f71a439a2ca4bb9de 100644 --- a/imperative/python/megengine/optimizer/sgd.py +++ b/imperative/python/megengine/optimizer/sgd.py @@ -16,7 +16,7 @@ from .optimizer import Optimizer class SGD(Optimizer): r"""Implements stochastic gradient descent. - + Nesterov momentum is based on the formula from `"On the importance of initialization and momentum in deep learning" `_ . @@ -25,6 +25,7 @@ class SGD(Optimizer): parameter groups. lr: learning rate. momentum: momentum factor. Default: 0.0 + nesterov: enables Nesterov momentum. Default: False weight_decay: weight decay (L2 penalty). Default: 0.0 """ @@ -33,6 +34,7 @@ class SGD(Optimizer): params: Union[Iterable[Parameter], dict], lr: float, momentum: float = 0.0, + nesterov: bool = False, weight_decay: float = 0.0, ): assert lr >= 0.0, "Invalid learning rate: {}".format(lr) @@ -40,9 +42,11 @@ class SGD(Optimizer): assert weight_decay >= 0.0, "Invalid weight_decay value: {}".format( weight_decay ) + assert not nesterov or momentum > 0.0, "Nesterov momentum requires a momentum" defaults = dict(lr=lr, momentum=momentum, weight_decay=weight_decay) super().__init__(params, defaults) + self.nesterov = nesterov self._disable_type_convert = True def _create_state(self, param_group): @@ -76,20 +80,22 @@ class SGD(Optimizer): grad = grad + param * _weight_decay if inplace_mode: - if momentum: + if momentum != 0.0: v = self._state[param]["momentum_buffer"] _inplace_add_(v, grad, alpha=_momentum, beta=c1) - _inplace_add_(param, v, alpha=c1, beta=_neg_lr) - else: - _inplace_add_(param, grad, alpha=c1, beta=_neg_lr) + if self.nesterov: + grad = grad + v * _momentum + else: + grad = v + _inplace_add_(param, grad, alpha=c1, beta=_neg_lr) continue - if momentum: + if momentum != 0.0: v = self._state[param]["momentum_buffer"] - # v = v * _momentum + grad v *= _momentum v += grad - - param -= _lr * v - else: - param -= _lr * grad + if self.nesterov: + grad = grad + v * _momentum + else: + grad = v + param -= _lr * grad diff --git a/imperative/python/test/integration/test_optimizer.py b/imperative/python/test/integration/test_optimizer.py index d5ca99a9b77107f0963bca364f4a8ab6d3af85fe..577c09f7397b5c2edc3b5a0bedc02c8fba570f53 100644 --- a/imperative/python/test/integration/test_optimizer.py +++ b/imperative/python/test/integration/test_optimizer.py @@ -124,6 +124,7 @@ def _test_optimizer(opt_str, test_case, check_class, update_lr=False): "case", [ {"momentum": 0.9, "lr": 0.01}, # SGD with momentum + {"momentum": 0.9, "lr": 0.01, "nesterov": True}, # with nesterov momentum {"lr": 0.01}, # simple SGD {"weight_decay": 0.1, "lr": 0.01}, # with weight_decay ], @@ -144,9 +145,12 @@ def test_sgd(monkeypatch, case, update_lr, inplace_mode): grad = param.grad.numpy() if hasattr(self, "weight_decay") and self.weight_decay != 0.0: grad = grad + ori_params[param] * self.weight_decay - if hasattr(self, "momentum"): + if hasattr(self, "momentum") and self.momentum != 0.0: self.slots[param] = grad + self.slots[param] * self.momentum - delta = -self.lr * self.slots[param] + if hasattr(self, "nesterov") and self.nesterov: + delta = -self.lr * (grad + self.slots[param] * self.momentum) + else: + delta = -self.lr * self.slots[param] else: delta = -self.lr * grad np.testing.assert_almost_equal(