Merge pull request #52 from kozistr/refactor/optimizers

[Refactor] BaseOptimizer

Author: kozistr

pytorch_optimizer/adabelief.py

@@ -3,10 +3,11 @@
 import torch
 from torch.optim.optimizer import Optimizer
 
+from pytorch_optimizer.base_optimizer import BaseOptimizer
 from pytorch_optimizer.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS, STATE
 
 
-class AdaBelief(Optimizer):
+class AdaBelief(Optimizer, BaseOptimizer):
     """
     Reference : https://github.com/juntang-zhuang/Adabelief-Optimizer
     Example :
@@ -37,7 +38,7 @@ def __init__(
         adamd_debias_term: bool = False,
         eps: float = 1e-16,
     ):
-        """AdaBelief
+        """AdaBelief optimizer
         :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups
         :param lr: float. learning rate
         :param betas: BETAS. coefficients used for computing running averages of gradient and the squared hessian trace
@@ -62,7 +63,7 @@ def __init__(
         self.adamd_debias_term = adamd_debias_term
         self.eps = eps
 
-        self.check_valid_parameters()
+        self.validate_parameters()
 
         defaults: DEFAULTS = dict(
             lr=lr,
@@ -75,17 +76,11 @@ def __init__(
         )
         super().__init__(params, defaults)
 
-    def check_valid_parameters(self):
-        if self.lr < 0.0:
-            raise ValueError(f'Invalid learning rate : {self.lr}')
-        if not 0.0 <= self.betas[0] < 1.0:
-            raise ValueError(f'Invalid beta_0 : {self.betas[0]}')
-        if not 0.0 <= self.betas[1] < 1.0:
-            raise ValueError(f'Invalid beta_1 : {self.betas[1]}')
-        if self.weight_decay < 0.0:
-            raise ValueError(f'Invalid weight_decay : {self.weight_decay}')
-        if self.eps < 0.0:
-            raise ValueError(f'Invalid eps : {self.eps}')
+    def validate_parameters(self):
+        self.validate_learning_rate(self.lr)
+        self.validate_betas(self.betas)
+        self.validate_weight_decay(self.weight_decay)
+        self.validate_epsilon(self.eps)
 
     def __setstate__(self, state: STATE):
         super().__setstate__(state)
@@ -125,7 +120,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                     grad = grad.float()
 
                 p_fp32 = p
-                if p.dtype in {torch.float16, torch.bfloat16}:
+                if p.dtype in (torch.float16, torch.bfloat16):
                     p_fp32 = p_fp32.float()
 
                 state = self.state[p]
@@ -158,14 +153,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 exp_avg_var.mul_(beta2).addcmul_(grad_residual, grad_residual, value=1.0 - beta2)
 
                 if group['amsgrad']:
-                    max_exp_avg_var = state['max_exp_avg_var']
-
-                    torch.max(
-                        max_exp_avg_var,
-                        exp_avg_var.add_(group['eps']),
-                        out=max_exp_avg_var,
-                    )
-
+                    max_exp_avg_var = torch.max(state['max_exp_avg_var'], exp_avg_var.add_(group['eps']))
                     de_nom = (max_exp_avg_var.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
                 else:
                     de_nom = (exp_avg_var.add_(group['eps']).sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
@@ -176,7 +164,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                         step_size /= bias_correction1
                     p_fp32.addcdiv_(exp_avg, de_nom, value=-step_size)
                 else:
-                    buffered = group['buffer'][int(state['step'] % 10)]
+                    buffered = group['buffer'][state['step'] % 10]
                     if state['step'] == buffered[0]:
                         n_sma, step_size = buffered[1], buffered[2]
                     else:
@@ -213,7 +201,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                     elif step_size > 0:
                         p_fp32.add_(exp_avg, alpha=-step_size * group['lr'])
 
-                if p.dtype in {torch.float16, torch.bfloat16}:
+                if p.dtype in (torch.float16, torch.bfloat16):
                     p.copy_(p_fp32)
 
         return loss

pytorch_optimizer/adabound.py

@@ -4,10 +4,11 @@
 import torch
 from torch.optim.optimizer import Optimizer
 
+from pytorch_optimizer.base_optimizer import BaseOptimizer
 from pytorch_optimizer.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS, STATE
 
 
-class AdaBound(Optimizer):
+class AdaBound(Optimizer, BaseOptimizer):
     """
     Reference : https://github.com/Luolc/AdaBound
     Example :
@@ -57,7 +58,7 @@ def __init__(
         self.fixed_decay = fixed_decay
         self.eps = eps
 
-        self.check_valid_parameters()
+        self.validate_parameters()
 
         defaults: DEFAULTS = dict(
             lr=lr,
@@ -73,17 +74,11 @@ def __init__(
 
         self.base_lrs: List[float] = [group['lr'] for group in self.param_groups]
 
-    def check_valid_parameters(self):
-        if self.lr < 0.0:
-            raise ValueError(f'Invalid learning rate : {self.lr}')
-        if self.weight_decay < 0.0:
-            raise ValueError(f'Invalid weight_decay : {self.weight_decay}')
-        if not 0.0 <= self.betas[0] < 1.0:
-            raise ValueError(f'Invalid beta_0 : {self.betas[0]}')
-        if not 0.0 <= self.betas[1] < 1.0:
-            raise ValueError(f'Invalid beta_1 : {self.betas[1]}')
-        if self.eps < 0.0:
-            raise ValueError(f'Invalid eps : {self.eps}')
+    def validate_parameters(self):
+        self.validate_learning_rate(self.lr)
+        self.validate_betas(self.betas)
+        self.validate_weight_decay(self.weight_decay)
+        self.validate_epsilon(self.eps)
 
     def __setstate__(self, state: STATE):
         super().__setstate__(state)

pytorch_optimizer/adahessian.py

@@ -3,10 +3,11 @@
 import torch
 from torch.optim import Optimizer
 
+from pytorch_optimizer.base_optimizer import BaseOptimizer
 from pytorch_optimizer.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
 
 
-class AdaHessian(Optimizer):
+class AdaHessian(Optimizer, BaseOptimizer):
     """
     Reference : https://github.com/davda54/ada-hessian
     Example :
@@ -59,7 +60,7 @@ def __init__(
         self.eps = eps
         self.seed = seed
 
-        self.check_valid_parameters()
+        self.validate_parameters()
 
         # use a separate generator that deterministically generates
         # the same `z`s across all GPUs in case of distributed training
@@ -79,19 +80,12 @@ def __init__(
             p.hess = 0.0
             self.state[p]['hessian_step'] = 0
 
-    def check_valid_parameters(self):
-        if self.lr < 0.0:
-            raise ValueError(f'Invalid learning rate : {self.lr}')
-        if self.weight_decay < 0.0:
-            raise ValueError(f'Invalid weight_decay : {self.weight_decay}')
-        if not 0.0 <= self.betas[0] < 1.0:
-            raise ValueError(f'Invalid beta_0 : {self.betas[0]}')
-        if not 0.0 <= self.betas[1] < 1.0:
-            raise ValueError(f'Invalid beta_1 : {self.betas[1]}')
-        if not 0.0 <= self.hessian_power <= 1.0:
-            raise ValueError(f'Invalid hessian_power : {self.hessian_power}')
-        if self.eps < 0.0:
-            raise ValueError(f'Invalid eps : {self.eps}')
+    def validate_parameters(self):
+        self.validate_learning_rate(self.lr)
+        self.validate_betas(self.betas)
+        self.validate_weight_decay(self.weight_decay)
+        self.validate_hessian_power(self.hessian_power)
+        self.validate_epsilon(self.eps)
 
     def get_params(self) -> Iterable:
         """Gets all parameters in all param_groups with gradients"""
@@ -104,7 +98,9 @@ def zero_hessian(self):
                 p.hess.zero_()
 
     def set_hessian(self):
-        """Computes the Hutchinson approximation of the hessian trace and accumulates it for each trainable parameter"""
+        """Computes the Hutchinson approximation of the hessian trace
+        and accumulates it for each trainable parameter
+        """
         params = []
         for p in self.get_params():
             if p.grad is None:

pytorch_optimizer/adamp.py

@@ -5,11 +5,12 @@
 import torch.nn.functional as F
 from torch.optim.optimizer import Optimizer
 
+from pytorch_optimizer.base_optimizer import BaseOptimizer
 from pytorch_optimizer.gc import centralize_gradient
 from pytorch_optimizer.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
 
 
-class AdamP(Optimizer):
+class AdamP(Optimizer, BaseOptimizer):
     """
     Reference : https://github.com/clovaai/AdamP
     Example :
@@ -58,7 +59,7 @@ def __init__(
         self.wd_ratio = wd_ratio
         self.use_gc = use_gc
 
-        self.check_valid_parameters()
+        self.validate_parameters()
 
         defaults: DEFAULTS = dict(
             lr=lr,
@@ -72,19 +73,12 @@ def __init__(
         )
         super().__init__(params, defaults)
 
-    def check_valid_parameters(self):
-        if self.lr < 0.0:
-            raise ValueError(f'Invalid learning rate : {self.lr}')
-        if not 0.0 <= self.betas[0] < 1.0:
-            raise ValueError(f'Invalid beta_0 : {self.betas[0]}')
-        if not 0.0 <= self.betas[1] < 1.0:
-            raise ValueError(f'Invalid beta_1 : {self.betas[1]}')
-        if self.weight_decay < 0.0:
-            raise ValueError(f'Invalid weight_decay : {self.weight_decay}')
-        if self.eps < 0.0:
-            raise ValueError(f'Invalid eps : {self.eps}')
-        if not 0.0 <= self.wd_ratio < 1.0:
-            raise ValueError(f'Invalid wd_ratio : {self.wd_ratio}')
+    def validate_parameters(self):
+        self.validate_learning_rate(self.lr)
+        self.validate_betas(self.betas)
+        self.validate_weight_decay(self.weight_decay)
+        self.validate_weight_decay_ratio(self.wd_ratio)
+        self.validate_epsilon(self.eps)
 
     @staticmethod
     def channel_view(x: torch.Tensor) -> torch.Tensor:

pytorch_optimizer/base_optimizer.py

@@ -0,0 +1,66 @@
+from abc import ABC, abstractmethod
+
+from pytorch_optimizer.types import BETAS
+
+
+class BaseOptimizer(ABC):
+    @staticmethod
+    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_beta0(beta_0: float):
+        if not 0.0 <= beta_0 < 1.0:
+            raise ValueError(f'[-] beta0 {beta_0} must be in the range [0, 1)')
+
+    @staticmethod
+    def validate_betas(betas: BETAS):
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f'[-] beta1 {betas[0]} must be in the range [0, 1)')
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f'[-] beta2 {betas[1]} must be in the range [0, 1)')
+
+    @staticmethod
+    def validate_weight_decay(weight_decay: float):
+        if weight_decay < 0.0:
+            raise ValueError(f'[-] weight_decay {weight_decay} must be non-negative')
+
+    @staticmethod
+    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:
+            raise ValueError(f'[-] trust_coefficient {trust_coefficient} must be non-negative')
+
+    @staticmethod
+    def validate_momentum(momentum: float):
+        if not 0.0 <= momentum < 1.0:
+            raise ValueError(f'[-] momentum {momentum} must be in the range [0, 1)')
+
+    @staticmethod
+    def validate_lookahead_k(k: int):
+        if k < 0:
+            raise ValueError(f'[-] k {k} must be non-negative')
+
+    @staticmethod
+    def validate_rho(rho: float):
+        if rho < 0.0:
+            raise ValueError(f'[-] rho {rho} must be non-negative')
+
+    @staticmethod
+    def validate_epsilon(epsilon: float):
+        if epsilon < 0.0:
+            raise ValueError(f'[-] epsilon {epsilon} must be non-negative')
+
+    @abstractmethod
+    def validate_parameters(self):
+        raise NotImplementedError

pytorch_optimizer/diffgrad.py

@@ -3,10 +3,11 @@
 import torch
 from torch.optim.optimizer import Optimizer
 
+from pytorch_optimizer.base_optimizer import BaseOptimizer
 from pytorch_optimizer.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS, STATE
 
 
-class DiffGrad(Optimizer):
+class DiffGrad(Optimizer, BaseOptimizer):
     """
     Reference : https://github.com/shivram1987/diffGrad
     Example :
@@ -31,7 +32,7 @@ def __init__(
         weight_decay: float = 0.0,
         adamd_debias_term: bool = False,
     ):
-        """DiffGrad
+        """DiffGrad optimizer
         :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups
         :param lr: float. learning rate
         :param betas: BETAS. coefficients used for computing running averages of gradient and the squared hessian trace
@@ -44,24 +45,18 @@ def __init__(
         self.betas = betas
         self.weight_decay = weight_decay
 
-        self.check_valid_parameters()
+        self.validate_parameters()
 
         defaults: DEFAULTS = dict(
             lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, adamd_debias_term=adamd_debias_term
         )
         super().__init__(params, defaults)
 
-    def check_valid_parameters(self):
-        if self.lr < 0.0:
-            raise ValueError(f'Invalid learning rate : {self.lr}')
-        if self.weight_decay < 0.0:
-            raise ValueError(f'Invalid weight_decay : {self.weight_decay}')
-        if not 0.0 <= self.betas[0] < 1.0:
-            raise ValueError(f'Invalid beta_0 : {self.betas[0]}')
-        if not 0.0 <= self.betas[1] < 1.0:
-            raise ValueError(f'Invalid beta_1 : {self.betas[1]}')
-        if self.eps < 0.0:
-            raise ValueError(f'Invalid eps : {self.eps}')
+    def validate_parameters(self):
+        self.validate_learning_rate(self.lr)
+        self.validate_betas(self.betas)
+        self.validate_weight_decay(self.weight_decay)
+        self.validate_epsilon(self.eps)
 
     def __setstate__(self, state: STATE):
         super().__setstate__(state)