Add more lookahead optimizer tests (#53)

Author: jettify

tests/test_basic.py

@@ -38,6 +38,11 @@ def ids(v):
     return n
 
 
+def build_lookahead(*a, **kw):
+    base = optim.Yogi(*a, **kw)
+    return optim.Lookahead(base)
+
+
 optimizers = [
     (
         optim.NovoGrad,
@@ -51,6 +56,7 @@ def ids(v):
     (optim.AdaBound, {'lr': 1.0}, 800),
     (optim.Yogi, {'lr': 1.0}, 500),
     (optim.AccSGD, {'lr': 0.015}, 800),
+    (build_lookahead, {'lr': 1.0}, 500),
 ]
 
 
@@ -69,3 +75,6 @@ def test_benchmark_function(case, optimizer_config):
         f.backward(retain_graph=True)
         optimizer.step()
     assert torch.allclose(x, x_min, atol=0.001)
+
+    name = optimizer.__class__.__name__
+    assert name in optimizer.__repr__()

tests/test_optimizer_with_nn.py

@@ -45,6 +45,11 @@ def ids(v):
     return f'{v[0].__name__} {v[1:]}'
 
 
+def build_lookahead(*a, **kw):
+    base = optim.Yogi(*a, **kw)
+    return optim.Lookahead(base)
+
+
 optimizers = [
     (optim.NovoGrad, {'lr': 0.01, 'weight_decay': 1e-3}, 200),
     (optim.Lamb, {'lr': 0.01, 'weight_decay': 1e-3}, 200),
@@ -55,6 +60,7 @@ def ids(v):
     (optim.Yogi, {'lr': 0.1, 'weight_decay': 1e-3}, 200),
     (optim.RAdam, {'lr': 1.0, 'weight_decay': 1e-3}, 200),
     (optim.AccSGD, {'lr': 1.0, 'weight_decay': 1e-3}, 200),
+    (build_lookahead, {'lr': 0.1, 'weight_decay': 1e-3}, 200),
 ]
 
 

torch_optimizer/lookahead.py

@@ -24,7 +24,7 @@ class Lookahead(Optimizer):
     Example:
         >>> import torch_optimizer as optim
         >>> yogi = optim.Yogi(model.parameters(), lr=0.1)
-        >>> optimizer = optim.Lookahead(yogi, k=5)
+        >>> optimizer = optim.Lookahead(yogi, k=5, alpha=0.5)
         >>> optimizer.zero_grad()
         >>> loss_fn(model(input), target).backward()
         >>> optimizer.step()
@@ -116,3 +116,14 @@ def load_state_dict(self, state_dict: State) -> None:
     def zero_grad(self) -> None:
         r"""Clears the gradients of all optimized :class:`torch.Tensor` s."""
         self.optimizer.zero_grad()
+
+    def __repr__(self) -> str:
+        base_str = self.optimizer.__repr__()
+        format_string = self.__class__.__name__ + ' ('
+        format_string += '\n'
+        format_string += f'k: {self.k}\n'
+        format_string += f'alpha: {self.alpha}\n'
+        format_string += base_str
+        format_string += '\n'
+        format_string += ')'
+        return format_string

torch_optimizer/novograd.py

@@ -1,5 +1,5 @@
 import torch
-from torch.optim import Optimizer
+from torch.optim.optimizer import Optimizer
 
 from .types import Betas2, OptFloat, OptLossClosure, Params
 

torch_optimizer/sgdw.py

@@ -35,9 +35,9 @@ def __init__(
         self,
         params: Params,
         lr: float = 1e-3,
-        momentum: float = 0,
-        dampening: float = 0,
-        weight_decay: float = 1e-2,
+        momentum: float = 0.0,
+        dampening: float = 0.0,
+        weight_decay: float = 0.0,
         nesterov: bool = False,
     ) -> None:
         if not 0.0 <= lr: