Merge pull request #116 from kozistr/feature/optimizers

[Feature] Implement Ali-G optimizer

Author: kozistr

README.rst

@@ -65,8 +65,25 @@ Also, you can load the optimizer via `torch.hub`
     opt = torch.hub.load('kozistr/pytorch_optimizer', 'adamp')
     optimizer = opt(model.parameters())
 
+If you want to build the optimizer with parameters & configs, there's `create_optimizer()` API.
 
-And you can check the supported optimizers & lr schedulers.
+::
+
+    from pytorch_optimizer import create_optimizer
+
+    optimizer = create_optimizer(
+        model,
+        'adamp',
+        lr=1e-2,
+        weight_decay=1e-3,
+        use_gc=True,
+        use_lookahead=True,
+    )
+
+Supported Optimizers
+--------------------
+
+You can check the supported optimizers & lr schedulers.
 
 ::
 
@@ -75,10 +92,6 @@ And you can check the supported optimizers & lr schedulers.
     supported_optimizers = get_supported_optimizers()
     supported_lr_schedulers = get_supported_lr_schedulers()
 
-
-Supported Optimizers
---------------------
-
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | Optimizer    | Description                                                                                     | Official Code                                                                     | Paper                                                                                         |
 +==============+=================================================================================================+===================================================================================+===============================================================================================+
@@ -124,6 +137,8 @@ Supported Optimizers
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | Lion         | *Symbolic Discovery of Optimization Algorithms*                                                 | `github <https://github.com/google/automl/tree/master/lion>`__                    | `https://arxiv.org/abs/2302.06675 <https://arxiv.org/abs/2302.06675>`__                       |
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| Ali-G        | *Adaptive Learning Rates for Interpolation with Gradients*                                      | `github <https://github.com/oval-group/ali-g>`__                                  | `https://arxiv.org/abs/1906.05661 <https://arxiv.org/abs/1906.05661>`__                       |
++--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------
@@ -327,6 +342,8 @@ Citations
 
 `Lion <https://github.com/google/automl/tree/master/lion#citation>`__
 
+`Ali-G <https://github.com/oval-group/ali-g#adaptive-learning-rates-for-interpolation-with-gradients>`__
+
 Citation
 --------
 
@@ -338,7 +355,7 @@ Or you can get from "cite this repository" button.
     @software{Kim_pytorch_optimizer_Bunch_of_2022,
         author = {Kim, Hyeongchan},
         month = {1},
-        title = {{pytorch_optimizer: Bunch of optimizer implementations in PyTorch with clean-code, strict types}},
+        title = {{pytorch_optimizer: optimizer & lr scheduler implementations in PyTorch}},
         version = {1.0.0},
         year = {2022}
     }

docs/optimizer_api.rst

@@ -272,3 +272,11 @@ Lion
 
 .. autoclass:: pytorch_optimizer.Lion
     :members:
+
+.. _AliG:
+
+AliG
+----
+
+.. autoclass:: pytorch_optimizer.AliG
+    :members:

pyproject.toml

@@ -89,6 +89,7 @@ target-version = "py39"
 "./tests/test_load_lr_schedulers.py" = ["D", "S101"]
 "./tests/test_lr_schedulers.py" = ["D"]
 "./tests/test_lr_scheduler_parameters.py" = ["D", "S101"]
+"./tests/test_create_optimizer.py" = ["D"]
 "./pytorch_optimizer/__init__.py" = ["F401"]
 "./pytorch_optimizer/lr_scheduler/__init__.py" = ["F401"]
 

pytorch_optimizer/__init__.py

@@ -1,7 +1,9 @@
 # ruff: noqa
 from typing import Dict, List
 
-from pytorch_optimizer.base.types import OPTIMIZER, SCHEDULER
+from torch import nn
+
+from pytorch_optimizer.base.types import OPTIMIZER, PARAMETERS, SCHEDULER
 from pytorch_optimizer.experimental.deberta_v3_lr_scheduler import deberta_v3_large_lr_scheduler
 from pytorch_optimizer.lr_scheduler import (
     ConstantLR,
@@ -23,6 +25,7 @@
 from pytorch_optimizer.optimizer.adan import Adan
 from pytorch_optimizer.optimizer.adapnm import AdaPNM
 from pytorch_optimizer.optimizer.agc import agc
+from pytorch_optimizer.optimizer.alig import AliG
 from pytorch_optimizer.optimizer.apollo import Apollo
 from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptSGD
 from pytorch_optimizer.optimizer.diffgrad import DiffGrad
@@ -100,6 +103,7 @@
     Apollo,
     NovoGrad,
     Lion,
+    AliG,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 
@@ -129,6 +133,49 @@ def load_optimizer(optimizer: str) -> OPTIMIZER:
     return OPTIMIZERS[optimizer]
 
 
+def create_optimizer(
+    model: nn.Module,
+    optimizer_name: str,
+    lr: float = 1e-3,
+    weight_decay: float = 0.0,
+    wd_ban_list: List[str] = ('bias', 'LayerNorm.bias', 'LayerNorm.weight'),
+    use_lookahead: bool = False,
+    **kwargs,
+):
+    r"""Build optimizer.
+
+    :param model: nn.Module. model.
+    :param optimizer_name: str. name of optimizer.
+    :param lr: float. learning rate.
+    :param weight_decay: float. weight decay.
+    :param wd_ban_list: List[str]. weight decay ban list by layer.
+    :param use_lookahead: bool. use lookahead.
+    """
+    optimizer_name = optimizer_name.lower()
+
+    if weight_decay > 0.0:
+        parameters = get_optimizer_parameters(model, weight_decay, wd_ban_list)
+    else:
+        parameters = model.parameters()
+
+    optimizer = load_optimizer(optimizer_name)
+
+    if optimizer_name == 'alig':
+        optimizer = optimizer(parameters, max_lr=lr, **kwargs)
+    else:
+        optimizer = optimizer(parameters, lr=lr, **kwargs)
+
+    if use_lookahead:
+        optimizer = Lookahead(
+            optimizer,
+            k=kwargs['k'] if 'k' in kwargs else 5,
+            alpha=kwargs['alpha'] if 'alpha' in kwargs else 0.5,
+            pullback_momentum=kwargs['pullback_momentum'] if 'pullback_momentum' in kwargs else 'none',
+        )
+
+    return optimizer
+
+
 def load_lr_scheduler(lr_scheduler: str) -> SCHEDULER:
     lr_scheduler: str = lr_scheduler.lower()
 

pytorch_optimizer/optimizer/adams.py

@@ -39,12 +39,7 @@ def __init__(
 
         self.validate_parameters()
 
-        defaults: DEFAULTS = {
-            'lr': lr,
-            'betas': betas,
-            'weight_decay': weight_decay,
-            'eps': eps,
-        }
+        defaults: DEFAULTS = {'lr': lr, 'betas': betas, 'weight_decay': weight_decay, 'eps': eps}
         super().__init__(params, defaults)
 
     def validate_parameters(self):

pytorch_optimizer/optimizer/alig.py

@@ -0,0 +1,115 @@
+from typing import Callable, Optional
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, PARAMETERS
+
+
+class AliG(Optimizer, BaseOptimizer):
+    r"""Adaptive Learning Rates for Interpolation with Gradients.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param max_lr: Optional[float]. max learning rate.
+    :param projection_fn : Callable. projection function to enforce constraints.
+    :param momentum: float. momentum.
+    :param adjusted_momentum: bool. if True, use pytorch-like momentum, instead of standard Nesterov momentum.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        max_lr: Optional[float] = None,
+        projection_fn: Optional[Callable] = None,
+        momentum: float = 0.0,
+        adjusted_momentum: bool = False,
+        eps: float = 1e-5,
+    ):
+        self.max_lr = max_lr
+        self.projection_fn = projection_fn
+        self.momentum = momentum
+        self.adjusted_momentum = adjusted_momentum
+        self.eps = eps
+
+        self.validate_parameters()
+
+        defaults: DEFAULTS = {'max_lr': max_lr, 'momentum': momentum}
+        super().__init__(params, defaults)
+
+        if self.projection_fn is not None:
+            self.projection_fn()
+
+    def validate_parameters(self):
+        self.validate_momentum(self.momentum)
+        self.validate_epsilon(self.eps)
+
+    @property
+    def __str__(self) -> str:
+        return 'AliG'
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            for p in group['params']:
+                state = self.state[p]
+
+                state['momentum_buffer'] = torch.zeros_like(p)
+
+    @torch.no_grad()
+    def compute_step_size(self, loss: float) -> float:
+        r"""Compute step_size."""
+        global_grad_norm: float = 0
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is not None:
+                    global_grad_norm += p.grad.norm().pow(2).item()
+
+        return loss / (global_grad_norm + self.eps)
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None) -> LOSS:
+        if closure is None:
+            raise ValueError('[-] AliG optimizer needs closure. (eg. `optimizer.step(lambda: float(loss))`).')
+
+        loss = closure()
+
+        un_clipped_step_size: float = self.compute_step_size(loss)
+
+        for group in self.param_groups:
+            step_size = group['step_size'] = (
+                min(un_clipped_step_size, group['max_lr']) if group['max_lr'] is not None else un_clipped_step_size
+            )
+            momentum = group['momentum']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise NoSparseGradientError(self.__str__)
+
+                state = self.state[p]
+                if len(state) == 0 and momentum > 0.0:
+                    state['momentum_buffer'] = torch.zeros_like(p)
+
+                p.add_(grad, alpha=-step_size)
+
+                if momentum > 0.0:
+                    buffer = state['momentum_buffer']
+
+                    if self.adjusted_momentum:
+                        buffer.mul_(momentum).sub_(grad)
+                        p.add_(buffer, alpha=step_size * momentum)
+                    else:
+                        buffer.mul_(momentum).add_(grad, alpha=-step_size)
+                        p.add_(buffer, alpha=momentum)
+
+            if self.projection_fn is not None:
+                self.projection_fn()
+
+        return loss

pytorch_optimizer/optimizer/gsam.py

@@ -33,7 +33,7 @@ def loss_fn(predictions, targets):
                 lr_scheduler.step()
                 optimizer.update_rho_t()
 
-    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
     :param base_optimizer: Optimizer. base optimizer.
     :param model: nn.Module. model.
     :param alpha: float. rho alpha.

pytorch_optimizer/optimizer/lion.py

@@ -31,11 +31,7 @@ def __init__(
 
         self.validate_parameters()
 
-        defaults: DEFAULTS = {
-            'lr': lr,
-            'betas': betas,
-            'weight_decay': weight_decay,
-        }
+        defaults: DEFAULTS = {'lr': lr, 'betas': betas, 'weight_decay': weight_decay}
         super().__init__(params, defaults)
 
     def validate_parameters(self):

pytorch_optimizer/optimizer/madgrad.py

@@ -42,12 +42,7 @@ def __init__(
 
         self.validate_parameters()
 
-        defaults: DEFAULTS = {
-            'lr': lr,
-            'weight_decay': weight_decay,
-            'momentum': momentum,
-            'eps': eps,
-        }
+        defaults: DEFAULTS = {'lr': lr, 'weight_decay': weight_decay, 'momentum': momentum, 'eps': eps}
         super().__init__(params, defaults)
 
     def validate_parameters(self):

pytorch_optimizer/optimizer/novograd.py

@@ -39,12 +39,7 @@ def __init__(
 
         self.validate_parameters()
 
-        defaults: DEFAULTS = {
-            'lr': lr,
-            'betas': betas,
-            'weight_decay': weight_decay,
-            'eps': eps,
-        }
+        defaults: DEFAULTS = {'lr': lr, 'betas': betas, 'weight_decay': weight_decay, 'eps': eps}
         super().__init__(params, defaults)
 
     def validate_parameters(self):