Merge pull request #59 from kozistr/feature/optimizers

[Feature] Remove AdaHessian optimizer

Author: kozistr

README.rst

@@ -76,6 +76,8 @@ Supported Optimizers
 +--------------+----------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | Shampoo      | *Preconditioned Stochastic Tensor Optimization*                                        | `github <https://github.com/moskomule/shampoo.pytorch>`__                         | `https://arxiv.org/abs/1802.09568 <https://arxiv.org/abs/1802.09568>`__                       |
 +--------------+----------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| Nero         | *Learning by Turning: Neural Architecture Aware Optimisation*                          | `github <https://github.com/jxbz/nero>`__                                         | `https://arxiv.org/abs/2102.07227 <https://arxiv.org/abs/2102.07227>`__                       |
++--------------+----------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------
@@ -482,6 +484,17 @@ Shampoo: Preconditioned Stochastic Tensor Optimization
       organization={PMLR}
     }
 
+Nero: Learning by Turning: Neural Architecture Aware Optimisation
+
+::
+
+    @misc{nero2021,
+      title={Learning by Turning: Neural Architecture Aware Optimisation},
+      author={Yang Liu and Jeremy Bernstein and Markus Meister and Yisong Yue},
+      year={2021},
+      eprint={arXiv:2102.07227}
+    }
+
 Author
 ------
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "1.0.0"
+version = "1.1.0"
 description = "Bunch of optimizer implementations in PyTorch with clean-code, strict types. Also, including useful optimization ideas."
 license = "Apache-2.0"
 authors = ["kozistr <[email protected]>"]

pytorch_optimizer/__init__.py

@@ -1,7 +1,6 @@
 # pylint: disable=unused-import
 from pytorch_optimizer.adabelief import AdaBelief
 from pytorch_optimizer.adabound import AdaBound
-from pytorch_optimizer.adahessian import AdaHessian
 from pytorch_optimizer.adamp import AdamP
 from pytorch_optimizer.adapnm import AdaPNM
 from pytorch_optimizer.agc import agc
@@ -14,6 +13,7 @@
 from pytorch_optimizer.lars import LARS
 from pytorch_optimizer.lookahead import Lookahead
 from pytorch_optimizer.madgrad import MADGRAD
+from pytorch_optimizer.nero import Nero
 from pytorch_optimizer.optimizers import load_optimizers
 from pytorch_optimizer.pcgrad import PCGrad
 from pytorch_optimizer.pnm import PNM

pytorch_optimizer/adahessian.py

@@ -11,6 +11,11 @@ def validate_learning_rate(learning_rate: float):
         if learning_rate < 0.0:
             raise ValueError(f'[-] learning rate {learning_rate} must be positive')
 
+    @staticmethod
+    def validate_beta(beta: float):
+        if not 0.0 <= beta <= 1.0:
+            raise ValueError(f'[-] beta {beta} must be in the range [0, 1]')
+
     @staticmethod
     def validate_beta0(beta_0: float):
         if not 0.0 <= beta_0 <= 1.0:
@@ -39,11 +44,6 @@ def validate_weight_decay_ratio(weight_decay_ratio: float):
         if not 0.0 <= weight_decay_ratio < 1.0:
             raise ValueError(f'[-] weight_decay_ratio {weight_decay_ratio} must be in the range [0, 1)')
 
-    @staticmethod
-    def validate_hessian_power(hessian_power: float):
-        if not 0.0 <= hessian_power <= 1.0:
-            raise ValueError(f'[-] hessian_power {hessian_power} must be in the range [0, 1]')
-
     @staticmethod
     def validate_trust_coefficient(trust_coefficient: float):
         if trust_coefficient < 0.0:

pytorch_optimizer/base_optimizer.py

@@ -0,0 +1,103 @@
+import torch
+from torch.optim.optimizer import Optimizer
+
+from pytorch_optimizer.base_optimizer import BaseOptimizer
+from pytorch_optimizer.types import CLOSURE, DEFAULTS, LOSS, PARAMETERS
+from pytorch_optimizer.utils import neuron_mean, neuron_norm
+
+
+class Nero(Optimizer, BaseOptimizer):
+    """
+    Reference : https://github.com/jxbz/nero
+    Example :
+        from pytorch_optimizer import Nero
+        ...
+        model = YourModel()
+        optimizer = Nero(model.parameters())
+        ...
+        for input, output in data:
+          optimizer.zero_grad()
+          loss = loss_function(output, model(input))
+          loss.backward()
+          optimizer.step()
+    """
+
+    def __init__(self, params: PARAMETERS, lr: float = 0.01, beta: float = 0.999, constraints: bool = True):
+        """AdamP optimizer
+        :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups
+        :param lr: float. learning rate
+        :param beta: float. coefficients used for computing running averages of gradient and the squared hessian trace
+        :param constraints: bool.
+        """
+        self.lr = lr
+        self.beta = beta
+
+        self.validate_parameters()
+
+        defaults: DEFAULTS = dict(lr=lr, constraints=constraints)
+        super().__init__(params, defaults)
+
+    def validate_parameters(self):
+        self.validate_learning_rate(self.lr)
+        self.validate_beta(self.beta)
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            for p in group['params']:
+                if group['constraints'] and p.dim() > 1:
+                    p.sub_(neuron_mean(p))
+                    p.div_(neuron_norm(p))
+
+                state = self.state[p]
+
+                state['step'] = 0
+                state['exp_avg_sq'] = torch.zeros_like(neuron_norm(p))
+                state['scale'] = neuron_norm(p).mean()
+
+                if state['scale'] == 0.0:
+                    state['scale'] = 0.01
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None) -> LOSS:
+        loss: LOSS = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError('Nero does not support sparse gradients')
+
+                state = self.state[p]
+                if len(state) == 0:
+                    if group['constraints'] and p.dim() > 1:
+                        p.sub_(neuron_mean(p))
+                        p.div_(neuron_norm(p))
+
+                    state['step'] = 0
+                    state['exp_avg_sq'] = torch.zeros_like(neuron_norm(p))
+                    state['scale'] = neuron_norm(p).mean()
+                    if state['scale'] == 0.0:
+                        state['scale'] = 0.01
+
+                state['step'] += 1
+
+                bias_correction: float = 1.0 - self.beta ** state['step']
+                state['exp_avg_sq'] = self.beta * state['exp_avg_sq'] + (1.0 - self.beta) * neuron_norm(grad) ** 2
+
+                grad_normed = grad / (state['exp_avg_sq'] / bias_correction).sqrt()
+                grad_normed[torch.isnan(grad_normed)] = 0.0
+
+                p.sub_(group['lr'] * state['scale'] * grad_normed)
+
+                if group['constraints'] and p.dim() > 1:
+                    p.sub_(neuron_mean(p))
+                    p.div_(neuron_norm(p))
+
+        return loss