feat: adagrad optimizer support

CHANGELOG.md

@@ -13,9 +13,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Implement AdaGrad optimizer and exponential learning rate decay schedule by [@Benjamin-eecs](https://github.com/Benjamin-eecs) and [@XuehaiPan](https://github.com/XuehaiPan) in [#80](https://github.com/metaopt/torchopt/pull/80).
 - Enable tests on Windows by [@XuehaiPan](https://github.com/XuehaiPan) in [#140](https://github.com/metaopt/torchopt/pull/140).
 - Add `ruff` and `flake8` plugins integration by [@XuehaiPan](https://github.com/XuehaiPan) in [#138](https://github.com/metaopt/torchopt/pull/138) and [#139](https://github.com/metaopt/torchopt/pull/139).
-- Add more documentation on implicit differentiation by [@Benjamin-eecs](https://github.com/Benjamin-eecs) and [@XuehaiPan](https://github.com/XuehaiPan).
+- Add more documentation on implicit differentiation by [@Benjamin-eecs](https://github.com/Benjamin-eecs) and [@XuehaiPan](https://github.com/XuehaiPan) in [#143](https://github.com/metaopt/torchopt/pull/143).
 
 ### Changed
 

docs/source/api/api.rst

@@ -30,17 +30,23 @@ Functional Optimizers
 .. autosummary::
 
     FuncOptimizer
+    adagrad
     adam
-    sgd
-    rmsprop
     adamw
+    rmsprop
+    sgd
 
 Wrapper for Function Optimizer
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: FuncOptimizer
     :members:
 
+Functional AdaGrad Optimizer
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: adagrad
+
 Functional Adam Optimizer
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -51,16 +57,16 @@ Functional AdamW Optimizer
 
 .. autofunction:: adamw
 
-Functional SGD Optimizer
-~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: sgd
-
 Functional RMSProp Optimizer
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autofunction:: rmsprop
 
+Functional SGD Optimizer
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: sgd
+
 ------
 
 Classic Optimizers
@@ -70,10 +76,16 @@ Classic Optimizers
 
 .. autosummary::
 
+    AdaGrad
     Adam
-    SGD
-    RMSProp
     AdamW
+    RMSProp
+    SGD
+
+Classic AdaGrad Optimizer
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: AdaGrad
 
 Classic Adam Optimizer
 ~~~~~~~~~~~~~~~~~~~~~~
@@ -85,16 +97,16 @@ Classic AdamW Optimizer
 
 .. autoclass:: AdamW
 
-Classic SGD Optimizer
-~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: SGD
-
 Classic RMSProp Optimizer
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: RMSProp
 
+Classic SGD Optimizer
+~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: SGD
+
 ------
 
 Differentiable Meta-Optimizers
@@ -104,10 +116,16 @@ Differentiable Meta-Optimizers
 
 .. autosummary::
 
+    MetaAdaGrad
     MetaAdam
-    MetaSGD
-    MetaRMSProp
     MetaAdamW
+    MetaRMSProp
+    MetaSGD
+
+Differentiable Meta-AdaGrad Optimizer
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: MetaAdaGrad
 
 Differentiable Meta-Adam Optimizer
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -119,16 +137,16 @@ Differentiable Meta-AdamW Optimizer
 
 .. autoclass:: MetaAdamW
 
-Differentiable Meta-SGD Optimizer
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: MetaSGD
-
 Differentiable Meta-RMSProp Optimizer
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: MetaRMSProp
 
+Differentiable Meta-SGD Optimizer
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: MetaSGD
+
 ------
 
 Implicit Differentiation

docs/source/explicit_diff/explicit_diff.rst

@@ -53,10 +53,11 @@ For PyTorch-like API (e.g., ``step()``), we designed a base class :class:`torcho
 .. autosummary::
 
     torchopt.MetaOptimizer
+    torchopt.MetaAdaGrad
     torchopt.MetaAdam
-    torchopt.MetaSGD
-    torchopt.MetaRMSProp
     torchopt.MetaAdamW
+    torchopt.MetaRMSProp
+    torchopt.MetaSGD
 
 By combining low-level API :class:`torchopt.MetaOptimizer` with the previous functional optimizer, we can achieve high-level API:
 

docs/source/implicit_diff/implicit_diff.rst

@@ -38,7 +38,7 @@ In `IMAML <https://arxiv.org/abs/1909.04630>`_, the function :math:`F` in the fi
 Fixed-point Iteration
 ~~~~~~~~~~~~~~~~~~~~~
 
-Sometimes the inner-level optimal solution can also be achieved by fixed point where the optionality :math:`T` takes the form:
+Sometimes the inner-level optimal solution can also be achieved by fixed point where the optimality :math:`T` takes the form:
 
 .. math::
 

docs/source/optimizer/optim.rst

@@ -18,10 +18,11 @@ Currently, TorchOpt supports 4 functional optimizers: :func:`sgd`, :func:`adam`,
 .. autosummary::
 
     torchopt.FuncOptimizer
+    torchopt.adagrad
     torchopt.adam
-    torchopt.sgd
-    torchopt.rmsprop
     torchopt.adamw
+    torchopt.rmsprop
+    torchopt.sgd
 
 Apply Parameter Updates
 -----------------------
@@ -84,10 +85,12 @@ We offer original PyTorch APIs (e.g., ``zero_grad()`` or ``step()``) for traditi
 .. autosummary::
 
     torchopt.Optimizer
+    torchopt.AdaGrad
     torchopt.Adam
-    torchopt.SGD
-    torchopt.RMSProp
     torchopt.AdamW
+    torchopt.RMSProp
+    torchopt.SGD
+
 
 By combining low-level API :class:`torchopt.Optimizer` with the previous functional optimizer, we can achieve high-level API:
 

docs/source/spelling_wordlist.txt

@@ -172,3 +172,6 @@ abc
 ABCMeta
 subclasscheck
 ctx
+Duchi
+invertible
+AdaGrad

tests/test_alias.py

@@ -450,6 +450,70 @@ def test_adam_accelerated_cuda(
     _set_use_chain_flat(True)
 
 
+@helpers.parametrize(
+    dtype=[torch.float64],
+    lr=[1e-2, 1e-3, 1e-4],
+    lr_decay=[0.0, 1e-2],
+    initial_accumulator_value=[0.0, 1e-1],
+    eps=[1e-8],
+    inplace=[True, False],
+    weight_decay=[0.0, 1e-2],
+    maximize=[False, True],
+    use_chain_flat=[True, False],
+)
+def test_adagrad(
+    dtype: torch.dtype,
+    lr: float,
+    lr_decay: float,
+    initial_accumulator_value: float,
+    eps: float,
+    inplace: bool,
+    weight_decay: float,
+    maximize: bool,
+    use_chain_flat: bool,
+) -> None:
+    _set_use_chain_flat(use_chain_flat)
+
+    model, model_ref, model_base, loader = helpers.get_models(device='cpu', dtype=dtype)
+
+    fmodel, params, buffers = functorch.make_functional_with_buffers(model)
+    optim = torchopt.adagrad(
+        lr=lr,
+        lr_decay=lr_decay,
+        weight_decay=weight_decay,
+        initial_accumulator_value=initial_accumulator_value,
+        eps=eps,
+        maximize=maximize,
+    )
+    optim_state = optim.init(params)
+    optim_ref = torch.optim.Adagrad(
+        model_ref.parameters(),
+        lr=lr,
+        lr_decay=lr_decay,
+        weight_decay=weight_decay,
+        initial_accumulator_value=initial_accumulator_value,
+        eps=eps,
+        maximize=maximize,
+    )
+    for xs, ys in loader:
+        xs = xs.to(dtype=dtype)
+        pred = fmodel(params, buffers, xs)
+        pred_ref = model_ref(xs)
+        loss = F.cross_entropy(pred, ys)
+        loss_ref = F.cross_entropy(pred_ref, ys)
+
+        grads = torch.autograd.grad(loss, params, allow_unused=True)
+        updates, optim_state = optim.update(grads, optim_state, params=params, inplace=inplace)
+        params = torchopt.apply_updates(params, updates, inplace=inplace)
+
+        optim_ref.zero_grad()
+        loss_ref.backward()
+        optim_ref.step()
+
+    helpers.assert_model_all_close((params, buffers), model_ref, model_base, dtype=dtype)
+    _set_use_chain_flat(True)
+
+
 @helpers.parametrize(
     dtype=[torch.float64],
     lr=[1e-2, 1e-3, 1e-4],

tests/test_import.py

@@ -25,6 +25,7 @@ def test_accelerated_op_import() -> None:
 
 
 def test_alias_import() -> None:
+    torchopt.adagrad
     torchopt.adam
     torchopt.adamw
     torchopt.rmsprop
@@ -33,8 +34,8 @@ def test_alias_import() -> None:
     torchopt.alias.adamw
     torchopt.alias.rmsprop
     torchopt.alias.sgd
-    from torchopt import adam, adamw, rmsprop, sgd
-    from torchopt.alias import adam, adamw, rmsprop, sgd
+    from torchopt import adagrad, adam, adamw, rmsprop, sgd
+    from torchopt.alias import adagrad, adam, adamw, rmsprop, sgd
 
 
 def test_diff_import() -> None:
@@ -107,17 +108,23 @@ def test_nn_import() -> None:
 
 def test_optim_import() -> None:
     torchopt.FuncOptimizer
+    torchopt.MetaAdaGrad
+    torchopt.MetaAdagrad
     torchopt.MetaAdam
     torchopt.MetaAdamW
     torchopt.MetaRMSProp
     torchopt.MetaRMSprop
     torchopt.MetaSGD
+    torchopt.AdaGrad
+    torchopt.Adagrad
     torchopt.Adam
     torchopt.AdamW
     torchopt.Optimizer
     torchopt.RMSProp
     torchopt.RMSprop
     torchopt.SGD
+    torchopt.optim.meta.MetaAdaGrad
+    torchopt.optim.meta.MetaAdagrad
     torchopt.optim.meta.MetaAdam
     torchopt.optim.meta.MetaAdamW
     torchopt.optim.meta.MetaRMSProp
@@ -132,21 +139,27 @@ def test_optim_import() -> None:
     torchopt.optim.func.FuncOptimizer
     from torchopt import (
         SGD,
+        AdaGrad,
+        Adagrad,
         Adam,
         AdamW,
         FuncOptimizer,
+        MetaAdaGrad,
+        MetaAdagrad,
         MetaAdam,
         MetaAdamW,
         MetaOptimizer,
-        MetaRMSProp,
         MetaRMSprop,
+        MetaRMSProp,
         MetaSGD,
         Optimizer,
         RMSProp,
     )
     from torchopt.optim import SGD, Adam, AdamW, FuncOptimizer, Optimizer, RMSProp
     from torchopt.optim.func import FuncOptimizer
     from torchopt.optim.meta import (
+        MetaAdaGrad,
+        MetaAdagrad,
         MetaAdam,
         MetaAdamW,
         MetaOptimizer,