Merge pull request #45 from kozistr/test/cases

[Test] Add more test cases

Author: kozistr

pytorch_optimizer/adabelief.py

@@ -38,14 +38,14 @@ def __init__(
         adamd_debias_term: bool = False,
         eps: float = 1e-16,
     ):
-        """
+        """AdaBelief
         :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
         :param weight_decay: float. weight decay (L2 penalty)
         :param n_sma_threshold: (recommended is 5)
         :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW
-        :param fixed_decay: bool.
+        :param fixed_decay: bool. fix weight decay
         :param rectify: bool. perform the rectified update similar to RAdam
         :param degenerated_to_sgd: bool. perform SGD update when variance of gradient is high
         :param amsgrad: bool. whether to use the AMSBound variant
@@ -63,6 +63,8 @@ def __init__(
         self.adamd_debias_term = adamd_debias_term
         self.eps = eps
 
+        self.check_valid_parameters()
+
         buffer: BUFFER = [[None, None, None] for _ in range(10)]
 
         if is_valid_parameters(params):
@@ -81,6 +83,18 @@ 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 __setstate__(self, state: STATE):
         super().__setstate__(state)
         for group in self.param_groups:

pytorch_optimizer/adabound.py

@@ -37,15 +37,15 @@ def __init__(
         adamd_debias_term: bool = False,
         eps: float = 1e-8,
     ):
-        """
+        """AdaBound
         :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups
         :param lr: float. learning rate
         :param final_lr: float. final learning rate
         :param betas: BETAS. coefficients used for computing running averages of gradient and the squared hessian trace
         :param gamma: float. convergence speed of the bound functions
         :param weight_decay: float. weight decay (L2 penalty)
         :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW
-        :param fixed_decay: bool.
+        :param fixed_decay: bool. fix weight decay
         :param amsbound: bool. whether to use the AMSBound variant
         :param adamd_debias_term: bool. Only correct the denominator to avoid inflating step sizes early in training
         :param eps: float. term added to the denominator to improve numerical stability
@@ -57,6 +57,8 @@ def __init__(
         self.fixed_decay = fixed_decay
         self.eps = eps
 
+        self.check_valid_parameters()
+
         defaults: DEFAULTS = dict(
             lr=lr,
             betas=betas,

pytorch_optimizer/pcgrad.py

@@ -1,6 +1,6 @@
 import random
 from copy import deepcopy
-from typing import Iterable, List
+from typing import Iterable, List, Tuple
 
 import numpy as np
 import torch
@@ -35,12 +35,12 @@ def check_valid_parameters(self):
             raise ValueError(f'invalid reduction : {self.reduction}')
 
     @staticmethod
-    def flatten_grad(grads) -> torch.Tensor:
+    def flatten_grad(grads: List[torch.Tensor]) -> torch.Tensor:
         return torch.cat([g.flatten() for g in grads])
 
     @staticmethod
     def un_flatten_grad(grads, shapes) -> List[torch.Tensor]:
-        un_flatten_grad = []
+        un_flatten_grad: List[torch.Tensor] = []
         idx: int = 0
         for shape in shapes:
             length = np.prod(shape)
@@ -54,39 +54,40 @@ def zero_grad(self):
     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']:
                 p.grad = grads[idx]
                 idx += 1
 
-    def retrieve_grad(self):
+    def retrieve_grad(self) -> Tuple[List[torch.Tensor], List[int], List[torch.Tensor]]:
         """get the gradient of the parameters of the network with specific objective"""
         grad, shape, has_grad = [], [], []
         for group in self.optimizer.param_groups:
             for p in group['params']:
                 if p.grad is None:
                     shape.append(p.shape)
-                    grad.append(torch.zeros_like(p).to(p.device))
-                    has_grad.append(torch.zeros_like(p).to(p.device))
+                    grad.append(torch.zeros_like(p, device=p.device))
+                    has_grad.append(torch.zeros_like(p, device=p.device))
                     continue
 
                 shape.append(p.grad.shape)
                 grad.append(p.grad.clone())
-                has_grad.append(torch.ones_like(p).to(p.device))
+                has_grad.append(torch.ones_like(p, device=p.device))
 
         return grad, shape, has_grad
 
-    def pack_grad(self, objectives: Iterable[nn.Module]):
+    def pack_grad(
+        self, objectives: Iterable[nn.Module]
+    ) -> 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[float]. a list of objectives
+        :param objectives: Iterable[nn.Module]. a list of objectives
         :return:
         """
         grads, shapes, has_grads = [], [], []
         for objective in objectives:
-            self.zero_grad()
-
+            self.optimizer.zero_grad(set_to_none=True)
             objective.backward(retain_graph=True)
 
             grad, shape, has_grad = self.retrieve_grad()
@@ -98,7 +99,7 @@ def pack_grad(self, objectives: Iterable[nn.Module]):
         return grads, shapes, has_grads
 
     def project_conflicting(self, grads, has_grads) -> 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:
@@ -114,12 +115,10 @@ def project_conflicting(self, grads, has_grads) -> torch.Tensor:
                     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[shared] = torch.stack([g[shared] for g in pc_grad])
-
         if self.reduction == 'mean':
-            merged_grad = merged_grad.mean(dim=0)
-        else:  # self.reduction == 'sum'
-            merged_grad = merged_grad.sum(dim=0)
+            merged_grad[shared] = torch.stack([g[shared] for g in pc_grad]).mean(dim=0)
+        else:
+            merged_grad[shared] = torch.stack([g[shared] for g in pc_grad]).sum(dim=0)
 
         merged_grad[~shared] = torch.stack([g[~shared] for g in pc_grad]).sum(dim=0)
 

pytorch_optimizer/radam.py

@@ -35,13 +35,13 @@ def __init__(
         adamd_debias_term: bool = False,
         eps: float = 1e-8,
     ):
-        """
+        """RAdam
         :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups
-        :param lr: float. learning rate.
+        :param lr: float. learning rate
         :param betas: BETAS. coefficients used for computing running averages of gradient and the squared hessian trace
         :param weight_decay: float. weight decay (L2 penalty)
         :param n_sma_threshold: int. (recommended is 5)
-        :param degenerated_to_sgd: float.
+        :param degenerated_to_sgd: float. degenerated to SGD
         :param adamd_debias_term: bool. Only correct the denominator to avoid inflating step sizes early in training
         :param eps: float. term added to the denominator to improve numerical stability
         """

pytorch_optimizer/ranger21.py

@@ -92,6 +92,8 @@ def __init__(
         self.norm_loss_factor = norm_loss_factor
         self.eps = eps
 
+        self.check_valid_parameters()
+
         # lookahead
         self.lookahead_step: int = 0
 
@@ -124,6 +126,18 @@ def __init__(
         self.start_warm_down: int = num_iterations - self.num_warm_down_iterations
         self.warm_down_lr_delta: float = self.starting_lr - self.min_lr
 
+    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 __setstate__(self, state: STATE):
         super().__setstate__(state)
 

tests/test_load_optimizers.py

@@ -36,8 +36,5 @@ def test_load_optimizers_valid(valid_optimizer_names):
 
 @pytest.mark.parametrize('invalid_optimizer_names', INVALID_OPTIMIZER_NAMES)
 def test_load_optimizers_invalid(invalid_optimizer_names):
-    try:
+    with pytest.raises(NotImplementedError):
         load_optimizers(invalid_optimizer_names)
-    except NotImplementedError:
-        return True
-    return False

tests/test_optimizer_parameters.py

@@ -0,0 +1,65 @@
+from typing import List
+
+import pytest
+
+from pytorch_optimizer import load_optimizers
+
+OPTIMIZER_NAMES: List[str] = [
+    'adamp',
+    'sgdp',
+    'madgrad',
+    'ranger',
+    'ranger21',
+    'radam',
+    'adabound',
+    'adahessian',
+    'adabelief',
+    'diffgrad',
+    'diffrgrad',
+    'lamb',
+]
+
+BETA_OPTIMIZER_NAMES: List[str] = [
+    'adabelief',
+    'adabound',
+    'adahessian',
+    'adamp',
+    'diffgrad',
+    'diffrgrad',
+    'lamb',
+    'radam',
+    'ranger',
+    'ranger21',
+]
+
+
+@pytest.mark.parametrize('optimizer_names', OPTIMIZER_NAMES)
+def test_learning_rate(optimizer_names):
+    with pytest.raises(ValueError):
+        optimizer = load_optimizers(optimizer_names)
+        optimizer(None, lr=-1e-2)
+
+
+@pytest.mark.parametrize('optimizer_names', OPTIMIZER_NAMES)
+def test_epsilon(optimizer_names):
+    with pytest.raises(ValueError):
+        optimizer = load_optimizers(optimizer_names)
+        optimizer(None, eps=-1e-6)
+
+
+@pytest.mark.parametrize('optimizer_names', OPTIMIZER_NAMES)
+def test_weight_decay(optimizer_names):
+    with pytest.raises(ValueError):
+        optimizer = load_optimizers(optimizer_names)
+        optimizer(None, weight_decay=-1e-3)
+
+
+@pytest.mark.parametrize('optimizer_names', BETA_OPTIMIZER_NAMES)
+def test_betas(optimizer_names):
+    with pytest.raises(ValueError):
+        optimizer = load_optimizers(optimizer_names)
+        optimizer(None, betas=(-0.1, 0.1))
+
+    with pytest.raises(ValueError):
+        optimizer = load_optimizers(optimizer_names)
+        optimizer(None, betas=(0.1, -0.1))

tests/test_optimizers.py

@@ -17,6 +17,7 @@
     DiffRGrad,
     Lamb,
     Lookahead,
+    PCGrad,
     RAdam,
     Ranger,
     Ranger21,
@@ -39,6 +40,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
+class MultiHeadLogisticRegression(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(2, 2)
+        self.head1 = nn.Linear(2, 1)
+        self.head2 = nn.Linear(2, 1)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        x = self.fc1(x)
+        x = F.relu(x)
+        return self.head1(x), self.head2(x)
+
+
 def make_dataset(num_samples: int = 100, dims: int = 2, seed: int = 42) -> Tuple[torch.Tensor, torch.Tensor]:
     rng = np.random.RandomState(seed)
 
@@ -180,4 +194,37 @@ def test_sam_optimizers(optimizer_config):
         loss_fn(y_data, model(x_data)).backward()
         optimizer.second_step(zero_grad=True)
 
+        if init_loss == np.inf:
+            init_loss = loss
+
+    assert init_loss > 2.0 * loss
+
+
+@pytest.mark.parametrize('optimizer_config', FP32_OPTIMIZERS, ids=ids)
+def test_pc_grad_optimizers(optimizer_config):
+    torch.manual_seed(42)
+
+    x_data, y_data = make_dataset()
+
+    model: nn.Module = MultiHeadLogisticRegression()
+    loss_fn_1: nn.Module = nn.BCEWithLogitsLoss()
+    loss_fn_2: nn.Module = nn.L1Loss()
+
+    optimizer_class, config, iterations = optimizer_config
+    optimizer = PCGrad(optimizer_class(model.parameters(), **config))
+
+    loss: float = np.inf
+    init_loss: float = np.inf
+    for _ in range(iterations):
+        optimizer.zero_grad()
+        y_pred_1, y_pred_2 = model(x_data)
+        loss1, loss2 = loss_fn_1(y_pred_1, y_data), loss_fn_2(y_pred_2, y_data)
+
+        loss = (loss1 + loss2) / 2.0
+        if init_loss == np.inf:
+            init_loss = loss
+
+        optimizer.pc_backward([loss1, loss2])
+        optimizer.step()
+
     assert init_loss > 2.0 * loss