torchzero 0.1.7__py3-none-any.whl → 0.1.8__py3-none-any.whl

torchzero/core/module.py

@@ -212,6 +212,22 @@ class OptimizerModule(TensorListOptimizer, ABC): # type:ignore
         if self._initialized: return super().__repr__()
         return f"uninitialized {self.__class__.__name__}()"
 
+    def state_dict(self):
+        state_dict = {}
+        state_dict['__self__'] = super().state_dict()
+        for k,v in self.children.items():
+            state_dict[k] = v.state_dict()
+        return state_dict
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        super().load_state_dict(state_dict['__self__'])
+        for k, v in self.children.items():
+            if k in state_dict:
+                v.load_state_dict(state_dict[k])
+            else:
+                warnings.warn(f"Tried to load state dict for {k}: {v.__class__.__name__}, but it is not present in state_dict with {list(state_dict.keys()) = }")
+
+
     def set_params(self, params: ParamsT):
         """
         Set parameters to this module. Use this to set per-parameter group settings.

torchzero/optim/modular.py

@@ -2,6 +2,7 @@ from collections import abc
 import warnings
 from inspect import cleandoc
 import torch
+from typing import Any
 
 from ..core import OptimizerModule, TensorListOptimizer, OptimizationVars, _Chain, _Chainable
 from ..utils.python_tools import flatten
@@ -67,6 +68,21 @@ class Modular(TensorListOptimizer):
             for hook in module.post_init_hooks:
                 hook(self, module)
 
+    def state_dict(self):
+        state_dict = {}
+        state_dict['__self__'] = super().state_dict()
+        for i,v in enumerate(self.unrolled_modules):
+            state_dict[str(i)] = v.state_dict()
+        return state_dict
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        super().load_state_dict(state_dict['__self__'])
+        for i,v in enumerate(self.unrolled_modules):
+            if str(i) in state_dict:
+                v.load_state_dict(state_dict[str(i)])
+            else:
+                warnings.warn(f"Tried to load state dict for {i}th module: {v.__class__.__name__}, but it is not present in state_dict with {list(state_dict.keys()) = }")
+
     def get_lr_module(self, last=True) -> OptimizerModule:
         """
         Retrieves the module in the chain that controls the learning rate.

torchzero-0.1.8.dist-info/METADATA

@@ -0,0 +1,130 @@
+Metadata-Version: 2.2
+Name: torchzero
+Version: 0.1.8
+Summary: Modular optimization library for PyTorch.
+Author-email: Ivan Nikishev <nkshv2@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2024 inikishev
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: Homepage, https://github.com/inikishev/torchzero
+Project-URL: Repository, https://github.com/inikishev/torchzero
+Project-URL: Issues, https://github.com/inikishev/torchzero/isses
+Keywords: optimization,optimizers,torch,neural networks,zeroth order,second order
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: numpy
+Requires-Dist: typing_extensions
+
+![example workflow](https://github.com/inikishev/torchzero/actions/workflows/tests.yml/badge.svg)
+
+# torchzero
+
+`torchzero` implements a large number of chainable optimization modules that can be chained together to create custom optimizers:
+
+```py
+import torchzero as tz
+
+optimizer = tz.Modular(
+    model.parameters(),
+    tz.m.Adam(),
+    tz.m.Cautious(),
+    tz.m.LR(1e-3),
+    tz.m.WeightDecay(1e-4)
+)
+
+# standard training loop
+for batch in dataset:
+    preds = model(batch)
+    loss = criterion(preds)
+    optimizer.zero_grad()
+    optimizer.step()
+```
+
+Each module takes the output of the previous module and applies a further transformation. This modular design avoids redundant code, such as reimplementing cautioning, orthogonalization, laplacian smoothing, etc for every optimizer. It is also easy to experiment with grafting, interpolation between different optimizers, and perhaps some weirder combinations like nested momentum.
+
+Modules are not limited to gradient transformations. They can perform other operations like line searches, exponential moving average (EMA) and stochastic weight averaging (SWA), gradient accumulation, gradient approximation, and more.
+
+There are over 100 modules, all accessible within the `tz.m` namespace. For example, the Adam update rule is available as `tz.m.Adam`. Complete list of modules is available in [documentation](https://torchzero.readthedocs.io/en/latest/autoapi/torchzero/modules/index.html).
+
+## Closure
+
+Some modules and optimizers in torchzero, particularly line-search methods and gradient approximation modules, require a closure function. This is similar to how `torch.optim.LBFGS` works in PyTorch. In torchzero, closure needs to accept a boolean backward argument (though the argument can have any name). When `backward=True`, the closure should zero out old gradients using `opt.zero_grad()`, and compute new gradients using `loss.backward()`.
+
+```py
+def closure(backward = True):
+    preds = model(inputs)
+    loss = loss_fn(preds, targets)
+
+    if backward:
+        optimizer.zero_grad()
+        loss.backward()
+    return loss
+
+optimizer.step(closure)
+```
+
+If you intend to use gradient-free methods, `backward` argument is still required in the closure. Simply leave it unused. Gradient-free and gradient approximation methods always call closure with `backward=False`.
+
+All built-in pytorch optimizers, as well as most custom ones, support closure too. So the code above will work with all other optimizers out of the box, and you can switch between different optimizers without rewriting your training loop.
+
+# Documentation
+
+For more information on how to create, use and extend torchzero modules, please refer to the documentation at [torchzero.readthedocs.io](https://torchzero.readthedocs.io/en/latest/index.html).
+
+# Extra
+
+Some other optimization related things in torchzero:
+
+### scipy.optimize.minimize wrapper
+
+scipy.optimize.minimize wrapper with support for both gradient and hessian via batched autograd
+
+```py
+from torchzero.optim.wrappers.scipy import ScipyMinimize
+opt = ScipyMinimize(model.parameters(), method = 'trust-krylov')
+```
+
+Use as any other closure-based optimizer, but make sure closure accepts `backward` argument. Note that it performs full minimization on each step.
+
+### Nevergrad wrapper
+
+[Nevergrad](https://github.com/facebookresearch/nevergrad) is an optimization library by facebook with an insane number of gradient free methods.
+
+```py
+from torchzero.optim.wrappers.nevergrad import NevergradOptimizer
+opt = NevergradOptimizer(bench.parameters(), ng.optimizers.NGOptBase, budget = 1000)
+```
+
+Use as any other closure-based optimizer, but make sure closure accepts `backward` argument.
+
+### NLopt wrapper
+
+[NLopt](https://nlopt.readthedocs.io/en/latest/NLopt_Algorithms/) is another optimization library similar to scipy.optimize.minimize, with a large number of both gradient based and gradient free methods.
+
+```py
+from torchzero.optim.wrappers.nlopt import NLOptOptimizer
+opt = NLOptOptimizer(bench.parameters(), 'LD_TNEWTON_PRECOND_RESTART', maxeval = 1000)
+```
+
+Use as any other closure-based optimizer, but make sure closure accepts `backward` argument. Note that it performs full minimization on each step.

{torchzero-0.1.7.dist-info → torchzero-0.1.8.dist-info}/RECORD

@@ -1,7 +1,7 @@
 torchzero/__init__.py,sha256=CCIYfhGNYMnRP_cdXL7DgocxkEWYUZYgB3Sf1T5tdYY,203
 torchzero/tensorlist.py,sha256=V9m5zJ44PtiJxTOiM7cADxYaWK-ogv5Tk_KnyQqm1oo,41601
 torchzero/core/__init__.py,sha256=hab7HAep0JIVeJ-EQhcOAB9oKIIo2MuCVn7yS3BFVYA,266
-torchzero/core/module.py,sha256=Dhf8Rn6zKpzzO3rzyxzJnRAe_tOKS35AgH-FQQKR16A,22362
+torchzero/core/module.py,sha256=iFgRX_PY2Y0WD8_w5q6ygry1VhvpsE3jO-ySS_YQkuQ,23023
 torchzero/core/tensorlist_optimizer.py,sha256=MenRzyNPQJWx26j4Tj2vbX2pEcRu_26HQPIpUu_nsFc,9957
 torchzero/modules/__init__.py,sha256=5f8kt2mMn1eo9YcjXc3ESW-bqMQIRf646V3zlr8UAO4,571
 torchzero/modules/adaptive/__init__.py,sha256=YBVDXCosr4-C-GFCoreHS3DFyHiYMhCbOWgdhVVaZ_E,161
@@ -69,7 +69,7 @@ torchzero/modules/weight_averaging/__init__.py,sha256=nJJRs68AV2G4rGwiiHNRfm6Xmt
 torchzero/modules/weight_averaging/ema.py,sha256=xdVnKC8PxCQCec4ad3ncvznvVsQM5O6EVfsOKsRr18k,2854
 torchzero/modules/weight_averaging/swa.py,sha256=syio5qq1vf76d5bAJU-I8Zc8U12bR2n8C9hEznCgr7s,6764
 torchzero/optim/__init__.py,sha256=vk6pIYJHWAGYJMdtJ1otsmVph-pdL5HwBg-CTeBCGso,253
-torchzero/optim/modular.py,sha256=tiuKXpLqq6jWwWWrfC_QU9XkErl2Saaw_xGau0frwzQ,6060
+torchzero/optim/modular.py,sha256=tg-VUcxDAhBzM4uR7SFkf5I2hV_SScfqXCbZpvV1Yzc,6788
 torchzero/optim/experimental/__init__.py,sha256=RqNzJu5mVl3T0u7cf4TBzSiA20M1kxTZVYWjSVhEHuU,585
 torchzero/optim/experimental/experimental.py,sha256=tMHZVbEXm3s6mMr7unFSvk_Jks3uAaAG0fzsH6gr098,10928
 torchzero/optim/experimental/ray_search.py,sha256=GYyssL64D6RiImrZ2tchoZJ04x9rX-Bp1y2nQXEGxX0,2662
@@ -97,8 +97,8 @@ torchzero/utils/compile.py,sha256=pYEyX8P26iCb_hFqAXC8IP2SSQrRfC7ZDhXS0vVCsfY,12
 torchzero/utils/derivatives.py,sha256=koLmuUcVcX41SrH_9rvfJyMXyHyocNLuZ-C8Kr2B7hk,4844
 torchzero/utils/python_tools.py,sha256=kkyDhoP695HhapfKrdjcrRbRAbcvB0ArP1pkxuVUlf0,1192
 torchzero/utils/torch_tools.py,sha256=sSBY5Bmk9LOAgPtaq-6TK4wDgPXsg6FIWxv8CVDx82k,3580
-torchzero-0.1.7.dist-info/LICENSE,sha256=r9ZciAoZoqKC_FNADE0ORukj1p1XhLXEbegdsAyqhJs,1087
-torchzero-0.1.7.dist-info/METADATA,sha256=kA1ASlxBXLMOh_T2pgyk2OHsI50sOwguDfDmjg8fgAU,6096
-torchzero-0.1.7.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-torchzero-0.1.7.dist-info/top_level.txt,sha256=isztuDR1ZGo8p2tORLa-vNuomcbLj7Xd208lhd-pVPs,10
-torchzero-0.1.7.dist-info/RECORD,,
+torchzero-0.1.8.dist-info/LICENSE,sha256=r9ZciAoZoqKC_FNADE0ORukj1p1XhLXEbegdsAyqhJs,1087
+torchzero-0.1.8.dist-info/METADATA,sha256=gtnLFgZ4XQwI7GO5U_p8fUFrt1X8V_Lyv9bvISGnSro,6058
+torchzero-0.1.8.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+torchzero-0.1.8.dist-info/top_level.txt,sha256=isztuDR1ZGo8p2tORLa-vNuomcbLj7Xd208lhd-pVPs,10
+torchzero-0.1.8.dist-info/RECORD,,

torchzero-0.1.7.dist-info/METADATA

@@ -1,120 +0,0 @@
-Metadata-Version: 2.2
-Name: torchzero
-Version: 0.1.7
-Summary: Modular optimization library for PyTorch.
-Author-email: Ivan Nikishev <nkshv2@gmail.com>
-License: MIT License
-        
-        Copyright (c) 2024 inikishev
-        
-        Permission is hereby granted, free of charge, to any person obtaining a copy
-        of this software and associated documentation files (the "Software"), to deal
-        in the Software without restriction, including without limitation the rights
-        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-        copies of the Software, and to permit persons to whom the Software is
-        furnished to do so, subject to the following conditions:
-        
-        The above copyright notice and this permission notice shall be included in all
-        copies or substantial portions of the Software.
-        
-        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-        SOFTWARE.
-        
-Project-URL: Homepage, https://github.com/inikishev/torchzero
-Project-URL: Repository, https://github.com/inikishev/torchzero
-Project-URL: Issues, https://github.com/inikishev/torchzero/isses
-Keywords: optimization,optimizers,torch,neural networks,zeroth order,second order
-Requires-Python: >=3.10
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: torch
-Requires-Dist: numpy
-Requires-Dist: typing_extensions
-
-![example workflow](https://github.com/inikishev/torchzero/actions/workflows/tests.yml/badge.svg)
-
-# torchzero
-
-This is a work-in-progress optimizers library for pytorch with composable zeroth, first, second order and quasi newton methods, gradient approximation, line searches and a whole lot of other stuff.
-
-Most optimizers are modular, meaning you can chain them like this:
-
-```py
-optimizer = torchzero.optim.Modular(model.parameters(), [*list of modules*])`
-```
-
-For example you might use `[ClipNorm(4), LR(1e-3), NesterovMomentum(0.9)]` for standard SGD with gradient clipping and nesterov momentum. Move `ClipNorm` to the end to clip the update instead of the gradients. If you don't have access to gradients, add a `RandomizedFDM()` at the beginning to approximate them via randomized finite differences. Add `Cautious()` to make the optimizer cautious.
-
-Each new module takes previous module update and works on it. That way there is no need to reimplement stuff like laplacian smoothing for all optimizers, and it is easy to experiment with grafting, interpolation between different optimizers, and perhaps some weirder combinations like nested momentum.
-
-# How to use
-
-All modules are defined in `torchzero.modules`. You can generally mix and match them however you want. Some pre-made optimizers are available in `torchzero.optim`.
-
-Some optimizers require closure, which should look like this:
-
-```py
-def closure(backward = True):
-  preds = model(inputs)
-  loss = loss_fn(preds, targets)
-
-  # if you can't call loss.backward(), and instead use gradient-free methods,
-  # they always call closure with backward=False.
-  # so you can remove the part below, but keep the unused backward argument.
-  if backward:
-    optimizer.zero_grad()
-    loss.backward()
-  return loss
-
-optimizer.step(closure)
-```
-
-This closure will also work with all built in pytorch optimizers, including LBFGS, all optimizers in this library, as well as most custom ones.
-
-# Contents
-
-Docs are available at [torchzero.readthedocs.io](https://torchzero.readthedocs.io/en/latest/). A preliminary list of all modules is available here <https://torchzero.readthedocs.io/en/latest/autoapi/torchzero/modules/index.html#classes>. Some of the implemented algorithms:
-
-- SGD/Rprop/RMSProp/AdaGrad/Adam as composable modules. They are also tested to exactly match built in pytorch versions.
-- Cautious Optimizers (<https://huggingface.co/papers/2411.16085>)
-- Optimizer grafting (<https://openreview.net/forum?id=FpKgG31Z_i9>)
-- Laplacian smoothing (<https://arxiv.org/abs/1806.06317>)
-- Polyak momentum, nesterov momentum
-- Gradient norm and value clipping, gradient normalization
-- Gradient centralization (<https://arxiv.org/abs/2004.01461>)
-- Learning rate droput (<https://pubmed.ncbi.nlm.nih.gov/35286266/>).
-- Forward gradient (<https://arxiv.org/abs/2202.08587>)
-- Gradient approximation via finite difference or randomized finite difference, which includes SPSA, RDSA, FDSA and Gaussian smoothing (<https://arxiv.org/abs/2211.13566v3>)
-- Various line searches
-- Exact Newton's method (with Levenberg-Marquardt regularization), newton with hessian approximation via finite difference, subspace finite differences newton.
-- Directional newton via one additional forward pass
-
-All modules should be quite fast, especially on models with many different parameters, due to `_foreach` operations.
-
-I am getting to the point where I can start focusing on good docs and tests. As of now, the code should be considered experimental, untested and subject to change, so feel free but be careful if using this for actual project.
-
-# Wrappers
-
-### scipy.optimize.minimize wrapper
-
-scipy.optimize.minimize wrapper with support for both gradient and hessian via batched autograd
-
-```py
-from torchzero.optim.wrappers.scipy import ScipyMinimize
-opt = ScipyMinimize(model.parameters(), method = 'trust-krylov')
-```
-
-Use as any other optimizer (make sure closure accepts `backward` argument like one from **How to use**). Note that it performs full minimization on each step.
-
-### Nevergrad wrapper
-
-```py
-opt = NevergradOptimizer(bench.parameters(), ng.optimizers.NGOptBase, budget = 1000)
-```
-
-Use as any other optimizer (make sure closure accepts `backward` argument like one from **How to use**).