update: fp16, SAM test cases

Author: kozistr

tests/test_optimizers.py

@@ -8,6 +8,7 @@
 
 from pytorch_optimizer import (
     MADGRAD,
+    SAM,
     SGDP,
     AdaBelief,
     AdaBound,
@@ -19,6 +20,7 @@
     RAdam,
     Ranger,
     Ranger21,
+    SafeFP16Optimizer,
 )
 
 __REFERENCE__ = 'https://github.com/jettify/pytorch-optimizer/blob/master/tests/test_optimizer_with_nn.py'
@@ -66,7 +68,7 @@ def build_lookahead(*parameters, **kwargs):
     return Lookahead(AdamP(*parameters, **kwargs))
 
 
-OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
+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),
     (AdaBound, {'lr': 1e-2, 'gamma': 0.1, 'weight_decay': 1e-3}, 200),
@@ -81,17 +83,31 @@ def build_lookahead(*parameters, **kwargs):
     (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),
+    (AdaBelief, {'lr': 5e-1, 'weight_decay': 1e-3}, 200),
+    (AdaBound, {'lr': 5e-1, 'gamma': 0.1, 'weight_decay': 1e-3}, 200),
+    (AdamP, {'lr': 5e-1, 'weight_decay': 1e-3}, 500),
+    (DiffGrad, {'lr': 15 - 1, 'weight_decay': 1e-3}, 500),
+    (DiffRGrad, {'lr': 1e-1, 'weight_decay': 1e-3}, 200),
+    (Lamb, {'lr': 1e-1, 'weight_decay': 1e-3}, 200),
+    (RAdam, {'lr': 1e-1, 'weight_decay': 1e-3}, 200),
+    (SGDP, {'lr': 5e-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),
+]
+
 
-@pytest.mark.parametrize('optimizer_config', OPTIMIZERS, ids=ids)
-def test_optimizers(optimizer_config):
+@pytest.mark.parametrize('optimizer_fp32_config', FP32_OPTIMIZERS, ids=ids)
+def test_f32_optimizers(optimizer_fp32_config):
     torch.manual_seed(42)
 
     x_data, y_data = make_dataset()
 
     model: nn.Module = LogisticRegression()
     loss_fn: nn.Module = nn.BCEWithLogitsLoss()
 
-    optimizer_class, config, iterations = optimizer_config
+    optimizer_class, config, iterations = optimizer_fp32_config
     optimizer = optimizer_class(model.parameters(), **config)
 
     loss: float = np.inf
@@ -110,3 +126,58 @@ def test_optimizers(optimizer_config):
         optimizer.step()
 
     assert init_loss > 2.0 * loss
+
+
+@pytest.mark.parametrize('optimizer_fp16_config', FP16_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()
+
+    optimizer_class, config, iterations = optimizer_fp16_config
+    optimizer = SafeFP16Optimizer(optimizer_class(model.parameters(), **config))
+
+    loss: float = np.inf
+    init_loss: float = np.inf
+    for _ in range(1000):
+        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 init_loss - 0.01 > loss
+
+
+@pytest.mark.parametrize('optimizer_config', FP32_OPTIMIZERS, ids=ids)
+def test_sam_optimizers(optimizer_config):
+    torch.manual_seed(42)
+
+    x_data, y_data = make_dataset()
+
+    model: nn.Module = LogisticRegression()
+    loss_fn: nn.Module = nn.BCEWithLogitsLoss()
+
+    optimizer_class, config, iterations = optimizer_config
+    optimizer = SAM(model.parameters(), optimizer_class, **config)
+
+    loss: float = np.inf
+    init_loss: float = np.inf
+    for _ in range(iterations):
+        loss = loss_fn(y_data, model(x_data))
+        loss.backward()
+        optimizer.first_step(zero_grad=True)
+
+        loss_fn(y_data, model(x_data)).backward()
+        optimizer.second_step(zero_grad=True)
+
+    assert init_loss > 2.0 * loss