Merge pull request #202 from kozistr/fix/lookahead-optimizer

[Refactor] Lookahead optimizer

Author: kozistr

CONTRIBUTING.md

@@ -7,6 +7,8 @@ Contributions to `pytorch-optimizer` for code, documentation, and tests are alwa
 Currently, `black` and `ruff` are used to format & lint the code. Here are the [lint options](https://github.com/kozistr/pytorch_optimizer/blob/main/pyproject.toml#L69)
 Or you just simply run `make format` and `make check` on the project root.
 
+You can create the environment with `make init` or just install the pip packages to your computer.
+
 A few differences from the default `black` (or another style guide) are
 
 1. line-length is **119** characters.

Makefile

@@ -1,8 +1,8 @@
 .PHONY: init format test check requirements docs
 
 init:
-	python -m pip install -q -U poetry
-	python -m poetry install
+	python -m pip install -q -U poetry isort black ruff pytest pytest-cov
+	python -m poetry install --dev
 
 format:
 	isort --profile black -l 119 pytorch_optimizer tests hubconf.py

docs/changelogs/v2.11.2.md

@@ -0,0 +1,14 @@
+## Change Log
+
+### Fix
+
+* Fix Lookahead optimizer (#200, #201, #202)
+  * When using PyTorch Lightning which expects your optimiser to be a subclass of `Optimizer`.
+
+### Contributions
+
+thanks to @georg-wolflein
+
+### Diff
+
+[2.11.1...2.11.2](https://github.com/kozistr/pytorch_optimizer/compare/v2.11.1...v2.11.2)

poetry.lock

@@ -60,7 +60,7 @@ black = [
     { version = "==23.3.0", python = ">=3.7,<3.8" },
     { version = "^23.7.0", python = ">=3.8"}
 ]
-ruff = "^0.0.284"
+ruff = "^0.0.286"
 pytest = "^7.4.0"
 pytest-cov = "^4.1.0"
 

pyproject.toml

@@ -1,14 +1,14 @@
-from collections import defaultdict, OrderedDict
-from typing import Dict, Callable
+from collections import defaultdict
+from typing import Callable, Dict
 
 import torch
 from torch.optim import Optimizer
 
 from pytorch_optimizer.base.optimizer import BaseOptimizer
-from pytorch_optimizer.base.types import CLOSURE, LOSS, OPTIMIZER, STATE, DEFAULTS
+from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, OPTIMIZER, STATE
 
 
-class Lookahead(BaseOptimizer, Optimizer):
+class Lookahead(Optimizer, BaseOptimizer):
     r"""k steps forward, 1 step back.
 
     :param optimizer: OPTIMIZER. base optimizer.
@@ -28,6 +28,9 @@ def __init__(
         self.validate_range(alpha, 'alpha', 0.0, 1.0)
         self.validate_options(pullback_momentum, 'pullback_momentum', ['none', 'reset', 'pullback'])
 
+        self._optimizer_step_pre_hooks: Dict[int, Callable] = {}
+        self._optimizer_step_post_hooks: Dict[int, Callable] = {}
+
         self.alpha = alpha
         self.k = k
         self.pullback_momentum = pullback_momentum
@@ -47,15 +50,13 @@ def __init__(
                 state['slow_params'].copy_(p)
                 if self.pullback_momentum == 'pullback':
                     state['slow_momentum'] = torch.zeros_like(p)
-        
-        # Instead of calling super().__init__, we set the attributes ourselves
-        self._optimizer_step_pre_hooks: Dict[int, Callable] = OrderedDict()
-        self._optimizer_step_post_hooks: Dict[int, Callable] = OrderedDict()
+
         self.defaults: DEFAULTS = {
+            'lookahead_alpha': alpha,
+            'lookahead_k': k,
+            'lookahead_pullback_momentum': pullback_momentum,
             **optimizer.defaults,
-            **dict(lookahead_alpha=alpha, lookahead_k=k, lookahead_pullback_momentum=pullback_momentum),
         }
-        
 
     def __getstate__(self):
         return {

pytorch_optimizer/optimizer/lookahead.py

@@ -2,11 +2,11 @@
 
 black==23.3.0 ; python_full_version >= "3.7.2" and python_version < "3.8"
 black==23.7.0 ; python_version >= "3.8" and python_full_version < "4.0.0"
-click==8.1.6 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
+click==8.1.7 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 colorama==0.4.6 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0" and (sys_platform == "win32" or platform_system == "Windows")
 coverage[toml]==7.2.7 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
-exceptiongroup==1.1.2 ; python_full_version >= "3.7.2" and python_version < "3.11"
-filelock==3.12.2 ; python_version >= "3.8" and python_full_version < "4.0.0"
+exceptiongroup==1.1.3 ; python_full_version >= "3.7.2" and python_version < "3.11"
+filelock==3.12.3 ; python_version >= "3.8" and python_full_version < "4.0.0"
 importlib-metadata==6.7.0 ; python_full_version >= "3.7.2" and python_version < "3.8"
 iniconfig==2.0.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 isort==5.11.5 ; python_full_version >= "3.7.2" and python_version < "3.8"
@@ -24,7 +24,7 @@ platformdirs==3.10.0 ; python_full_version >= "3.7.2" and python_full_version <
 pluggy==1.2.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 pytest-cov==4.1.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 pytest==7.4.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
-ruff==0.0.284 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
+ruff==0.0.286 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 sympy==1.12 ; python_version >= "3.8" and python_full_version < "4.0.0"
 tomli==2.0.1 ; python_full_version >= "3.7.2" and python_version < "3.11"
 torch==1.13.1+cpu ; python_full_version >= "3.7.2" and python_version < "3.8"

requirements-dev.txt

@@ -1,6 +1,6 @@
 --extra-index-url https://download.pytorch.org/whl/cpu
 
-filelock==3.12.2 ; python_version >= "3.8" and python_full_version < "4.0.0"
+filelock==3.12.3 ; python_version >= "3.8" and python_full_version < "4.0.0"
 jinja2==3.1.2 ; python_version >= "3.8" and python_full_version < "4.0.0"
 markupsafe==2.1.3 ; python_version >= "3.8" and python_full_version < "4.0.0"
 mpmath==1.3.0 ; python_version >= "3.8" and python_full_version < "4.0.0"