torchzero


Nametorchzero JSON
Version 0.3.11 PyPI version JSON
download
home_pageNone
SummaryModular optimization library for PyTorch.
upload_time2025-07-14 17:50:38
maintainerNone
docs_urlNone
authorNone
requires_python>=3.10
licenseMIT 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.
keywords optimization optimizers torch neural networks zeroth order second order
VCS
bugtrack_url
requirements No requirements were recorded.
Travis-CI No Travis.
coveralls test coverage No coveralls.
            ![tests](https://github.com/inikishev/torchzero/actions/workflows/tests.yml/badge.svg)

# torchzero

**Modular optimization library for PyTorch**

`torchzero` is a PyTorch library providing a highly modular framework for creating and experimenting with a huge number of various optimization algorithms - various momentum techniques, gradient clipping, gradient approximations, line searches, quasi newton methods and more. All algorithms are implemented as modules that can be chained together freely.

## Installation

```bash
pip install torchzero
```

pip version is always the latest one. Or install from this repo

```bash
pip install git+https://github.com/inikishev/torchzero
```

**Dependencies:**

* Python >= 3.10
* `torch`
* `numpy`
* `typing_extensions`

## Quick Start / Usage Example

Basic example:

```python
import torch
from torch import nn
import torchzero as tz

# Define a simple model
model = nn.Linear(10, 1)
criterion = nn.MSELoss()
inputs = torch.randn(5, 10)
targets = torch.randn(5, 1)

# Create an optimizer
# The order of modules matters:
# 1. ClipValue: clips gradients to (-10, 10) range.
# 2. Adam: applies Adam update rule to clipped gradients.
# 3. NormalizeByEMA: stabilizes the update by normalizing it to an exponential
# moving average of past updates.
# 4. WeightDecay - decoupled weight decay (can also move after LR to fully decouple)
# 5. LR: Scales the computed update by the learning rate (supports LR schedulers).
optimizer = tz.Modular(
    model.parameters(),
    tz.m.ClipValue(10),
    tz.m.Adam(),
    tz.m.NormalizeByEMA(max_ema_growth=1.1),
    tz.m.WeightDecay(1e-4),
    tz.m.LR(1e-1),
)

# Standard training loop
for epoch in range(100):
    optimizer.zero_grad()
    output = model(inputs)
    loss = criterion(output, targets)
    loss.backward()
    optimizer.step()
    if (epoch+1) % 10 == 0: print(f"Epoch {epoch+1}, Loss: {loss.item()}")
```

## Overview of Available Modules

`torchzero` provides a huge number of various modules:

* **Optimizers**: Optimization algorithms.
  * `Adam`, `Adan`, `Adagrad`, `ESGD`, `FullMatrixAdagrad`, `LMAdagrad`, `AdaHessian`, `AdaptiveHeavyBall`, `OrthoGrad`, `Lion`, `MARS`, `MatrixMomentum`, `AdaptiveMatrixMomentum`, `Muon`, `RMSprop`, `Rprop`, `SAM`, `ASAM`, `MSAM`, `Shampoo`, `SOAP`, `SophiaH`.

  Additionally many other optimizers can be easily defined via modules:
  * Grams: `[tz.m.Adam(), tz.m.GradSign()]`
  * LaProp: `[tz.m.RMSprop(), tz.m.EMA(0.9)]`
  * Signum: `[tz.m.HeavyBall(), tz.m.Sign()]`
  * Efficient full-matrix version of any diagonal optimizer, like Adam: `[tz.m.LMAdagrad(beta=0.999, inner=tz.m.EMA(0.9)), tz.m.Debias(0.9, 0.999)]`
  * Cautious version of any optimizer, like SOAP: `[tz.m.SOAP(), tz.m.Cautious()]`

* **Momentum**:
  * `HeavyBall`: Classic momentum (Polyak's momentum).
  * `NAG`: Nesterov Accelerated Gradient.
  * `EMA`: Exponential moving average.
  * `Averaging` (`MedianAveraging`, `WeightedAveraging`): Simple, median, or weighted averaging of updates.
  * `Cautious`, `ScaleByGradCosineSimilarity`: Momentum cautioning.

* **Stabilization**: Gradient stabilization techniques.
  * `ClipNorm`: Clips gradient L2 norm.
  * `ClipValue`: Clips gradient values element-wise.
  * `Normalize`: Normalizes gradients to unit norm.
  * `Centralize`: Centralizes gradients by subtracting the mean.
  * `ClipNormByEMA`, `NormalizeByEMA`, `ClipValueByEMA`: Clipping/Normalization based on EMA of past values.
  * `ClipNormGrowth`, `ClipValueGrowth`: Limits norm or value growth.

* **Gradient approximations**: Methods for approximating gradients.
  * `FDM`: Finite difference method.
  * `RandomizedFDM` (`MeZO`, `SPSA`, `RDSA`, `Gaussian smoothing`): Randomized finite difference methods (also subspaces).
  * `ForwardGradient`: Randomized gradient approximation via forward mode automatic differentiation.

* **Second order**: Second order methods.
  * `Newton`: Classic Newton's method.
  * `InverseFreeNewton`: Inverse-free version of Newton's method.
  * `NewtonCG`: Matrix-free newton's method with conjugate gradient or minimal residual solvers.
  * `TruncatedNewtonCG`: Steihaug-Toint Trust-region NewtonCG via a truncated CG solver.
  * `NystromSketchAndSolve`: Nyström sketch-and-solve method.
  * `NystromPCG`: NewtonCG with Nyström preconditioning.
  * `HigherOrderNewton`: Higher order Newton's method with trust region.

* **Quasi-Newton**: Approximate second-order optimization methods.
  * `LBFGS`: Limited-memory BFGS.
  * `LSR1`: Limited-memory SR1.
  * `OnlineLBFGS`: Online LBFGS.
  * `BFGS`, `DFP`, `ICUM`, `PSB`, `SR1`, `SSVM`, `BroydenBad`, `BroydenGood`, `FletcherVMM`, `GradientCorrection`, `Greenstadt1`, `Greenstadt2`, `Horisho`, `McCormick`, `NewSSM`, `Pearson`, `ProjectedNewtonRaphson`, `ThomasOptimalMethod`, `ShorR`: Full-matrix quasi-newton methods.
  * `DiagonalBFGS`, `DiagonalSR1`, `DiagonalQuasiCauchi`, `DiagonalWeightedQuasiCauchi`, `DNRTR`, `NewDQN`: Diagonal quasi-newton methods.
  * `PolakRibiere`, `FletcherReeves`, `HestenesStiefel`, `DaiYuan`, `LiuStorey`, `ConjugateDescent`, `HagerZhang`, `HybridHS_DY`, `ProjectedGradientMethod`: Conjugate gradient methods.

* **Trust Region** Trust region can work with exact hessian or any of the quasi-newton methods (L-BFGS support is WIP)
  * `TrustCG`: Trust-region, uses a Steihaug-Toint truncated CG solver.
  * `CubicRegularization`: Cubic regularization, works better with exact hessian.

* **Line Search**:
  * `Backtracking`, `AdaptiveBacktracking`: Backtracking line searches (adaptive is my own).
  * `StrongWolfe`: Cubic interpolation line search satisfying strong Wolfe conditions.
  * `ScipyMinimizeScalar`: Wrapper for SciPy's scalar minimization for line search.

* **Learning Rate**:
  * `LR`: Controls learning rate and adds support for LR schedulers.
  * `PolyakStepSize`: Polyak's subgradient method.
  * `BarzilaiBorwein`: Barzilai-Borwein step-size.
  * `Warmup`, `WarmupNormCLip`: Learning rate warmup.

* **Projections**: This can implement things like GaLore but I haven't done that yet.
  <!-- * `FFTProjection`, `DCTProjection`: Use any update rule in Fourier or DCT domain (doesn't seem to help though).
  * `VectorProjection`, `TensorizeProjection`, `BlockPartition`, `TensorNormsProjection`: Structural projection methods (for block BFGS etc.). -->
  This is WIP
  * `To`: this casts everything to any other dtype and device for other modules, e.g. if you want better precision
  * `ViewAsReal`: put if you have complex paramters.

* **Smoothing**: Smoothing-based optimization methods.
  * `LaplacianSmoothing`: Laplacian smoothing for gradients (implements Laplacian Smooth GD).
  * `GaussianHomotopy`: Smoothing via randomized Gaussian homotopy.

* **Weight Decay**:.
  * `WeightDecay`: Standard L2 or L1 weight decay.

* **Ops**: This has low level operations, also stuff like grafting and gradient accumulation.

* **Wrappers**.
  * `Wrap`: Wraps any PyTorch optimizer, allowing to use it as a module.

* **Experimental**: various horrible atrocities

A complete list of modules is available in the [documentation](https://torchzero.readthedocs.io/en/latest/autoapi/torchzero/modules/index.html).

## Advanced Usage

### Closure

Certain modules, particularly line searches and gradient approximations require a closure, similar to L-BFGS in PyTorch. Also some modules require closure to accept an additional `backward` argument, refer to example below:

```python
# basic training loop
for inputs, targets in dataloader:

    def closure(backward=True): # make sure it is True by default
        preds = model(inputs)
        loss = criterion(preds, targets)

        if backward: # gradient approximations always call with backward=False.
            optimizer.zero_grad()
            loss.backward()

        return loss

    loss = optimizer.step(closure)
```

The code above will also work with any other optimizer because all PyTorch optimizers and most custom ones support closure, so there is no need to rewrite training loop.

Non-batched example (rosenbrock):

```py
import torchzero as tz

def rosen(x, y):
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2

W = torch.tensor([-1.1, 2.5], requires_grad=True)

def closure(backward=True):
    loss = rosen(*W)
    if backward:
        W.grad = None # same as opt.zero_grad()
        loss.backward()
    return loss

opt = tz.Modular([W], tz.m.NewtonCG(), tz.m.StrongWolfe())
for step in range(20):
    loss = opt.step(closure)
    print(f'{step} - {loss}')
```

### Module combinations

There are practically no rules to the ordering of the modules - anything will work. For example any method can be made zeroth order by putting it after some gradient approximation module such as GaussianSmoothing:

```python
opt = tz.Modular(
    bench.parameters(),
    tz.m.GaussianSmoothing(h=0.01, n_samples=10),
    tz.m.NewtonCG(hvp_method='forward'),
    tz.m.AdaptiveBacktracking(),
)
```

GaussianSmoothing actually creates a new **closure** which approximates the gradient. To NewtonCG this closure is just like
any other closure, so it works seamlessly.

Any module can be projected (this is how it will work once I implement GaLore, but I haven't done that yet):

```python
tz.m.GaLore([tz.m.GraftModules(tz.m.Shampoo(), tz.m.RMSprop()), tz.m.LR(1e-2)])
```

### Low level modules

torchzero provides a lot of low-level modules that can be used to recreate update rules, or combine existing update rules
in new ways. Here are some equivalent ways to make Adam in order of their involvement:

```python
tz.m.Adam()
```

```python
# Adam is debiased RMSprop applied to EMA
tz.m.RMSprop(0.999, debiased=True, init='zeros', inner=tz.m.EMA(0.9))
```

```python
tz.m.DivModules(
    tz.m.EMA(0.9, debiased=True),
    [tz.m.SqrtEMASquared(0.999, debiased=True), tz.m.Add(1e-8)]
)
```

```python
tz.m.DivModules(
    [tz.m.EMA(0.9), tz.m.Debias(beta1=0.9, beta2=0.999)],
    [tz.m.EMASquared(0.999), tz.m.Sqrt(), tz.m.Add(1e-8)]
)
```

```python
tz.m.DivModules(
    [tz.m.EMA(0.9), tz.m.Debias(beta1=0.9)],
    [
        tz.m.Pow(2),
        tz.m.EMA(0.999),
        tz.m.AccumulateMaximum() if amsgrad else tz.m.Identity(),
        tz.m.Sqrt(),
        tz.m.Debias2(beta=0.999),
        tz.m.Add(1e-8)]
)
```

### Quick guide to implementing new modules

Modules are quite similar to torch.optim.Optimizer, the main difference is that everything is stored in the Vars object,
not in the module itself. Also both per-parameter settings and state are stored in per-parameter dictionaries. Feel free to modify the example below.

```python
import torch
from torchzero.core import Module, Var

class HeavyBall(Module):
    def __init__(self, momentum: float = 0.9, dampening: float = 0):
        defaults = dict(momentum=momentum, dampening=dampening)
        super().__init__(defaults)

    def step(self, var: Var):
        # Var object holds all attributes used for optimization - parameters, gradient, update, etc.
        # a module takes a Var object, modifies it or creates a new one, and returns it
        # Var has a bunch of attributes, including parameters, gradients, update, closure, loss
        # for now we are only interested in update, and we will apply the heavyball rule to it.

        params = var.params
        update = var.get_update() # list of tensors

        exp_avg_list = []
        for p, u in zip(params, update):
            state = self.state[p]
            settings = self.settings[p]
            momentum = settings['momentum']
            dampening = settings['dampening']

            if 'momentum_buffer' not in state:
                state['momentum_buffer'] = torch.zeros_like(p)

            buf = state['momentum_buffer']
            u *= 1 - dampening

            buf.mul_(momentum).add_(u)

            # clone because further modules might modify exp_avg in-place
            # and it is part of self.state
            exp_avg_list.append(buf.clone())

        # set new update to var
        var.update = exp_avg_list
        return var
```

More in-depth guide will be available in the documentation in the future.

## Other stuff

There are also wrappers providing `torch.optim.Optimizer` interface for various other libraries. When using those, make sure closure has `backward` argument as described in **Advanced Usage**.

---

### Scipy

#### torchzero.optim.wrappers.scipy.ScipyMinimize

A wrapper for `scipy.optimize.minimize` with gradients and hessians supplied by pytorch autograd. Scipy provides implementations of the following methods: `'nelder-mead', 'powell', 'cg', 'bfgs', 'newton-cg', 'l-bfgs-b', 'tnc', 'cobyla', 'cobyqa', 'slsqp', 'trust-constr', 'dogleg', 'trust-ncg', 'trust-exact', 'trust-krylov'`.

#### torchzero.optim.wrappers.scipy.ScipyDE, ScipyDualAnnealing, ScipySHGO, ScipyDIRECT, ScipyBrute

Equivalent wrappers for other derivative free solvers available in `scipy.optimize`

---

### NLOpt

#### torchzero.optim.wrappers.nlopt.NLOptWrapper

A wrapper for [NLOpt](https://github.com/stevengj/nlopt) with gradients supplied by pytorch autograd. NLOpt is another popular library with many gradient based and gradient free [algorithms](https://nlopt.readthedocs.io/en/latest/NLopt_Algorithms/)

---

### Nevergrad

#### torchzero.optim.wrappers.nevergrad.NevergradWrapper

A wrapper for [nevergrad](https://facebookresearch.github.io/nevergrad/) which has a huge library of gradient free [algorithms](https://facebookresearch.github.io/nevergrad/optimizers_ref.html#optimizers)

---

### fast-cma-es

#### torchzero.optim.wrappers.fcmaes.FcmaesWrapper

A wrapper for [fast-cma-es](https://github.com/dietmarwo/fast-cma-es), which implements various gradient free algorithms. Notably it includes [BITEOPT](https://github.com/avaneev/biteopt) which seems to have very good performance in benchmarks.

# License

This project is licensed under the MIT License

# Project Links

The documentation is available at <https://torchzero.readthedocs.io/en/latest/>

            

Raw data

            {
    "_id": null,
    "home_page": null,
    "name": "torchzero",
    "maintainer": null,
    "docs_url": null,
    "requires_python": ">=3.10",
    "maintainer_email": null,
    "keywords": "optimization, optimizers, torch, neural networks, zeroth order, second order",
    "author": null,
    "author_email": "Ivan Nikishev <nkshv2@gmail.com>",
    "download_url": "https://files.pythonhosted.org/packages/bd/63/0e5ab34eafea7a35e270e701bfe2b5d172781565befc5a1779506a3a188e/torchzero-0.3.11.tar.gz",
    "platform": null,
    "description": "![tests](https://github.com/inikishev/torchzero/actions/workflows/tests.yml/badge.svg)\n\n# torchzero\n\n**Modular optimization library for PyTorch**\n\n`torchzero` is a PyTorch library providing a highly modular framework for creating and experimenting with a huge number of various optimization algorithms - various momentum techniques, gradient clipping, gradient approximations, line searches, quasi newton methods and more. All algorithms are implemented as modules that can be chained together freely.\n\n## Installation\n\n```bash\npip install torchzero\n```\n\npip version is always the latest one. Or install from this repo\n\n```bash\npip install git+https://github.com/inikishev/torchzero\n```\n\n**Dependencies:**\n\n* Python >= 3.10\n* `torch`\n* `numpy`\n* `typing_extensions`\n\n## Quick Start / Usage Example\n\nBasic example:\n\n```python\nimport torch\nfrom torch import nn\nimport torchzero as tz\n\n# Define a simple model\nmodel = nn.Linear(10, 1)\ncriterion = nn.MSELoss()\ninputs = torch.randn(5, 10)\ntargets = torch.randn(5, 1)\n\n# Create an optimizer\n# The order of modules matters:\n# 1. ClipValue: clips gradients to (-10, 10) range.\n# 2. Adam: applies Adam update rule to clipped gradients.\n# 3. NormalizeByEMA: stabilizes the update by normalizing it to an exponential\n# moving average of past updates.\n# 4. WeightDecay - decoupled weight decay (can also move after LR to fully decouple)\n# 5. LR: Scales the computed update by the learning rate (supports LR schedulers).\noptimizer = tz.Modular(\n    model.parameters(),\n    tz.m.ClipValue(10),\n    tz.m.Adam(),\n    tz.m.NormalizeByEMA(max_ema_growth=1.1),\n    tz.m.WeightDecay(1e-4),\n    tz.m.LR(1e-1),\n)\n\n# Standard training loop\nfor epoch in range(100):\n    optimizer.zero_grad()\n    output = model(inputs)\n    loss = criterion(output, targets)\n    loss.backward()\n    optimizer.step()\n    if (epoch+1) % 10 == 0: print(f\"Epoch {epoch+1}, Loss: {loss.item()}\")\n```\n\n## Overview of Available Modules\n\n`torchzero` provides a huge number of various modules:\n\n* **Optimizers**: Optimization algorithms.\n  * `Adam`, `Adan`, `Adagrad`, `ESGD`, `FullMatrixAdagrad`, `LMAdagrad`, `AdaHessian`, `AdaptiveHeavyBall`, `OrthoGrad`, `Lion`, `MARS`, `MatrixMomentum`, `AdaptiveMatrixMomentum`, `Muon`, `RMSprop`, `Rprop`, `SAM`, `ASAM`, `MSAM`, `Shampoo`, `SOAP`, `SophiaH`.\n\n  Additionally many other optimizers can be easily defined via modules:\n  * Grams: `[tz.m.Adam(), tz.m.GradSign()]`\n  * LaProp: `[tz.m.RMSprop(), tz.m.EMA(0.9)]`\n  * Signum: `[tz.m.HeavyBall(), tz.m.Sign()]`\n  * Efficient full-matrix version of any diagonal optimizer, like Adam: `[tz.m.LMAdagrad(beta=0.999, inner=tz.m.EMA(0.9)), tz.m.Debias(0.9, 0.999)]`\n  * Cautious version of any optimizer, like SOAP: `[tz.m.SOAP(), tz.m.Cautious()]`\n\n* **Momentum**:\n  * `HeavyBall`: Classic momentum (Polyak's momentum).\n  * `NAG`: Nesterov Accelerated Gradient.\n  * `EMA`: Exponential moving average.\n  * `Averaging` (`MedianAveraging`, `WeightedAveraging`): Simple, median, or weighted averaging of updates.\n  * `Cautious`, `ScaleByGradCosineSimilarity`: Momentum cautioning.\n\n* **Stabilization**: Gradient stabilization techniques.\n  * `ClipNorm`: Clips gradient L2 norm.\n  * `ClipValue`: Clips gradient values element-wise.\n  * `Normalize`: Normalizes gradients to unit norm.\n  * `Centralize`: Centralizes gradients by subtracting the mean.\n  * `ClipNormByEMA`, `NormalizeByEMA`, `ClipValueByEMA`: Clipping/Normalization based on EMA of past values.\n  * `ClipNormGrowth`, `ClipValueGrowth`: Limits norm or value growth.\n\n* **Gradient approximations**: Methods for approximating gradients.\n  * `FDM`: Finite difference method.\n  * `RandomizedFDM` (`MeZO`, `SPSA`, `RDSA`, `Gaussian smoothing`): Randomized finite difference methods (also subspaces).\n  * `ForwardGradient`: Randomized gradient approximation via forward mode automatic differentiation.\n\n* **Second order**: Second order methods.\n  * `Newton`: Classic Newton's method.\n  * `InverseFreeNewton`: Inverse-free version of Newton's method.\n  * `NewtonCG`: Matrix-free newton's method with conjugate gradient or minimal residual solvers.\n  * `TruncatedNewtonCG`: Steihaug-Toint Trust-region NewtonCG via a truncated CG solver.\n  * `NystromSketchAndSolve`: Nystr\u00f6m sketch-and-solve method.\n  * `NystromPCG`: NewtonCG with Nystr\u00f6m preconditioning.\n  * `HigherOrderNewton`: Higher order Newton's method with trust region.\n\n* **Quasi-Newton**: Approximate second-order optimization methods.\n  * `LBFGS`: Limited-memory BFGS.\n  * `LSR1`: Limited-memory SR1.\n  * `OnlineLBFGS`: Online LBFGS.\n  * `BFGS`, `DFP`, `ICUM`, `PSB`, `SR1`, `SSVM`, `BroydenBad`, `BroydenGood`, `FletcherVMM`, `GradientCorrection`, `Greenstadt1`, `Greenstadt2`, `Horisho`, `McCormick`, `NewSSM`, `Pearson`, `ProjectedNewtonRaphson`, `ThomasOptimalMethod`, `ShorR`: Full-matrix quasi-newton methods.\n  * `DiagonalBFGS`, `DiagonalSR1`, `DiagonalQuasiCauchi`, `DiagonalWeightedQuasiCauchi`, `DNRTR`, `NewDQN`: Diagonal quasi-newton methods.\n  * `PolakRibiere`, `FletcherReeves`, `HestenesStiefel`, `DaiYuan`, `LiuStorey`, `ConjugateDescent`, `HagerZhang`, `HybridHS_DY`, `ProjectedGradientMethod`: Conjugate gradient methods.\n\n* **Trust Region** Trust region can work with exact hessian or any of the quasi-newton methods (L-BFGS support is WIP)\n  * `TrustCG`: Trust-region, uses a Steihaug-Toint truncated CG solver.\n  * `CubicRegularization`: Cubic regularization, works better with exact hessian.\n\n* **Line Search**:\n  * `Backtracking`, `AdaptiveBacktracking`: Backtracking line searches (adaptive is my own).\n  * `StrongWolfe`: Cubic interpolation line search satisfying strong Wolfe conditions.\n  * `ScipyMinimizeScalar`: Wrapper for SciPy's scalar minimization for line search.\n\n* **Learning Rate**:\n  * `LR`: Controls learning rate and adds support for LR schedulers.\n  * `PolyakStepSize`: Polyak's subgradient method.\n  * `BarzilaiBorwein`: Barzilai-Borwein step-size.\n  * `Warmup`, `WarmupNormCLip`: Learning rate warmup.\n\n* **Projections**: This can implement things like GaLore but I haven't done that yet.\n  <!-- * `FFTProjection`, `DCTProjection`: Use any update rule in Fourier or DCT domain (doesn't seem to help though).\n  * `VectorProjection`, `TensorizeProjection`, `BlockPartition`, `TensorNormsProjection`: Structural projection methods (for block BFGS etc.). -->\n  This is WIP\n  * `To`: this casts everything to any other dtype and device for other modules, e.g. if you want better precision\n  * `ViewAsReal`: put if you have complex paramters.\n\n* **Smoothing**: Smoothing-based optimization methods.\n  * `LaplacianSmoothing`: Laplacian smoothing for gradients (implements Laplacian Smooth GD).\n  * `GaussianHomotopy`: Smoothing via randomized Gaussian homotopy.\n\n* **Weight Decay**:.\n  * `WeightDecay`: Standard L2 or L1 weight decay.\n\n* **Ops**: This has low level operations, also stuff like grafting and gradient accumulation.\n\n* **Wrappers**.\n  * `Wrap`: Wraps any PyTorch optimizer, allowing to use it as a module.\n\n* **Experimental**: various horrible atrocities\n\nA complete list of modules is available in the [documentation](https://torchzero.readthedocs.io/en/latest/autoapi/torchzero/modules/index.html).\n\n## Advanced Usage\n\n### Closure\n\nCertain modules, particularly line searches and gradient approximations require a closure, similar to L-BFGS in PyTorch. Also some modules require closure to accept an additional `backward` argument, refer to example below:\n\n```python\n# basic training loop\nfor inputs, targets in dataloader:\n\n    def closure(backward=True): # make sure it is True by default\n        preds = model(inputs)\n        loss = criterion(preds, targets)\n\n        if backward: # gradient approximations always call with backward=False.\n            optimizer.zero_grad()\n            loss.backward()\n\n        return loss\n\n    loss = optimizer.step(closure)\n```\n\nThe code above will also work with any other optimizer because all PyTorch optimizers and most custom ones support closure, so there is no need to rewrite training loop.\n\nNon-batched example (rosenbrock):\n\n```py\nimport torchzero as tz\n\ndef rosen(x, y):\n    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2\n\nW = torch.tensor([-1.1, 2.5], requires_grad=True)\n\ndef closure(backward=True):\n    loss = rosen(*W)\n    if backward:\n        W.grad = None # same as opt.zero_grad()\n        loss.backward()\n    return loss\n\nopt = tz.Modular([W], tz.m.NewtonCG(), tz.m.StrongWolfe())\nfor step in range(20):\n    loss = opt.step(closure)\n    print(f'{step} - {loss}')\n```\n\n### Module combinations\n\nThere are practically no rules to the ordering of the modules - anything will work. For example any method can be made zeroth order by putting it after some gradient approximation module such as GaussianSmoothing:\n\n```python\nopt = tz.Modular(\n    bench.parameters(),\n    tz.m.GaussianSmoothing(h=0.01, n_samples=10),\n    tz.m.NewtonCG(hvp_method='forward'),\n    tz.m.AdaptiveBacktracking(),\n)\n```\n\nGaussianSmoothing actually creates a new **closure** which approximates the gradient. To NewtonCG this closure is just like\nany other closure, so it works seamlessly.\n\nAny module can be projected (this is how it will work once I implement GaLore, but I haven't done that yet):\n\n```python\ntz.m.GaLore([tz.m.GraftModules(tz.m.Shampoo(), tz.m.RMSprop()), tz.m.LR(1e-2)])\n```\n\n### Low level modules\n\ntorchzero provides a lot of low-level modules that can be used to recreate update rules, or combine existing update rules\nin new ways. Here are some equivalent ways to make Adam in order of their involvement:\n\n```python\ntz.m.Adam()\n```\n\n```python\n# Adam is debiased RMSprop applied to EMA\ntz.m.RMSprop(0.999, debiased=True, init='zeros', inner=tz.m.EMA(0.9))\n```\n\n```python\ntz.m.DivModules(\n    tz.m.EMA(0.9, debiased=True),\n    [tz.m.SqrtEMASquared(0.999, debiased=True), tz.m.Add(1e-8)]\n)\n```\n\n```python\ntz.m.DivModules(\n    [tz.m.EMA(0.9), tz.m.Debias(beta1=0.9, beta2=0.999)],\n    [tz.m.EMASquared(0.999), tz.m.Sqrt(), tz.m.Add(1e-8)]\n)\n```\n\n```python\ntz.m.DivModules(\n    [tz.m.EMA(0.9), tz.m.Debias(beta1=0.9)],\n    [\n        tz.m.Pow(2),\n        tz.m.EMA(0.999),\n        tz.m.AccumulateMaximum() if amsgrad else tz.m.Identity(),\n        tz.m.Sqrt(),\n        tz.m.Debias2(beta=0.999),\n        tz.m.Add(1e-8)]\n)\n```\n\n### Quick guide to implementing new modules\n\nModules are quite similar to torch.optim.Optimizer, the main difference is that everything is stored in the Vars object,\nnot in the module itself. Also both per-parameter settings and state are stored in per-parameter dictionaries. Feel free to modify the example below.\n\n```python\nimport torch\nfrom torchzero.core import Module, Var\n\nclass HeavyBall(Module):\n    def __init__(self, momentum: float = 0.9, dampening: float = 0):\n        defaults = dict(momentum=momentum, dampening=dampening)\n        super().__init__(defaults)\n\n    def step(self, var: Var):\n        # Var object holds all attributes used for optimization - parameters, gradient, update, etc.\n        # a module takes a Var object, modifies it or creates a new one, and returns it\n        # Var has a bunch of attributes, including parameters, gradients, update, closure, loss\n        # for now we are only interested in update, and we will apply the heavyball rule to it.\n\n        params = var.params\n        update = var.get_update() # list of tensors\n\n        exp_avg_list = []\n        for p, u in zip(params, update):\n            state = self.state[p]\n            settings = self.settings[p]\n            momentum = settings['momentum']\n            dampening = settings['dampening']\n\n            if 'momentum_buffer' not in state:\n                state['momentum_buffer'] = torch.zeros_like(p)\n\n            buf = state['momentum_buffer']\n            u *= 1 - dampening\n\n            buf.mul_(momentum).add_(u)\n\n            # clone because further modules might modify exp_avg in-place\n            # and it is part of self.state\n            exp_avg_list.append(buf.clone())\n\n        # set new update to var\n        var.update = exp_avg_list\n        return var\n```\n\nMore in-depth guide will be available in the documentation in the future.\n\n## Other stuff\n\nThere are also wrappers providing `torch.optim.Optimizer` interface for various other libraries. When using those, make sure closure has `backward` argument as described in **Advanced Usage**.\n\n---\n\n### Scipy\n\n#### torchzero.optim.wrappers.scipy.ScipyMinimize\n\nA wrapper for `scipy.optimize.minimize` with gradients and hessians supplied by pytorch autograd. Scipy provides implementations of the following methods: `'nelder-mead', 'powell', 'cg', 'bfgs', 'newton-cg', 'l-bfgs-b', 'tnc', 'cobyla', 'cobyqa', 'slsqp', 'trust-constr', 'dogleg', 'trust-ncg', 'trust-exact', 'trust-krylov'`.\n\n#### torchzero.optim.wrappers.scipy.ScipyDE, ScipyDualAnnealing, ScipySHGO, ScipyDIRECT, ScipyBrute\n\nEquivalent wrappers for other derivative free solvers available in `scipy.optimize`\n\n---\n\n### NLOpt\n\n#### torchzero.optim.wrappers.nlopt.NLOptWrapper\n\nA wrapper for [NLOpt](https://github.com/stevengj/nlopt) with gradients supplied by pytorch autograd. NLOpt is another popular library with many gradient based and gradient free [algorithms](https://nlopt.readthedocs.io/en/latest/NLopt_Algorithms/)\n\n---\n\n### Nevergrad\n\n#### torchzero.optim.wrappers.nevergrad.NevergradWrapper\n\nA wrapper for [nevergrad](https://facebookresearch.github.io/nevergrad/) which has a huge library of gradient free [algorithms](https://facebookresearch.github.io/nevergrad/optimizers_ref.html#optimizers)\n\n---\n\n### fast-cma-es\n\n#### torchzero.optim.wrappers.fcmaes.FcmaesWrapper\n\nA wrapper for [fast-cma-es](https://github.com/dietmarwo/fast-cma-es), which implements various gradient free algorithms. Notably it includes [BITEOPT](https://github.com/avaneev/biteopt) which seems to have very good performance in benchmarks.\n\n# License\n\nThis project is licensed under the MIT License\n\n# Project Links\n\nThe documentation is available at <https://torchzero.readthedocs.io/en/latest/>\n",
    "bugtrack_url": null,
    "license": "MIT License\n        \n        Copyright (c) 2024 inikishev\n        \n        Permission is hereby granted, free of charge, to any person obtaining a copy\n        of this software and associated documentation files (the \"Software\"), to deal\n        in the Software without restriction, including without limitation the rights\n        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n        copies of the Software, and to permit persons to whom the Software is\n        furnished to do so, subject to the following conditions:\n        \n        The above copyright notice and this permission notice shall be included in all\n        copies or substantial portions of the Software.\n        \n        THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n        SOFTWARE.\n        ",
    "summary": "Modular optimization library for PyTorch.",
    "version": "0.3.11",
    "project_urls": {
        "Homepage": "https://github.com/inikishev/torchzero",
        "Issues": "https://github.com/inikishev/torchzero/isses",
        "Repository": "https://github.com/inikishev/torchzero"
    },
    "split_keywords": [
        "optimization",
        " optimizers",
        " torch",
        " neural networks",
        " zeroth order",
        " second order"
    ],
    "urls": [
        {
            "comment_text": null,
            "digests": {
                "blake2b_256": "d6b7bdf55fd2bff6316746eb4b47713a5463a6b772dfec01c30317c2ee8960e2",
                "md5": "9014e3f92b4d9cbcfda6744c6428391e",
                "sha256": "484f9c51e9411a4e1e292ac7399d1ddb9996f100f304b8bcfae47cd30ea634f5"
            },
            "downloads": -1,
            "filename": "torchzero-0.3.11-py3-none-any.whl",
            "has_sig": false,
            "md5_digest": "9014e3f92b4d9cbcfda6744c6428391e",
            "packagetype": "bdist_wheel",
            "python_version": "py3",
            "requires_python": ">=3.10",
            "size": 299782,
            "upload_time": "2025-07-14T17:50:36",
            "upload_time_iso_8601": "2025-07-14T17:50:36.889255Z",
            "url": "https://files.pythonhosted.org/packages/d6/b7/bdf55fd2bff6316746eb4b47713a5463a6b772dfec01c30317c2ee8960e2/torchzero-0.3.11-py3-none-any.whl",
            "yanked": false,
            "yanked_reason": null
        },
        {
            "comment_text": null,
            "digests": {
                "blake2b_256": "bd630e5ab34eafea7a35e270e701bfe2b5d172781565befc5a1779506a3a188e",
                "md5": "70c93fa9530f7a0b448c6e0bdb733735",
                "sha256": "38b9ea02d6fcf6b9b82e40ee65168c2724f97e1e5b15cf18c8eff67eecb5589f"
            },
            "downloads": -1,
            "filename": "torchzero-0.3.11.tar.gz",
            "has_sig": false,
            "md5_digest": "70c93fa9530f7a0b448c6e0bdb733735",
            "packagetype": "sdist",
            "python_version": "source",
            "requires_python": ">=3.10",
            "size": 225973,
            "upload_time": "2025-07-14T17:50:38",
            "upload_time_iso_8601": "2025-07-14T17:50:38.736470Z",
            "url": "https://files.pythonhosted.org/packages/bd/63/0e5ab34eafea7a35e270e701bfe2b5d172781565befc5a1779506a3a188e/torchzero-0.3.11.tar.gz",
            "yanked": false,
            "yanked_reason": null
        }
    ],
    "upload_time": "2025-07-14 17:50:38",
    "github": true,
    "gitlab": false,
    "bitbucket": false,
    "codeberg": false,
    "github_user": "inikishev",
    "github_project": "torchzero",
    "travis_ci": false,
    "coveralls": false,
    "github_actions": true,
    "lcname": "torchzero"
}
        
Elapsed time: 0.49966s