Merge pull request #47 from kozistr/feature/madgrad

[Fix] sparse gradient for MADGRAD

Author: kozistr

pytorch_optimizer/adamp.py

@@ -38,9 +38,9 @@ def __init__(
         adamd_debias_term: bool = False,
         eps: float = 1e-8,
     ):
-        """
+        """AdamP optimizer
         :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 delta: float. threshold that determines whether a set of parameters is scale invariant or not

pytorch_optimizer/diffgrad.py

@@ -31,9 +31,9 @@ def __init__(
         weight_decay: float = 0.0,
         adamd_debias_term: bool = False,
     ):
-        """
+        """DiffGrad
         :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 eps: float. term added to the denominator to improve numerical stability
         :param weight_decay: float. weight decay (L2 penalty)

pytorch_optimizer/diffrgrad.py

@@ -38,11 +38,11 @@ def __init__(
     ):
         """Blend RAdam with DiffGrad
         :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: bool..
+        :param degenerated_to_sgd: bool. 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/lookahead.py

@@ -31,8 +31,8 @@ def __init__(
         alpha: float = 0.5,
         pullback_momentum: str = 'none',
     ):
-        """
-        :param optimizer: Optimizer.
+        """Lookahead
+        :param optimizer: Optimizer. base optimizer
         :param k: int. number of lookahead steps
         :param alpha: float. linear interpolation factor
         :param pullback_momentum: str. change to inner optimizer momentum on interpolation update

pytorch_optimizer/sgdp.py

@@ -36,9 +36,9 @@ def __init__(
         wd_ratio: float = 0.1,
         nesterov: bool = False,
     ):
-        """
+        """SGDP optimizer
         :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups
-        :param lr: float. learning rate.
+        :param lr: float. learning rate
         :param momentum: float. momentum factor
         :param dampening: float. dampening for momentum
         :param eps: float. term added to the denominator to improve numerical stability

pytorch_optimizer/version.py

@@ -1 +1 @@
-__VERSION__ = '0.3.3'
+__VERSION__ = '0.3.4'

tests/test_optimizer_parameters.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from pytorch_optimizer import SAM, load_optimizers
+from pytorch_optimizer import SAM, Lookahead, load_optimizers
 
 OPTIMIZER_NAMES: List[str] = [
     'adamp',
@@ -67,6 +67,17 @@ def test_betas(optimizer_names):
         optimizer(None, betas=(0.1, -0.1))
 
 
-def test_rho():
+def test_sam_parameters():
     with pytest.raises(ValueError):
         SAM(None, load_optimizers('adamp'), rho=-0.1)
+
+
+def test_lookahead_parameters():
+    with pytest.raises(ValueError):
+        Lookahead(load_optimizers('adamp'), k=0)
+
+    with pytest.raises(ValueError):
+        Lookahead(load_optimizers('adamp'), alpha=0)
+
+    with pytest.raises(ValueError):
+        Lookahead(load_optimizers('adamp'), pullback_momentum='asdf')

tests/test_optimizers.py

@@ -197,8 +197,9 @@ def test_f16_optimizers(optimizer_fp16_config):
     assert init_loss - 0.01 > loss
 
 
+@pytest.mark.parametrize('adaptive', (False, True))
 @pytest.mark.parametrize('optimizer_sam_config', FP32_OPTIMIZERS, ids=ids)
-def test_sam_optimizers(optimizer_sam_config):
+def test_sam_optimizers(adaptive, optimizer_sam_config):
     torch.manual_seed(42)
 
     x_data, y_data = make_dataset()
@@ -207,7 +208,7 @@ def test_sam_optimizers(optimizer_sam_config):
     loss_fn: nn.Module = nn.BCEWithLogitsLoss()
 
     optimizer_class, config, iterations = optimizer_sam_config
-    optimizer = SAM(model.parameters(), optimizer_class, **config)
+    optimizer = SAM(model.parameters(), optimizer_class, **config, adaptive=adaptive)
 
     loss: float = np.inf
     init_loss: float = np.inf

tests/test_sparse_gradient.py

@@ -31,10 +31,9 @@ def test_sparse_not_supported(no_sparse_optimizer):
     grad = torch.randn(1, 1).to_sparse(1)
     param.grad = grad
 
-    optimizer = load_optimizers(optimizer=no_sparse_optimizer)([param])
-    optimizer.zero_grad()
-
     with pytest.raises(RuntimeError):
+        optimizer = load_optimizers(optimizer=no_sparse_optimizer)([param])
+        optimizer.zero_grad()
         optimizer.step()
 
 
@@ -47,3 +46,8 @@ def test_sparse_supported(sparse_optimizer):
     optimizer = load_optimizers(optimizer=sparse_optimizer)([param], momentum=0.0)
     optimizer.zero_grad()
     optimizer.step()
+
+    with pytest.raises(RuntimeError):
+        optimizer = load_optimizers(optimizer=sparse_optimizer)([param], momentum=0.0, weight_decay=1e-3)
+        optimizer.zero_grad()
+        optimizer.step()