Merge pull request #147 from kozistr/feature/pid-optimizer

[Feature] Implement PID optimizer

Author: kozistr

README.rst

@@ -16,7 +16,7 @@ pytorch-optimizer
 
 | **pytorch-optimizer** is optimizer & lr scheduler collections in PyTorch.
 | I just re-implemented (speed & memory tweaks, plug-ins) the algorithm while based on the original paper. Also, It includes useful and practical optimization ideas.
-| Currently, about 40 optimizers, 6 lr schedulers are supported!
+| Currently, 43 optimizers, 6 lr schedulers are supported!
 |
 | Highly inspired by `pytorch-optimizer <https://github.com/jettify/pytorch-optimizer>`__.
 
@@ -171,6 +171,8 @@ You can check the supported optimizers & lr schedulers.
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | QHAdam       | *Quasi-hyperbolic momentum and Adam for deep learning*                                          | `github <https://github.com/facebookresearch/qhoptim>`__                          | `https://arxiv.org/abs/1810.06801 <https://arxiv.org/abs/1810.06801>`__                       |
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| PID          | *A PID Controller Approach for Stochastic Optimization of Deep Networks*                        | `github <https://github.com/tensorboy/PIDOptimizer>`__                            | `CVPR 18 <http://www4.comp.polyu.edu.hk/~cslzhang/paper/CVPR18_PID.pdf>`__                    |
++--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------
@@ -404,6 +406,8 @@ Citations
 
 `QHAdam <https://github.com/facebookresearch/qhoptim#reference>`__
 
+`PID <https://github.com/tensorboy/PIDOptimizer#citation>`__
+
 Citation
 --------
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "2.7.0"
+version = "2.8.0"
 description = "optimizer & lr scheduler 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

@@ -46,6 +46,7 @@
 from pytorch_optimizer.optimizer.nero import Nero
 from pytorch_optimizer.optimizer.novograd import NovoGrad
 from pytorch_optimizer.optimizer.pcgrad import PCGrad
+from pytorch_optimizer.optimizer.pid import PID
 from pytorch_optimizer.optimizer.pnm import PNM
 from pytorch_optimizer.optimizer.qhadam import QHAdam
 from pytorch_optimizer.optimizer.qhm import QHM
@@ -89,6 +90,7 @@
 OPTIMIZER_LIST: List[OPTIMIZER] = [
     AdaBelief,
     AdaBound,
+    PID,
     AdamP,
     Adai,
     Adan,

pytorch_optimizer/optimizer/pid.py

@@ -0,0 +1,121 @@
+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 PID(Optimizer, BaseOptimizer):
+    r"""A PID Controller Approach for Stochastic Optimization of Deep Networks.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :param momentum: float. momentum factor.
+    :param dampening: float. dampening for momentum.
+    :param derivative: float. D part of the PID.
+    :param integral: float. I part of the PID.
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        momentum: float = 0.0,
+        dampening: float = 0.0,
+        derivative: float = 10.0,
+        integral: float = 5.0,
+        weight_decay: float = 0.0,
+        weight_decouple: bool = False,
+    ):
+        self.lr = lr
+        self.momentum = momentum
+        self.dampening = dampening
+        self.derivative = derivative
+        self.integral = integral
+        self.weight_decay = weight_decay
+
+        self.validate_parameters()
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'momentum': momentum,
+            'dampening': dampening,
+            'derivative': derivative,
+            'integral': integral,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+        }
+        super().__init__(params, defaults)
+
+    def validate_parameters(self):
+        self.validate_learning_rate(self.lr)
+        self.validate_momentum(self.momentum)
+        self.validate_weight_decay(self.weight_decay)
+
+    def __str__(self) -> str:
+        return 'PID'
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                state = self.state[p]
+
+                if group['momentum'] > 0.0:
+                    state['grad_buffer'] = torch.zeros_like(p)
+                    state['i_buffer'] = torch.zeros_like(p)
+                    state['d_buffer'] = torch.zeros_like(p)
+
+    @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:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise NoSparseGradientError(str(self))
+
+                state = self.state[p]
+
+                if len(state) == 0 and group['momentum'] > 0.0:
+                    state['grad_buffer'] = torch.zeros_like(p)
+                    state['i_buffer'] = torch.zeros_like(p)
+                    state['d_buffer'] = torch.zeros_like(p)
+
+                if group['weight_decouple']:
+                    p.mul_(1.0 - group['weight_decay'] * group['lr'])
+                elif group['weight_decay'] > 0.0:
+                    grad.add_(p, alpha=group['weight_decay'])
+
+                if group['momentum'] > 0.0:
+                    i_buf = state['i_buffer']
+                    i_buf.mul_(group['momentum']).add_(grad, alpha=1.0 - group['dampening'])
+
+                    g_buf, d_buf = state['grad_buffer'], state['d_buffer']
+                    d_buf.mul_(group['momentum'])
+
+                    if group['step'] > 1:
+                        d_buf.add_(grad - g_buf, alpha=1.0 - group['momentum'])
+                        g_buf.copy_(grad)
+
+                    grad.add_(i_buf, alpha=group['integral']).add_(d_buf, alpha=group['derivative'])
+
+                p.add_(grad, alpha=-group['lr'])
+
+        return loss

tests/constants.py

@@ -6,6 +6,7 @@
     MADGRAD,
     MSVAG,
     OPTIMIZERS,
+    PID,
     PNM,
     QHM,
     SGDP,
@@ -346,6 +347,8 @@
     (QHAdam, {'lr': 1e0, 'weight_decay': 1e-3, 'weight_decouple': True}, 5),
     (QHM, {'lr': 1e0, 'weight_decay': 1e-3}, 5),
     (QHM, {'lr': 1e0, 'weight_decay': 1e-3, 'weight_decouple': True}, 5),
+    (PID, {'lr': 1e0, 'momentum': 0.9, 'dampening': 1.0, 'weight_decay': 1e-3}, 5),
+    (PID, {'lr': 1e0, 'momentum': 0.9, 'dampening': 1.0, 'weight_decay': 1e-3, 'weight_decouple': True}, 5),
 ]
 ADANORM_SUPPORTED_OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
     (AdaBelief, {'lr': 5e-1, 'weight_decay': 1e-3, 'adanorm': True}, 10),

tests/test_general_optimizer_parameters.py

@@ -36,6 +36,7 @@ def test_epsilon(optimizer_name):
         'msvag',
         'aggmo',
         'qhm',
+        'pid',
     ):
         pytest.skip(f'skip {optimizer_name} optimizer')
 

tests/test_load_optimizers.py

@@ -16,4 +16,4 @@ def test_load_optimizers_invalid(invalid_optimizer_names):
 
 
 def test_get_supported_optimizers():
-    assert len(get_supported_optimizers()) == 42
+    assert len(get_supported_optimizers()) == 43