Merge pull request #106 from kozistr/deps/pytorch-version

[Deps] Support Pytorch 2.0

Author: kozistr

README.rst

@@ -31,7 +31,7 @@ Install
 
     $ pip3 install -U pytorch-optimizer
 
-or
+If there's a version issue when installing the package, try with `--no-deps` option.
 
 ::
 

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "2.4.1"
+version = "2.4.2"
 description = "optimizer & lr scheduler implementations in PyTorch with clean-code, strict types. Also, including useful optimization ideas."
 license = "Apache-2.0"
 authors = ["kozistr <[email protected]>"]
@@ -34,16 +34,16 @@ classifiers = [
 [tool.poetry.dependencies]
 python = "^3.7.2"
 numpy = [
-    { version = "=1.21.1", python = ">=3.7,<3.8"},
-    { version = "*", python = ">=3.8"},
+    { version = "=1.21.1", python = ">=3.7,<3.8" },
+    { version = "*", python = ">=3.8" },
 ]
-torch = { version = "^1.10", source = "torch"}
+torch = { version = ">=1.10", source = "torch" }
 
 [tool.poetry.dev-dependencies]
-isort = "^5.11.4"
-black = "^22.12.0"
-ruff = "^0.0.237"
-pytest = "^7.2.0"
+isort = "^5.11.5"
+black = "^23.1.0"
+ruff = "^0.0.244"
+pytest = "^7.2.1"
 pytest-cov = "^4.0.0"
 
 [[tool.poetry.source]]
@@ -53,7 +53,7 @@ secondary = true
 
 [tool.ruff]
 select = ["A", "B", "C4", "D", "E", "F", "G", "I", "N", "S", "T", "ISC", "ICN", "W", "INP", "PIE", "T20", "RET", "SIM", "TID", "ARG", "ERA", "RUF", "YTT", "PL"]
-ignore = ["D100", "D102", "D104", "D105", "D107", "D203", "D213", "PIE790", "PLR2004"]
+ignore = ["D100", "D102", "D104", "D105", "D107", "D203", "D213", "PIE790", "PLR0912", "PLR0913", "PLR0915", "PLR2004"]
 fixable = ["A", "B", "C", "D", "E", "F"]
 unfixable = ["F401"]
 exclude = [

pyproject.toml

@@ -119,7 +119,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                     pre_cond.add_(grad @ grad_t)
                     if state['step'] % self.preconditioning_compute_steps == 0:
-                        inv_pre_cond = compute_power_svd(pre_cond, -1.0 / order)
+                        inv_pre_cond.copy_(compute_power_svd(pre_cond, -1.0 / order))
 
                     if dim_id == order - 1:
                         grad = grad_t @ inv_pre_cond
@@ -151,8 +151,12 @@ class ScalableShampoo(Optimizer, BaseOptimizer):
     :param inverse_exponent_override: int. fixed exponent for pre-conditioner, if > 0.
     :param start_preconditioning_step: int.
     :param preconditioning_compute_steps: int. performance tuning params for controlling memory and compute
-        requirements. How often to compute pre-conditioner.
-    :param statistics_compute_steps: int. How often to compute statistics.
+        requirements. How often to compute pre-conditioner. Ideally, 1 is the best. However, the current implementation
+        doesn't work on the distributed environment (there are no statistics & pre-conditioners sync among replicas),
+        compute on the GPU (not CPU) and the precision is fp32 (not fp64).
+        Also, followed by the paper, `preconditioning_compute_steps` does not have a significant effect on the
+        performance. So, If you have a problem with the speed, try to set this step bigger (e.g. 1000).
+    :param statistics_compute_steps: int. How often to compute statistics. usually set to 1 (or 10).
     :param block_size: int. Block size for large layers (if > 0).
         Block size = 1 ==> Adagrad (Don't do this, extremely inefficient!)
         Block size should be as large as feasible under memory/time constraints.
@@ -166,8 +170,8 @@ class ScalableShampoo(Optimizer, BaseOptimizer):
     :param diagonal_eps: float. term added to the denominator to improve numerical stability.
     :param matrix_eps: float. term added to the denominator to improve numerical stability.
     :param use_svd: bool. use SVD instead of Schur-Newton method to calculate M^{-1/p}.
-        Theoretically, Schur-Newton method is faster than SVD method to calculate M^{-1/p}.
-        However, the inefficiency of the loop code, SVD is much faster than that.
+        Theoretically, Schur-Newton method is faster than SVD method. However, the inefficiency of the loop code and
+        proper svd kernel, SVD is much faster in some cases (usually in case of small models).
         see https://github.com/kozistr/pytorch_optimizer/pull/103
     """
 

pytorch_optimizer/optimizer/shampoo.py

@@ -1,26 +1,26 @@
 --extra-index-url https://download.pytorch.org/whl/cpu
 
 attrs==22.2.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
-black==22.12.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
+black==23.1.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 click==8.1.3 ; 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 python_full_version >= "3.7.2" and python_full_version < "4.0.0" and platform_system == "Windows"
 coverage[toml]==7.1.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 exceptiongroup==1.1.0 ; python_full_version >= "3.7.2" and python_version < "3.11"
 importlib-metadata==6.0.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.4 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
-mypy-extensions==0.4.3 ; 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_full_version < "4.0.0"
+mypy-extensions==1.0.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 numpy==1.21.1 ; python_full_version >= "3.7.2" and python_version < "3.8"
-numpy==1.24.1 ; python_version >= "3.8" and python_full_version < "4.0.0"
+numpy==1.24.2 ; python_version >= "3.8" and python_full_version < "4.0.0"
 packaging==23.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 pathspec==0.11.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
-platformdirs==2.6.2 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
+platformdirs==3.0.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 pluggy==1.0.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 pytest-cov==4.0.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 pytest==7.2.1 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
-ruff==0.0.237 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
-tomli==2.0.1 ; python_full_version >= "3.7.2" and python_full_version < "3.11.0a7"
+ruff==0.0.244 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
+tomli==2.0.1 ; python_full_version >= "3.7.2" and python_full_version <= "3.11.0a6"
 torch==1.13.1+cpu ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 typed-ast==1.5.4 ; python_version < "3.8" and implementation_name == "cpython" and python_full_version >= "3.7.2"
 typing-extensions==4.4.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
-zipp==3.12.0 ; python_full_version >= "3.7.2" and python_version < "3.8"
+zipp==3.13.0 ; 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
 
 numpy==1.21.1 ; python_full_version >= "3.7.2" and python_version < "3.8"
-numpy==1.24.1 ; python_version >= "3.8" and python_full_version < "4.0.0"
+numpy==1.24.2 ; python_version >= "3.8" and python_full_version < "4.0.0"
 torch==1.13.1+cpu ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"
 typing-extensions==4.4.0 ; python_full_version >= "3.7.2" and python_full_version < "4.0.0"