Merge pull request #48 from kozistr/refactor/test

[Test] Refactor test modules

Author: kozistr

pytorch_optimizer/lamb.py

@@ -90,7 +90,7 @@ def step(self, closure: CLOSURE = None) -> float:
 
         grad_norm: float = 1.0
         if self.pre_norm:
-            grad_norm = self.get_gradient_norm()
+            grad_norm = self.get_gradient_norm() + self.eps
 
         for group in self.param_groups:
             for p in group['params']:

pytorch_optimizer/lookahead.py

@@ -82,7 +82,7 @@ def update(self, group: Dict):
             if self.pullback_momentum == 'pullback':
                 internal_momentum = self.optimizer.state[fast]['momentum_buffer']
                 self.optimizer.state[fast]['momentum_buffer'] = internal_momentum.mul_(self.alpha).add_(
-                    1.0 - self.alpha, param_state['slow_mom']
+                    param_state['slow_mom'], alpha=1.0 - self.alpha
                 )
                 param_state['slow_mom'] = self.optimizer.state[fast]['momentum_buffer']
             elif self.pullback_momentum == 'reset':

pytorch_optimizer/madgrad.py

@@ -62,14 +62,6 @@ def check_valid_parameters(self):
         if self.eps < 0.0:
             raise ValueError(f'Invalid eps : {self.eps}')
 
-    @property
-    def supports_memory_efficient_fp16(self) -> bool:
-        return False
-
-    @property
-    def supports_flat_params(self) -> bool:
-        return True
-
     def step(self, closure: CLOSURE = None) -> LOSS:
         loss: LOSS = None
         if closure is not None:

pytorch_optimizer/ranger21.py

@@ -255,6 +255,9 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
         # Phase 2 - Apply weight decay and step
         for group in self.param_groups:
+            if group['params'][0].grad is None:
+                continue
+
             lr = group['lr']
             step = self.state[group['params'][0]]['step']
 

pytorch_optimizer/version.py

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

tests/test_optimizer_parameters.py

@@ -1,9 +1,21 @@
 from typing import List
 
 import pytest
+import torch
+from torch import nn
 
 from pytorch_optimizer import SAM, Lookahead, load_optimizers
 
+
+class Example(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(1, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.fc1(x)
+
+
 OPTIMIZER_NAMES: List[str] = [
     'adamp',
     'sgdp',
@@ -58,12 +70,12 @@ def test_weight_decay(optimizer_names):
 
 @pytest.mark.parametrize('optimizer_names', BETA_OPTIMIZER_NAMES)
 def test_betas(optimizer_names):
+    optimizer = load_optimizers(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))
 
 
@@ -73,11 +85,20 @@ def test_sam_parameters():
 
 
 def test_lookahead_parameters():
+    model: nn.Module = Example()
+    parameters = model.parameters()
+    optimizer = load_optimizers('adamp')(parameters)
+
+    pullback_momentum_list: List[str] = ['none', 'reset', 'pullback']
+    for pullback_momentum in pullback_momentum_list:
+        opt = Lookahead(optimizer, pullback_momentum=pullback_momentum)
+        opt.load_state_dict(opt.state_dict())
+
     with pytest.raises(ValueError):
-        Lookahead(load_optimizers('adamp'), k=0)
+        Lookahead(optimizer, k=0)
 
     with pytest.raises(ValueError):
-        Lookahead(load_optimizers('adamp'), alpha=0)
+        Lookahead(optimizer, alpha=0)
 
     with pytest.raises(ValueError):
-        Lookahead(load_optimizers('adamp'), pullback_momentum='asdf')
+        Lookahead(optimizer, pullback_momentum='invalid')

tests/test_optimizers.py

@@ -25,8 +25,6 @@
     SafeFP16Optimizer,
 )
 
-__REFERENCE__ = 'https://github.com/jettify/pytorch-optimizer/blob/master/tests/test_optimizer_with_nn.py'
-
 
 class LogisticRegression(nn.Module):
     def __init__(self):
@@ -83,41 +81,23 @@ def build_lookahead(*parameters, **kwargs):
     return Lookahead(AdamP(*parameters, **kwargs))
 
 
-FP32_OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
-    (build_lookahead, {'lr': 1e-2, 'weight_decay': 1e-3}, 200),
-    (AdaBelief, {'lr': 1e-2, 'weight_decay': 1e-3}, 200),
-    (AdaBelief, {'lr': 1e-2, 'weight_decay': 1e-3, 'amsgrad': True}, 200),
-    (AdaBelief, {'lr': 1e-2, 'weight_decay': 1e-3, 'weight_decouple': False}, 200),
-    (AdaBelief, {'lr': 1e-2, 'weight_decay': 1e-3, 'rectify': False}, 200),
-    (AdaBound, {'lr': 1e-2, 'gamma': 0.1, 'weight_decay': 1e-3}, 200),
-    (AdaBound, {'lr': 1e-2, 'gamma': 0.1, 'weight_decay': 1e-3, 'amsbound': True}, 200),
-    (AdamP, {'lr': 1e-3, 'weight_decay': 1e-3}, 800),
-    (DiffGrad, {'lr': 1e-2, 'weight_decay': 1e-3}, 200),
-    (DiffRGrad, {'lr': 1e-1, 'weight_decay': 1e-3}, 200),
-    (Lamb, {'lr': 1e-1, 'weight_decay': 1e-3}, 500),
-    (RaLamb, {'lr': 1e-3, 'weight_decay': 1e-3}, 500),
-    (MADGRAD, {'lr': 1e-2, 'weight_decay': 1e-3}, 200),
-    (RAdam, {'lr': 1e-1, 'weight_decay': 1e-3}, 200),
-    (SGDP, {'lr': 1e-1, 'weight_decay': 1e-3}, 200),
-    (Ranger, {'lr': 1e-1, 'weight_decay': 1e-3}, 200),
-    (Ranger21, {'lr': 5e-1, 'weight_decay': 1e-3, 'num_iterations': 500}, 500),
-]
-
-FP16_OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
-    (build_lookahead, {'lr': 5e-1, 'weight_decay': 1e-3}, 500),
+OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
+    (build_lookahead, {'lr': 5e-1, 'weight_decay': 1e-3}, 200),
     (AdaBelief, {'lr': 5e-1, 'weight_decay': 1e-3}, 200),
     (AdaBelief, {'lr': 5e-1, 'weight_decay': 1e-3, 'amsgrad': True}, 200),
     (AdaBelief, {'lr': 5e-1, 'weight_decay': 1e-3, 'weight_decouple': False}, 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': 1e-1, 'gamma': 0.1, 'weight_decay': 1e-3, 'amsbound': True}, 200),
-    (AdamP, {'lr': 5e-1, 'weight_decay': 1e-3}, 500),
-    (DiffGrad, {'lr': 5e-1, 'weight_decay': 1e-3}, 500),
-    (DiffRGrad, {'lr': 1e-1, 'weight_decay': 1e-3}, 200),
-    (Lamb, {'lr': 1e-1, 'weight_decay': 1e-3}, 200),
-    (RaLamb, {'lr': 1e-1, 'weight_decay': 1e-3}, 500),
+    (AdaBound, {'lr': 5e-1, 'gamma': 0.1, 'weight_decay': 1e-3, 'amsbound': True}, 200),
+    (AdamP, {'lr': 5e-1, 'weight_decay': 1e-3}, 200),
+    (DiffGrad, {'lr': 5e-1, 'weight_decay': 1e-3}, 200),
+    (DiffRGrad, {'lr': 5e-1, 'weight_decay': 1e-3}, 200),
+    (Lamb, {'lr': 1e-1, 'weight_decay': 1e-3}, 500),
+    (Lamb, {'lr': 1e-1, 'weight_decay': 1e-3, 'pre_norm': True, 'eps': 1e-8}, 500),
+    (RaLamb, {'lr': 1e-1, 'weight_decay': 1e-3}, 200),
+    (MADGRAD, {'lr': 1e-2, 'weight_decay': 1e-3}, 500),
     (RAdam, {'lr': 1e-1, 'weight_decay': 1e-3}, 200),
-    (SGDP, {'lr': 5e-1, 'weight_decay': 1e-3}, 500),
+    (SGDP, {'lr': 2e-1, 'weight_decay': 1e-3}, 500),
     (Ranger, {'lr': 5e-1, 'weight_decay': 1e-3}, 200),
     (Ranger21, {'lr': 5e-1, 'weight_decay': 1e-3, 'num_iterations': 500}, 500),
 ]
@@ -137,20 +117,33 @@ def build_lookahead(*parameters, **kwargs):
 ]
 
 
-@pytest.mark.parametrize('optimizer_fp32_config', FP32_OPTIMIZERS, ids=ids)
-def test_f32_optimizers(optimizer_fp32_config):
+def tensor_to_numpy(x: torch.Tensor) -> np.ndarray:
+    return x.detach().cpu().numpy()
+
+
+def build_environment(use_gpu: bool = False) -> Tuple[Tuple[torch.Tensor, torch.Tensor], nn.Module, nn.Module]:
     torch.manual_seed(42)
 
     x_data, y_data = make_dataset()
-
     model: nn.Module = LogisticRegression()
     loss_fn: nn.Module = nn.BCEWithLogitsLoss()
 
+    if use_gpu and torch.cuda.is_available():
+        x_data, y_data = x_data.cuda(), y_data.cuda()
+        model = model.cuda()
+        loss_fn = loss_fn.cuda()
+
+    return (x_data, y_data), model, loss_fn
+
+
+@pytest.mark.parametrize('optimizer_fp32_config', OPTIMIZERS, ids=ids)
+def test_f32_optimizers(optimizer_fp32_config):
+    (x_data, y_data), model, loss_fn = build_environment()
+
     optimizer_class, config, iterations = optimizer_fp32_config
     optimizer = optimizer_class(model.parameters(), **config)
 
-    loss: float = np.inf
-    init_loss: float = np.inf
+    init_loss, loss = np.inf, np.inf
     for _ in range(iterations):
         optimizer.zero_grad()
 
@@ -164,23 +157,20 @@ def test_f32_optimizers(optimizer_fp32_config):
 
         optimizer.step()
 
-    assert init_loss > 2.0 * loss
+    assert tensor_to_numpy(init_loss) > 2.0 * tensor_to_numpy(loss)
 
 
-@pytest.mark.parametrize('optimizer_fp16_config', FP16_OPTIMIZERS, ids=ids)
+@pytest.mark.parametrize('optimizer_fp16_config', OPTIMIZERS, ids=ids)
 def test_f16_optimizers(optimizer_fp16_config):
-    torch.manual_seed(42)
-
-    x_data, y_data = make_dataset()
-
-    model: nn.Module = LogisticRegression()
-    loss_fn: nn.Module = nn.BCEWithLogitsLoss()
+    (x_data, y_data), model, loss_fn = build_environment()
 
     optimizer_class, config, iterations = optimizer_fp16_config
+    if optimizer_class.__name__ == 'MADGRAD':
+        return True
+
     optimizer = SafeFP16Optimizer(optimizer_class(model.parameters(), **config))
 
-    loss: float = np.inf
-    init_loss: float = np.inf
+    init_loss, loss = np.inf, np.inf
     for _ in range(1000):
         optimizer.zero_grad()
 
@@ -194,24 +184,18 @@ def test_f16_optimizers(optimizer_fp16_config):
 
         optimizer.step()
 
-    assert init_loss - 0.01 > loss
+    assert tensor_to_numpy(init_loss) - 0.01 > tensor_to_numpy(loss)
 
 
 @pytest.mark.parametrize('adaptive', (False, True))
-@pytest.mark.parametrize('optimizer_sam_config', FP32_OPTIMIZERS, ids=ids)
+@pytest.mark.parametrize('optimizer_sam_config', OPTIMIZERS, ids=ids)
 def test_sam_optimizers(adaptive, optimizer_sam_config):
-    torch.manual_seed(42)
-
-    x_data, y_data = make_dataset()
-
-    model: nn.Module = LogisticRegression()
-    loss_fn: nn.Module = nn.BCEWithLogitsLoss()
+    (x_data, y_data), model, loss_fn = build_environment()
 
     optimizer_class, config, iterations = optimizer_sam_config
     optimizer = SAM(model.parameters(), optimizer_class, **config, adaptive=adaptive)
 
-    loss: float = np.inf
-    init_loss: float = np.inf
+    init_loss, loss = np.inf, np.inf
     for _ in range(iterations):
         loss = loss_fn(y_data, model(x_data))
         loss.backward()
@@ -223,10 +207,34 @@ def test_sam_optimizers(adaptive, optimizer_sam_config):
         if init_loss == np.inf:
             init_loss = loss
 
-    assert init_loss > 2.0 * loss
+    assert tensor_to_numpy(init_loss) > 2.0 * tensor_to_numpy(loss)
 
 
-@pytest.mark.parametrize('optimizer_pc_grad_config', FP32_OPTIMIZERS, ids=ids)
+@pytest.mark.parametrize('optimizer_adamd_config', ADAMD_SUPPORTED_OPTIMIZERS, ids=ids)
+def test_adamd_optimizers(optimizer_adamd_config):
+    (x_data, y_data), model, loss_fn = build_environment()
+
+    optimizer_class, config, iterations = optimizer_adamd_config
+    optimizer = optimizer_class(model.parameters(), **config)
+
+    init_loss, loss = np.inf, np.inf
+    for _ in range(iterations):
+        optimizer.zero_grad()
+
+        y_pred = model(x_data)
+        loss = loss_fn(y_pred, y_data)
+
+        if init_loss == np.inf:
+            init_loss = loss
+
+        loss.backward()
+
+        optimizer.step()
+
+    assert tensor_to_numpy(init_loss) > 2.0 * tensor_to_numpy(loss)
+
+
+@pytest.mark.parametrize('optimizer_pc_grad_config', OPTIMIZERS, ids=ids)
 def test_pc_grad_optimizers(optimizer_pc_grad_config):
     torch.manual_seed(42)
 
@@ -239,8 +247,7 @@ def test_pc_grad_optimizers(optimizer_pc_grad_config):
     optimizer_class, config, iterations = optimizer_pc_grad_config
     optimizer = PCGrad(optimizer_class(model.parameters(), **config))
 
-    loss: float = np.inf
-    init_loss: float = np.inf
+    init_loss, loss = np.inf, np.inf
     for _ in range(iterations):
         optimizer.zero_grad()
         y_pred_1, y_pred_2 = model(x_data)
@@ -253,23 +260,20 @@ def test_pc_grad_optimizers(optimizer_pc_grad_config):
         optimizer.pc_backward([loss1, loss2])
         optimizer.step()
 
-    assert init_loss > 2.0 * loss
-
+    assert tensor_to_numpy(init_loss) > 1.5 * tensor_to_numpy(loss)
 
-@pytest.mark.parametrize('optimizer_adamd_config', ADAMD_SUPPORTED_OPTIMIZERS, ids=ids)
-def test_adamd_optimizers(optimizer_adamd_config):
-    torch.manual_seed(42)
 
-    x_data, y_data = make_dataset()
+@pytest.mark.parametrize('optimizer_config', OPTIMIZERS, ids=ids)
+def test_no_gradients(optimizer_config):
+    (x_data, y_data), model, loss_fn = build_environment()
 
-    model: nn.Module = LogisticRegression()
-    loss_fn: nn.Module = nn.BCEWithLogitsLoss()
+    model.fc1.weight.requires_grad = False
+    model.fc1.bias.requires_grad = False
 
-    optimizer_class, config, iterations = optimizer_adamd_config
+    optimizer_class, config, iterations = optimizer_config
     optimizer = optimizer_class(model.parameters(), **config)
 
-    loss: float = np.inf
-    init_loss: float = np.inf
+    init_loss, loss = np.inf, np.inf
     for _ in range(iterations):
         optimizer.zero_grad()
 
@@ -283,4 +287,4 @@ def test_adamd_optimizers(optimizer_adamd_config):
 
         optimizer.step()
 
-    assert init_loss > 2.0 * loss
+    assert tensor_to_numpy(init_loss) >= tensor_to_numpy(loss)