update: AdamMini optimizer

Author: kozistr

pytorch_optimizer/__init__.py

@@ -40,6 +40,7 @@
 from pytorch_optimizer.optimizer.adahessian import AdaHessian
 from pytorch_optimizer.optimizer.adai import Adai
 from pytorch_optimizer.optimizer.adalite import Adalite
+from pytorch_optimizer.optimizer.adam_mini import AdamMini
 from pytorch_optimizer.optimizer.adamax import AdaMax
 from pytorch_optimizer.optimizer.adamod import AdaMod
 from pytorch_optimizer.optimizer.adamp import AdamP
@@ -203,6 +204,7 @@
     GrokFastAdamW,
     Kate,
     StableAdamW,
+    AdamMini,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

pytorch_optimizer/base/optimizer.py

@@ -258,6 +258,11 @@ def validate_learning_rate(learning_rate: Optional[float]) -> None:
         if learning_rate is not None and learning_rate < 0.0:
             raise NegativeLRError(learning_rate)
 
+    @staticmethod
+    def validate_mod(x: int, y: int) -> None:
+        if x % y != 0:
+            raise ValueError(f'[-] {x} must be divisible by {y}')
+
     def validate_betas(self, betas: BETAS) -> None:
         if betas[0] is not None:
             self.validate_range(betas[0], 'beta1', 0.0, 1.0, range_type='[]')

tests/constants.py

@@ -134,6 +134,7 @@
     'fadam',
     'grokfastadamw',
     'stableadamw',
+    'adammini',
 ]
 
 VALID_LR_SCHEDULER_NAMES: List[str] = [

tests/test_load_modules.py

@@ -38,7 +38,7 @@ def test_load_lr_scheduler_invalid(invalid_lr_scheduler_names):
 
 
 def test_get_supported_optimizers():
-    assert len(get_supported_optimizers()) == 70
+    assert len(get_supported_optimizers()) == 71
 
 
 def test_get_supported_lr_schedulers():

tests/test_optimizers.py

@@ -650,3 +650,11 @@ def test_stableadamw_optimizer(environment):
     optimizer = load_optimizer('StableAdamW')(model.parameters())
     optimizer.reset()
     optimizer.step()
+
+
+def test_adam_mini_optimizer(environment):
+    _, model, _ = environment
+
+    optimizer = load_optimizer('AdamMini')(model)
+    optimizer.reset()
+    optimizer.step()