torchzero 0.3.4__tar.gz → 0.3.6__tar.gz

{torchzero-0.3.4 → torchzero-0.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.3.4
+Version: 0.3.6
 Summary: Modular optimization library for PyTorch.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 License: MIT License
@@ -37,26 +37,27 @@ Requires-Dist: numpy
 Requires-Dist: typing_extensions
 Dynamic: license-file
 
+![tests](https://github.com/inikishev/torchzero/actions/workflows/tests.yml/badge.svg)
+
 # torchzero
 
 **Modular optimization library for PyTorch**
 
-<!-- [![PyPI version](https://img.shields.io/pypi/v/torchzero.svg)](https://pypi.org/project/torchzero/)
-[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
-[![Build Status](https://img.shields.io/github/actions/workflow/status/torchzero/torchzero/ci.yml?branch=main)](https://github.com/torchzero/torchzero/actions)
-[![Documentation Status](https://readthedocs.org/projects/torchzero/badge/?version=latest)](https://torchzero.readthedocs.io/en/latest/?badge=latest) -->
-
-`torchzero` is a Python library providing a highly modular framework for creating and experimenting with optimization algorithms in PyTorch. It allows users to easily combine and customize various components of optimizers, such as momentum techniques, gradient clipping, line searches and more.
+`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.
 
-NOTE: torchzero is in active development, currently docs are in a state of flux and pip version is extremely outdated.
+NOTE: torchzero is in active development, currently docs are in a state of flux.
 
 ## Installation
 
 ```bash
-pip install git+https://github.com/inikishev/torchzero
+pip install torchzero
 ```
 
-(please don't use pip version yet, it is very outdated)
+pip version is always the latest one. Or install from this repo
+
+```bash
+pip install git+https://github.com/inikishev/torchzero
+```
 
 **Dependencies:**
 
@@ -65,34 +66,9 @@ pip install git+https://github.com/inikishev/torchzero
 * `numpy`
 * `typing_extensions`
 
-## Core Concepts
-
-<!-- ### Modular Design
-
-`torchzero` is built around a few key abstractions:
-
-* **`Module`**: The base class for all components in `torchzero`. Each `Module` implements a `step(vars)` method that processes the optimization variables.
-* **`Modular`**: The main optimizer class that chains together a sequence of `Module`s. It orchestrates the flow of data through the modules in the order they are provided.
-* **`Transform`**: A special type of `Module` designed for tensor transformations. These are often used for operations like applying momentum or scaling gradients.
-* **`Preconditioner`**: A subclass of `Transform`, typically used for preconditioning gradients (e.g., Adam, RMSprop).
-
-### `Vars` Object
-
-The `Vars` object is a data carrier that passes essential information between modules during an optimization step. It typically holds:
-
-* `params`: The model parameters.
-* `grad`: Gradients of the parameters.
-* `update`: The update to be applied to the parameters.
-* `loss`: The current loss value.
-* `closure`: A function to re-evaluate the model and loss (used by some line search algorithms and other modules that might need to recompute gradients or loss).
-
-### `TensorList`
-
-`torchzero` uses a custom `TensorList` class for efficient batched operations on lists of tensors. This allows for optimized performance when dealing with multiple parameter groups or complex update rules. -->
-
 ## Quick Start / Usage Example
 
-Here's a basic example of how to use `torchzero`:
+Basic example:
 
 ```python
 import torch
@@ -107,13 +83,16 @@ targets = torch.randn(5, 1)
 
 # Create an optimizer
 # The order of modules matters:
-# 1. SOAP: Computes the update.
-# 2. NormalizeByEMA: stabilizes the update by normalizing to an exponential moving average of past updates.
-# 3. WeightDecay - semi-decoupled, because it is applied after SOAP, but before LR
-# 4. LR: Scales the computed update by the learning rate (supports LR schedulers).
+# 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.SOAP(),
+    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),
@@ -131,9 +110,9 @@ for epoch in range(100):
 
 ## Overview of Available Modules
 
-`torchzero` provides a rich set of pre-built modules. Here are some key categories and examples:
+`torchzero` provides a huge number of various modules:
 
-* **Optimizers (`torchzero/modules/optimizers/`)**: Optimization algorithms.
+* **Optimizers**: Optimization algorithms.
   * `Adam`.
   * `Shampoo`.
   * `SOAP` (my current recommendation).
@@ -152,77 +131,74 @@ for epoch in range(100):
   * Full matrix version of any diagonal optimizer, like Adam: `tz.m.FullMatrixAdagrad(beta=0.999, inner=tz.m.EMA(0.9))`
   * Cautious version of any optimizer, like SOAP: `[tz.m.SOAP(), tz.m.Cautious()]`
 
-* **Clipping (`torchzero/modules/clipping/`)**: Gradient clipping techniques.
+* **Momentum**:
+  * `NAG`: Nesterov Accelerated Gradient.
+  * `HeavyBall`: Classic momentum (Polyak's momentum).
+  * `EMA`: Exponential moving average.
+  * `Averaging` (`Medianveraging`, `WeightedAveraging`): Simple, median, or weighted averaging of updates.
+  * `Cautious`, `ScaleByGradCosineSimilarity`: Momentum cautioning.
+  * `MatrixMomentum`, `AdaptiveMatrixMomentum`: Second order momentum.
+
+* **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 Approximation (`torchzero/modules/grad_approximation/`)**: Methods for approximating gradients.
-  * `FDM`: Finite Difference Method.
-  * `RandomizedFDM` (`MeZO`, `SPSA`, `RDSA`, `Gaussian smoothing`): Randomized Finite Difference Methods (also subspaces).
+
+* **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.
-* **Line Search (`torchzero/modules/line_search/`)**: Techniques for finding optimal step sizes.
-  * `Backtracking`, `AdaptiveBacktracking`: Backtracking line searches.
+
+* **Second order**: Second order methods.
+  * `Newton`: Classic Newton's method.
+  * `NewtonCG`: Matrix-free newton's method with conjugate gradient solver.
+  * `NystromSketchAndSolve`: Nyström sketch-and-solve method.
+  * `NystromPCG`: NewtonCG with Nyström preconditioning (my current recommendation).
+
+* **Quasi-Newton**: Approximate second-order optimization methods.
+  * `LBFGS`: Limited-memory BFGS.
+  * `LSR1`: Limited-memory SR1.
+  * `OnlineLBFGS`: Online LBFGS.
+  * `BFGS`, `SR1`, `DFP`, `BroydenGood`, `BroydenBad`, `Greenstadt1`, `Greenstadt2`, `ColumnUpdatingMethod`, `ThomasOptimalMethod`, `PSB`, `Pearson2`, `SSVM`: Classic full-matrix quasi-newton methods.
+  * `PolakRibiere`, `FletcherReeves`, `HestenesStiefel`, `DaiYuan`, `LiuStorey`, `ConjugateDescent`, `HagerZhang`, `HybridHS_DY`: Conjugate gradient methods.
+
+* **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.
   * `TrustRegion`: First order trust region method.
-* **Learning Rate (`torchzero/modules/lr/`)**: Learning rate control.
-  * `LR`: Applies a fixed learning rate.
+
+* **Learning Rate**:
+  * `LR`: Controls learning rate and adds support for LR schedulers.
   * `PolyakStepSize`: Polyak's method.
   * `Warmup`: Learning rate warmup.
-* **Momentum (`torchzero/modules/momentum/`)**: Momentum-based update modifications.
-  * `NAG`: Nesterov Accelerated Gradient.
-  * `HeavyBall`: Classic momentum (Polyak's momentum).
-  * `EMA`: Exponential moving average.
-  * `Averaging` (`Medianveraging`, `WeightedAveraging`): Simple, median, or weighted averaging of updates.
-  * `Cautious`, `ScaleByGradCosineSimilarity`: Momentum cautioning.
-  * `MatrixMomentum`, `AdaptiveMatrixMomentum`: Second order momentum.
-  <!-- * `CoordinateMomentum`: Momentum via random coordinates. -->
-* **Projections (`torchzero/modules/projections/`)**: Gradient projection techniques.
-  * `FFTProjection`, `DCTProjection`: Use any update rule in Fourier or DCT domain.
-  * `VectorProjection`, `TensorizeProjection`, `BlockPartition`, `TensorNormsProjection`: Structural projection methods.
-  <!-- * *(Note: DCT and Galore were commented out in the `__init__.py` I read, might be experimental or moved).* -->
-* **Quasi-Newton (`torchzero/modules/quasi_newton/`)**: Approximate second-order optimization methods.
-  * `LBFGS`: Limited-memory BFGS.
-  * `LSR1`: Limited-memory SR1.
-  * `OnlineLBFGS`: Online LBFGS.
-  <!-- * `ModularLBFGS`: A modular L-BFGS implementation (from experimental). -->
-  * `BFGS`, `SR1`, `DFP`, `BroydenGood`, `BroydenBad`, `Greenstadt1`, `Greenstadt2`, `ColumnUpdatingMethod`, `ThomasOptimalMethod`, `PSB`, `Pearson2`, `SSVM`: Classic full-matrix Quasi-Newton update formulas.
-  * Conjugate Gradient methods: `PolakRibiere`, `FletcherReeves`, `HestenesStiefel`, `DaiYuan`, `LiuStorey`, `ConjugateDescent`, `HagerZhang`, `HybridHS_DY`.
-* **Second Order (`torchzero/modules/second_order/`)**: Second order methods.
-  * `Newton`: Classic Newton's method.
-  * `NewtonCG`: Matrix-free newton's method with conjugate gradient solver.
-  * `NystromSketchAndSolve`: Nyström sketch-and-solve method.
-  * `NystromPCG`: NewtonCG with Nyström preconditioning.
-* **Smoothing (`torchzero/modules/smoothing/`)**: Techniques for smoothing the loss landscape or gradients.
-  * `LaplacianSmoothing`: Laplacian smoothing for gradients.
+
+* **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.).
+
+* **Smoothing**: Smoothing-based optimization methods.
+  * `LaplacianSmoothing`: Laplacian smoothing for gradients (implements Laplacian Smooth GD).
   * `GaussianHomotopy`: Smoothing via randomized Gaussian homotopy.
-* **Weight Decay (`torchzero/modules/weight_decay/`)**: Weight decay implementations.
+
+* **Weight Decay**:.
   * `WeightDecay`: Standard L2 or L1 weight decay.
-  <!-- * `DirectWeightDecay`: Applies weight decay directly to weights.
-  * `decay_weights_`: Functional form for decaying weights. -->
-* **Ops (`torchzero/modules/ops/`)**: Various tensor operations and utilities.
-  * `GradientAccumulation`: easy way to add gradient accumulation.
-  * `Unary*` (e.g., `Abs`, `Sqrt`, `Sign`): Unary operations.
-  * `Binary*` (e.g., `Add`, `Mul`, `Graft`): Binary operations.
-  * `Multi*` (e.g., `ClipModules`, `LerpModules`): Operations on multiple module outputs.
-  * `Reduce*` (e.g., `Mean`, `Sum`, `WeightedMean`): Reduction operations on multiple module outputs.
-
-* **Wrappers (`torchzero/modules/wrappers/`)**.
+
+* **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 (`torchzero/modules/experimental/`)**: Experimental modules.
-  * `GradMin`: Attempts to minimize gradient norm.
-  * `ReduceOutwardLR`: Reduces learning rate for parameters with outward pointing gradients.
-  * `RandomSubspacePreconditioning`, `HistorySubspacePreconditioning`: Preconditioning techniques using random or historical subspaces. -->
+* **Experimental**: various horrible atrocities
 
 ## Advanced Usage
 
 ### Closure
 
-Certain modules, particularly line searches and gradient approximations require a closure, similar to L-BFGS in PyTorch. In TorchZero closure accepts an additional `backward` argument, refer to example below:
+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
@@ -232,7 +208,7 @@ for inputs, targets in dataloader:
         preds = model(inputs)
         loss = criterion(preds, targets)
 
-        if backward:
+        if backward: # gradient approximations always call with backward=False.
             optimizer.zero_grad()
             loss.backward()
 
@@ -241,7 +217,7 @@ for inputs, targets in dataloader:
     loss = optimizer.step(closure)
 ```
 
-Also the closure above works with all PyTorch optimizers and most custom ones, so there is no need to rewrite the training loop.
+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):
 
@@ -266,9 +242,31 @@ for step in range(20):
     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
+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
@@ -276,20 +274,21 @@ 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, amsgrad=amsgrad), tz.m.Add(1e-8)]
+    [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, amsgrad=amsgrad), tz.m.Sqrt(), tz.m.Add(1e-8)]
+    [tz.m.EMASquared(0.999), tz.m.Sqrt(), tz.m.Add(1e-8)]
 )
 ```
 
@@ -306,8 +305,6 @@ tz.m.DivModules(
 )
 ```
 
-There are practically no rules to the ordering of the modules - anything will work, even line search after line search or nested gaussian homotopy.
-
 ### 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,
@@ -354,14 +351,17 @@ class HeavyBall(Module):
         return vars
 ```
 
-There are a some specialized base modules.
+There are a some specialized base modules that make it much easier to implement some specific things.
 
 * `GradApproximator` for gradient approximations
 * `LineSearch` for line searches
-* `Preconditioner` for gradient preconditioners
+* `Preconditioner` for preconditioners
+* `Projection` for projections like GaLore or into fourier domain.
 * `QuasiNewtonH` for full-matrix quasi-newton methods that update hessian inverse approximation (because they are all very similar)
 * `ConguateGradientBase` for conjugate gradient methods, basically the only difference is how beta is calculated.
 
+The documentation on how to actually use them is to write itself in the near future.
+
 ## License
 
 This project is licensed under the MIT License
@@ -369,11 +369,11 @@ This project is licensed under the MIT License
 ## Project Links
 
 TODO (there are docs but from very old version)
-<!-- * **Homepage**: `https://torchzero.github.io/torchzero/` (Placeholder - update if available)
-* **Repository**: `https://github.com/torchzero/torchzero` (Assuming this is the correct path) -->
 
 ## Other stuff
 
 There are also wrappers providing `torch.optim.Optimizer` interface for for `scipy.optimize`, NLOpt and Nevergrad.
 
 They are in `torchzero.optim.wrappers.scipy.ScipyMinimize`, `torchzero.optim.wrappers.nlopt.NLOptOptimizer`, and `torchzero.optim.wrappers.nevergrad.NevergradOptimizer`. Make sure closure has `backward` argument as described in **Advanced Usage**.
+
+Apparently https://github.com/avaneev/biteopt is diabolical so I will add a wrapper for it too very soon.