Merge pull request #216 from kozistr/feature/wsam-optimizer

[Feature] Implement WSAM optimizer
kozistr · Dec 10, 2023 · 5ff4dee · 5ff4dee
2 parents b6efbbc + 47aac68
commit 5ff4dee
Show file tree

Hide file tree

Showing 10 changed files with 551 additions and 306 deletions.
diff --git a/README.md b/README.md
diff --git a/docs/optimizer.md b/docs/optimizer.md
@@ -272,6 +272,10 @@
     :docstring:
     :members:
 
+::: pytorch_optimizer.WSAM
+    :docstring:
+    :members:
+
 ::: pytorch_optimizer.Yogi
     :docstring:
     :members:
diff --git a/pyproject.toml b/pyproject.toml
@@ -16,8 +16,8 @@ keywords = [
     "DAdaptSGD", "DAdaptLion", "DiffGrad", "Fromage", "Gravity", "GSAM", "LARS", "Lamb", "Lion", "LOMO", "Lookahead",
     "MADGRAD", "MSVAG", "Nero", "NovoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam",
     "Ranger", "Ranger21", "RotoGrad", "SAM", "SGDP", "Shampoo", "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SopihaH",
-    "SRMM", "SWATS", "Tiger", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard",
-    "Bi-Tempered", "Tversky", "FocalTversky", "LovaszHinge", "bitsandbytes",
+    "SRMM", "SWATS", "Tiger", "WSAM", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM",
+    "Jaccard", "Bi-Tempered", "Tversky", "FocalTversky", "LovaszHinge", "bitsandbytes",
 ]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",
@@ -126,7 +126,6 @@ testpaths = "tests"
 
 [tool.coverage.run]
 omit = [
-    "./pytorch_optimizer/optimizer/gsam.py",
     "./pytorch_optimizer/optimizer/rotograd.py",
 ]
 

diff --git a/pytorch_optimizer/__init__.py b/pytorch_optimizer/__init__.py
@@ -55,7 +55,6 @@
 from pytorch_optimizer.optimizer.fromage import Fromage
 from pytorch_optimizer.optimizer.gc import centralize_gradient
 from pytorch_optimizer.optimizer.gravity import Gravity
-from pytorch_optimizer.optimizer.gsam import GSAM
 from pytorch_optimizer.optimizer.lamb import Lamb
 from pytorch_optimizer.optimizer.lars import LARS
 from pytorch_optimizer.optimizer.lion import Lion
@@ -76,7 +75,7 @@
 from pytorch_optimizer.optimizer.ranger import Ranger
 from pytorch_optimizer.optimizer.ranger21 import Ranger21
 from pytorch_optimizer.optimizer.rotograd import RotoGrad
-from pytorch_optimizer.optimizer.sam import SAM
+from pytorch_optimizer.optimizer.sam import GSAM, SAM, WSAM
 from pytorch_optimizer.optimizer.sgd import ASGD, SGDW, AccSGD, SignSGD
 from pytorch_optimizer.optimizer.sgdp import SGDP
 from pytorch_optimizer.optimizer.shampoo import ScalableShampoo, Shampoo

diff --git a/pytorch_optimizer/optimizer/gsam.py b/pytorch_optimizer/optimizer/gsam.py