Merge pull request #54 from kozistr/refactor/optimizers

kozistr · web-flow · commit d65706010435 · 2022-03-05T16:35:50.000+09:00
[Refactor] Optimizers
diff --git a/lint.py b/lint.py
@@ -14,7 +14,7 @@ def get_configuration() -> Namespace:
     parser.add_argument(
         '-t',
         '--threshold',
-        default=9.9,
+        default=9.95,
         type=float,
     )
 
diff --git a/pytorch_optimizer/adabelief.py b/pytorch_optimizer/adabelief.py
@@ -82,12 +82,6 @@ def validate_parameters(self):
         self.validate_weight_decay(self.weight_decay)
         self.validate_epsilon(self.eps)
 
-    def __setstate__(self, state: STATE):
-        super().__setstate__(state)
-        for group in self.param_groups:
-            group.setdefault('amsgrad', False)
-            group.setdefault('adamd_debias_term', False)
-
     @torch.no_grad()
     def reset(self):
         for group in self.param_groups:
@@ -152,11 +146,11 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 grad_residual = grad - exp_avg
                 exp_avg_var.mul_(beta2).addcmul_(grad_residual, grad_residual, value=1.0 - beta2)
 
+                exp_avg_var = exp_avg_var.add_(group['eps'])
                 if group['amsgrad']:
-                    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'])
+                    exp_avg_var = torch.max(state['max_exp_avg_var'], exp_avg_var)
+
+                de_nom = (exp_avg_var.add_(group['eps']).sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
 
                 if not self.rectify:
                     step_size = group['lr']
diff --git a/pytorch_optimizer/adabound.py b/pytorch_optimizer/adabound.py
@@ -80,11 +80,6 @@ def validate_parameters(self):
         self.validate_weight_decay(self.weight_decay)
         self.validate_epsilon(self.eps)
 
-    def __setstate__(self, state: STATE):
-        super().__setstate__(state)
-        for group in self.param_groups:
-            group.setdefault('amsbound', False)
-
     @torch.no_grad()
     def reset(self):
         for group in self.param_groups:
@@ -140,10 +135,9 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
 
                 if group['amsbound']:
-                    max_exp_avg_sq = torch.max(state['max_exp_avg_sq'], exp_avg_sq)
-                    de_nom = max_exp_avg_sq.sqrt().add_(group['eps'])
-                else:
-                    de_nom = exp_avg_sq.sqrt().add_(group['eps'])
+                    exp_avg_sq = torch.max(state['max_exp_avg_sq'], exp_avg_sq)
+
+                de_nom = exp_avg_sq.sqrt().add_(group['eps'])
 
                 bias_correction1 = 1.0 - beta1 ** state['step']
                 bias_correction2 = 1.0 - beta2 ** state['step']
diff --git a/pytorch_optimizer/diffgrad.py b/pytorch_optimizer/diffgrad.py
@@ -4,7 +4,7 @@
 from torch.optim.optimizer import Optimizer
 
 from pytorch_optimizer.base_optimizer import BaseOptimizer
-from pytorch_optimizer.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS, STATE
+from pytorch_optimizer.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
 
 
 class DiffGrad(Optimizer, BaseOptimizer):
diff --git a/pytorch_optimizer/diffrgrad.py b/pytorch_optimizer/diffrgrad.py
@@ -4,7 +4,7 @@
 from torch.optim.optimizer import Optimizer
 
 from pytorch_optimizer.base_optimizer import BaseOptimizer
-from pytorch_optimizer.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS, STATE
+from pytorch_optimizer.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
 
 
 class DiffRGrad(Optimizer, BaseOptimizer):
diff --git a/pytorch_optimizer/lars.py b/pytorch_optimizer/lars.py
@@ -89,7 +89,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                     grad = grad.add(p, alpha=g['weight_decay'])
                     param_norm = torch.norm(p)
                     update_norm = torch.norm(grad)
-                    one = torch.ones_like(param_norm)
+                    one = torch.ones_like(param_norm, device=param_norm.device)
 
                     q = torch.where(
                         param_norm > 0.0,
@@ -100,7 +100,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                 param_state = self.state[p]
                 if 'mu' not in param_state:
-                    param_state['mu'] = torch.zeros_like(p)
+                    param_state['mu'] = torch.zeros_like(p, device=p.device)
 
                 mu = param_state['mu']
                 mu.mul_(g['momentum']).add_(grad)
diff --git a/pytorch_optimizer/pcgrad.py b/pytorch_optimizer/pcgrad.py
@@ -39,15 +39,15 @@ def validate_parameters(self):
 
     @torch.no_grad()
     def reset(self):
-        pass
+        self.zero_grad()
 
     def zero_grad(self):
         return self.optimizer.zero_grad(set_to_none=True)
 
     def step(self):
         return self.optimizer.step()
 
-    def set_grad(self, grads):
+    def set_grad(self, grads: List[torch.Tensor]):
         idx: int = 0
         for group in self.optimizer.param_groups:
             for p in group['params']:
@@ -74,7 +74,7 @@ def retrieve_grad(self) -> Tuple[List[torch.Tensor], List[int], List[torch.Tenso
     def pack_grad(self, objectives: Iterable) -> Tuple[List[torch.Tensor], List[List[int]], List[torch.Tensor]]:
         """pack the gradient of the parameters of the network for each objective
         :param objectives: Iterable[nn.Module]. a list of objectives
-        :return:
+        :return: torch.Tensor. packed gradients
         """
         grads, shapes, has_grads = [], [], []
         for objective in objectives:
@@ -89,27 +89,29 @@ def pack_grad(self, objectives: Iterable) -> Tuple[List[torch.Tensor], List[List
 
         return grads, shapes, has_grads
 
-    def project_conflicting(self, grads, has_grads) -> torch.Tensor:
+    def project_conflicting(self, grads: List[torch.Tensor], has_grads: List[torch.Tensor]) -> torch.Tensor:
         """project conflicting
         :param grads: a list of the gradient of the parameters
         :param has_grads: a list of mask represent whether the parameter has gradient
-        :return:
+        :return: torch.Tensor. merged gradients
         """
-        shared = torch.stack(has_grads).prod(0).bool()
+        shared: torch.Tensor = torch.stack(has_grads).prod(0).bool()
 
-        pc_grad = deepcopy(grads)
+        pc_grad: List[torch.Tensor] = deepcopy(grads)
         for g_i in pc_grad:
             random.shuffle(grads)
             for g_j in grads:
-                g_i_g_j = torch.dot(g_i, g_j)
+                g_i_g_j: torch.Tensor = torch.dot(g_i, g_j)
                 if g_i_g_j < 0:
                     g_i -= g_i_g_j * g_j / (g_j.norm() ** 2)
 
-        merged_grad = torch.zeros_like(grads[0]).to(grads[0].device)
+        merged_grad: torch.Tensor = torch.zeros_like(grads[0], device=grads[0].device)
+
+        shared_pc_gradients: torch.Tensor = torch.stack([g[shared] for g in pc_grad])
         if self.reduction == 'mean':
-            merged_grad[shared] = torch.stack([g[shared] for g in pc_grad]).mean(dim=0)
+            merged_grad[shared] = shared_pc_gradients.mean(dim=0)
         else:
-            merged_grad[shared] = torch.stack([g[shared] for g in pc_grad]).sum(dim=0)
+            merged_grad[shared] = shared_pc_gradients.sum(dim=0)
 
         merged_grad[~shared] = torch.stack([g[~shared] for g in pc_grad]).sum(dim=0)
 
@@ -121,7 +123,7 @@ def pc_backward(self, objectives: Iterable[nn.Module]):
         :return:
         """
         grads, shapes, has_grads = self.pack_grad(objectives)
+
         pc_grad = self.project_conflicting(grads, has_grads)
         pc_grad = un_flatten_grad(pc_grad, shapes[0])
-
         self.set_grad(pc_grad)
diff --git a/pytorch_optimizer/radam.py b/pytorch_optimizer/radam.py
@@ -1,5 +1,4 @@
 import math
-from typing import Dict
 
 import torch
 from torch.optim.optimizer import Optimizer
@@ -153,14 +152,13 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                         step_size = -1
                     buffered[2] = step_size
 
+                if group['weight_decay'] != 0 and (n_sma >= self.n_sma_threshold or step_size > 0):
+                    p_fp32.add_(p_fp32, alpha=-group['weight_decay'] * group['lr'])
+
                 if n_sma >= self.n_sma_threshold:
-                    if group['weight_decay'] != 0:
-                        p_fp32.add_(p_fp32, alpha=-group['weight_decay'] * group['lr'])
                     de_nom = exp_avg_sq.sqrt().add_(group['eps'])
                     p_fp32.addcdiv_(exp_avg, de_nom, value=-step_size * group['lr'])
                 elif step_size > 0:
-                    if group['weight_decay'] != 0:
-                        p_fp32.add_(p_fp32, alpha=-group['weight_decay'] * group['lr'])
                     p_fp32.add_(exp_avg, alpha=-step_size * group['lr'])
 
                 if p.dtype in (torch.float16, torch.bfloat16):
diff --git a/pytorch_optimizer/ralamb.py b/pytorch_optimizer/ralamb.py
@@ -94,7 +94,7 @@ def get_gradient_norm(self) -> float:
                 if p.grad is None:
                     continue
 
-                norm_sq += torch.linalg.norm(p.grad).item() ** 2
+                norm_sq += torch.linalg.norm(p.grad).cpu().numpy() ** 2
 
         norm = math.sqrt(norm_sq)
 
@@ -147,7 +147,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
 
                 state['step'] += 1
-                buffered = group['buffer'][int(state['step'] % 10)]
+                buffered = group['buffer'][state['step'] % 10]
 
                 bias_correction1 = 1.0 - beta1 ** state['step']
 
diff --git a/pytorch_optimizer/ranger.py b/pytorch_optimizer/ranger.py
@@ -1,5 +1,4 @@
 import math
-from typing import Dict
 
 import torch
 from torch.optim.optimizer import Optimizer
diff --git a/pytorch_optimizer/ranger21.py b/pytorch_optimizer/ranger21.py
@@ -148,11 +148,11 @@ def reset(self):
 
     @staticmethod
     def build_warm_up_iterations(total_iterations: int, beta2: float, warm_up_pct: float = 0.22) -> int:
-        beta_warm_up_iterations: int = math.ceil(2.0 / (1.0 - beta2))  # default un-tuned linear warmup
-        beta_pct: float = beta_warm_up_iterations / total_iterations
+        warm_up_iterations: int = math.ceil(2.0 / (1.0 - beta2))  # default un-tuned linear warmup
+        beta_pct: float = warm_up_iterations / total_iterations
         if beta_pct > 0.45:
             return int(warm_up_pct * total_iterations)
-        return beta_warm_up_iterations
+        return warm_up_iterations
 
     @staticmethod
     def build_warm_down_iterations(total_iterations: int, warm_down_pct: float = 0.72) -> int:
@@ -187,13 +187,6 @@ def get_warm_down(self, lr: float, iteration: int) -> float:
 
         return new_lr
 
-    @staticmethod
-    def get_state_values(group, state: STATE):
-        beta1, beta2 = group['betas']
-        mean_avg = state['mean_avg']
-        variance_avg = state['variance_avg']
-        return beta1, beta2, mean_avg, variance_avg
-
     @torch.no_grad()
     def step(self, closure: CLOSURE = None) -> LOSS:
         loss: LOSS = None
diff --git a/pytorch_optimizer/types.py b/pytorch_optimizer/types.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Union
 
 import torch
 
diff --git a/pytorch_optimizer/utils.py b/pytorch_optimizer/utils.py
@@ -43,12 +43,12 @@ def flatten_grad(grads: List[torch.Tensor]) -> torch.Tensor:
 
 def un_flatten_grad(grads: torch.Tensor, shapes: List[int]) -> List[torch.Tensor]:
     idx: int = 0
-    un_flatten_grad: List[torch.Tensor] = []
+    un_flatten_grads: List[torch.Tensor] = []
     for shape in shapes:
         length = np.prod(shape)
-        un_flatten_grad.append(grads[idx : idx + length].view(shape).clone())
+        un_flatten_grads.append(grads[idx : idx + length].view(shape).clone())
         idx += length
-    return un_flatten_grad
+    return un_flatten_grads
 
 
 def channel_view(x: torch.Tensor) -> torch.Tensor:
diff --git a/pytorch_optimizer/version.py b/pytorch_optimizer/version.py
@@ -1 +1 @@
-__VERSION__ = '0.4.1'
+__VERSION__ = '0.4.2'
diff --git a/tests/test_optimizer_parameters.py b/tests/test_optimizer_parameters.py
@@ -3,7 +3,7 @@
 import pytest
 from torch import nn
 
-from pytorch_optimizer import SAM, AdamP, Lookahead, PCGrad, SafeFP16Optimizer, load_optimizers
+from pytorch_optimizer import SAM, AdamP, Lookahead, PCGrad, Ranger21, SafeFP16Optimizer, load_optimizers
 from tests.utils import Example
 
 OPTIMIZER_NAMES: List[str] = [
@@ -173,3 +173,9 @@ def test_safe_fp16_methods():
         optimizer.set_lr(lr=5e-1)
 
     assert optimizer.loss_scale == 2.0 ** (15 - 1)
+
+
+def test_ranger21_methods():
+    assert Ranger21.build_warm_up_iterations(1000, 0.999) == 220
+    assert Ranger21.build_warm_up_iterations(4500, 0.999) == 2000
+    assert Ranger21.build_warm_down_iterations(1000) == 280
diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py
@@ -41,6 +41,8 @@
     (AdaBelief, {'lr': 5e-1, 'weight_decay': 1e-3, 'fixed_decay': True}, 200),
     (AdaBelief, {'lr': 5e-1, 'weight_decay': 1e-3, 'rectify': False}, 200),
     (AdaBound, {'lr': 5e-1, 'gamma': 0.1, 'weight_decay': 1e-3}, 200),
+    (AdaBound, {'lr': 5e-1, 'gamma': 0.1, 'weight_decay': 1e-3, 'fixed_decay': True}, 200),
+    (AdaBound, {'lr': 5e-1, 'gamma': 0.1, 'weight_decay': 1e-3, 'weight_decouple': False}, 200),
     (AdaBound, {'lr': 5e-1, 'gamma': 0.1, 'weight_decay': 1e-3, 'amsbound': True}, 200),
     (AdamP, {'lr': 5e-1, 'weight_decay': 1e-3}, 200),
     (AdamP, {'lr': 5e-1, 'weight_decay': 1e-3, 'use_gc': True}, 200),
@@ -61,6 +63,7 @@
     (RAdam, {'lr': 1e-1, 'weight_decay': 1e-3}, 200),
     (RAdam, {'lr': 1e-1, 'weight_decay': 1e-3, 'degenerated_to_sgd': True}, 200),
     (SGDP, {'lr': 2e-1, 'weight_decay': 1e-3}, 500),
+    (SGDP, {'lr': 2e-1, 'weight_decay': 1e-3, 'nesterov': True}, 500),
     (Ranger, {'lr': 5e-1, 'weight_decay': 1e-3}, 200),
     (Ranger21, {'lr': 5e-1, 'weight_decay': 1e-3, 'num_iterations': 500}, 500),
 ]
@@ -248,8 +251,9 @@ def test_adamd_optimizers(optimizer_adamd_config):
     assert tensor_to_numpy(init_loss) > 2.0 * tensor_to_numpy(loss)
 
 
+@pytest.mark.parametrize('reduction', ('mean', 'sum'))
 @pytest.mark.parametrize('optimizer_pc_grad_config', OPTIMIZERS, ids=ids)
-def test_pc_grad_optimizers(optimizer_pc_grad_config):
+def test_pc_grad_optimizers(reduction, optimizer_pc_grad_config):
     torch.manual_seed(42)
 
     x_data, y_data = make_dataset()
@@ -259,7 +263,7 @@ def test_pc_grad_optimizers(optimizer_pc_grad_config):
     loss_fn_2: nn.Module = nn.L1Loss()
 
     optimizer_class, config, iterations = optimizer_pc_grad_config
-    optimizer = PCGrad(optimizer_class(model.parameters(), **config))
+    optimizer = PCGrad(optimizer_class(model.parameters(), **config), reduction=reduction)
 
     if optimizer_class.__name__ == 'RaLamb' and 'pre_norm' in config:
         return True

Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@ def get_configuration() -> Namespace:`
`14`	`14`	`parser.add_argument(`
`15`	`15`	`'-t',`
`16`	`16`	`'--threshold',`
`17`		`- default=9.9,`
	`17`	`+ default=9.95,`
`18`	`18`	`type=float,`
`19`	`19`	`)`
`20`	`20`
Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,4 @@`
`1`	`1`	`import math`
`2`		`-from typing import Dict`
`3`	`2`
`4`	`3`	`import torch`
`5`	`4`	`from torch.optim.optimizer import Optimizer`