diff --git a/.deepsource.toml b/.deepsource.toml new file mode 100644 index 00000000..8294772c --- /dev/null +++ b/.deepsource.toml @@ -0,0 +1,17 @@ +version = 1 + +test_patterns = ["tests/test_*.py"] + +exclude_patterns = [ + "tests/test_optimizer.py", + "docs/conf.py", + "tests/utils.py", + "examples/*" +] + +[[analyzers]] +name = "python" +enabled = true + + [analyzers.meta] + runtime_version = "3.x.x" diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 00000000..51a941fb --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,21 @@ +version: 2 +updates: +- package-ecosystem: pip + directory: "/" + schedule: + interval: weekly + day: "monday" + time: "10:00" + open-pull-requests-limit: 10 + ignore: + - dependency-name: numpy + versions: + - 1.20.0 + - 1.20.1 + - dependency-name: isort + versions: + - 5.7.0 + - dependency-name: ipython + versions: + - 7.19.0 + - 7.20.0 diff --git a/.github/workflows/codeql-analysis.yml b/.github/workflows/codeql-analysis.yml new file mode 100644 index 00000000..d87d25cc --- /dev/null +++ b/.github/workflows/codeql-analysis.yml @@ -0,0 +1,67 @@ +# For most projects, this workflow file will not need changing; you simply need +# to commit it to your repository. +# +# You may wish to alter this file to override the set of languages analyzed, +# or to provide custom queries or build logic. +# +# ******** NOTE ******** +# We have attempted to detect the languages in your repository. Please check +# the `language` matrix defined below to confirm you have the correct set of +# supported CodeQL languages. +# +name: "CodeQL" + +on: + push: + branches: [ master ] + pull_request: + # The branches below must be a subset of the branches above + branches: [ master ] + schedule: + - cron: '39 10 * * 5' + +jobs: + analyze: + name: Analyze + runs-on: ubuntu-latest + + strategy: + fail-fast: false + matrix: + language: [ 'python' ] + # CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python' ] + # Learn more: + # https://docs.github.com/en/free-pro-team@latest/github/finding-security-vulnerabilities-and-errors-in-your-code/configuring-code-scanning#changing-the-languages-that-are-analyzed + + steps: + - name: Checkout repository + uses: actions/checkout@v2 + + # Initializes the CodeQL tools for scanning. + - name: Initialize CodeQL + uses: github/codeql-action/init@v1 + with: + languages: ${{ matrix.language }} + # If you wish to specify custom queries, you can do so here or in a config file. + # By default, queries listed here will override any specified in a config file. + # Prefix the list here with "+" to use these queries and those in the config file. + # queries: ./path/to/local/query, your-org/your-repo/queries@main + + # Autobuild attempts to build any compiled languages (C/C++, C#, or Java). + # If this step fails, then you should remove it and run the build manually (see below) + - name: Autobuild + uses: github/codeql-action/autobuild@v1 + + # ℹī¸ Command-line programs to run using the OS shell. + # 📚 https://git.io/JvXDl + + # ✏ī¸ If the Autobuild fails above, remove it and uncomment the following three lines + # and modify them (or add more) to build your code if your project + # uses a compiled language + + #- run: | + # make bootstrap + # make release + + - name: Perform CodeQL Analysis + uses: github/codeql-action/analyze@v1 diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml new file mode 100644 index 00000000..a8f8c80c --- /dev/null +++ b/.github/workflows/python-package.yml @@ -0,0 +1,71 @@ +# This workflow will install Python dependencies, run tests and lint with a variety of Python versions +# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions + +name: CI + +on: + push: + branches: + - master + tags: [ 'v*' ] + pull_request: + branches: + - master + schedule: + - cron: '0 6 * * 1' # Weekly Mon 6AM UTC build + + +jobs: + test: + + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ['3.7', '3.8', '3.9'] + + steps: + - uses: actions/checkout@v2 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install codecov + pip install -r requirements-dev.txt + - name: Lint + run: | + make lint + make checkbuild + + - name: Test + run: | + make cov + codecov + + deploy: + name: Deploy + runs-on: ubuntu-latest + needs: test + # Run only on pushing a tag + if: github.event_name == 'push' && contains(github.ref, 'refs/tags/') + steps: + - name: Checkout + uses: actions/checkout@v2 + - name: Setup Python 3.8 + uses: actions/setup-python@v2 + with: + python-version: 3.8 + - name: Install dependencies + run: + python -m pip install -U pip wheel twine + - name: Make dists + run: + python setup.py sdist bdist_wheel + - name: PyPI upload + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_TOKEN }} + run: | + twine upload dist/* diff --git a/.gitignore b/.gitignore index 7b1bbaa0..ba33694e 100644 --- a/.gitignore +++ b/.gitignore @@ -60,3 +60,7 @@ target/ .coverage.* coverage .mypy_cache/ +.DS_Store +tags +cscope.* +TODO diff --git a/.travis.yml b/.travis.yml deleted file mode 100644 index 9a4720c5..00000000 --- a/.travis.yml +++ /dev/null @@ -1,20 +0,0 @@ -dist: xenial -language: python -sudo: required - - -python: - - '3.7' - - '3.6' - - -install: - - pip install --upgrade setuptools - - pip install codecov - - pip install -r requirements-dev.txt - -script: - make cov - -after_success: - codecov diff --git a/CHANGES.rst b/CHANGES.rst index 5dec5649..8dfff3d6 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -1,6 +1,25 @@ Changes ------- -0.0.1 (YYYY-MM-DD) +0.3.1 (YYYY-MM-DD) +------------------ +* Deprecate RAdam optimizer. + +0.3.0 (2021-10-30) +------------------ +* Revert for Drop RAdam. + +0.2.0 (2021-10-25) +------------------ +* Drop RAdam optimizer since it is included in pytorch. +* Do not include tests as installable package. +* Preserver memory layout where possible. +* Add MADGRAD optimizer. + +0.1.0 (2021-01-01) ------------------ * Initial release. +* Added support for A2GradExp, A2GradInc, A2GradUni, AccSGD, AdaBelief, + AdaBound, AdaMod, Adafactor, Adahessian, AdamP, AggMo, Apollo, + DiffGrad, Lamb, Lookahead, NovoGrad, PID, QHAdam, QHM, RAdam, Ranger, + RangerQH, RangerVA, SGDP, SGDW, SWATS, Shampoo, Yogi. diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 00000000..53d09c8d --- /dev/null +++ b/CITATION.cff @@ -0,0 +1,9 @@ +cff-version: 1.2.0 +message: "If you use this software, please cite it as below." +authors: + - family-names: Novik + given-names: Mykola + orcid: https://orcid.org/0000-0002-0890-1159 +title: "torch-optimizer -- collection of optimization algorithms for PyTorch." +version: 1.0.1 +date-released: 2020-01-11 diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst new file mode 100644 index 00000000..d110944e --- /dev/null +++ b/CONTRIBUTING.rst @@ -0,0 +1,53 @@ +Contributing +============ + +Running Tests +------------- + +.. _GitHub: https://github.com/jettify/pytorch-optimizer +.. _PyTorch: https://github.com/pytorch/pytorch + +Thanks for your interest in contributing to ``pytorch-optimizer``, there are multiple +ways and places you can contribute. + +First of all just clone repository:: + + $ git clone git@github.com:jettify/pytorch-optimizer.git + +Create virtualenv with python3.5 (older version are not supported). For example +using *virtualenvwrapper* commands could look like:: + + $ cd pytorch-optimizer + $ mkvirtualenv --python=`which python3.7` pytorch-optimizer + + +After that please install libraries required for development:: + + $ pip install -r requirements-dev.txt + $ pip install -e . + +Congratulations, you are ready to run the test suite:: + + $ make cov + +To run individual use following command:: + + $ py.test -sv tests/test_basic.py -k test_name + + +Reporting an Issue +------------------ +If you have found issue with `pytorch-optimizer` please do +not hesitate to file an issue on the GitHub_ project. When filing your +issue please make sure you can express the issue with a reproducible test +case. + +When reporting an issue we also need as much information about your environment +that you can include. We never know what information will be pertinent when +trying narrow down the issue. Please include at least the following +information: + +* Version of `pytorch-optimizer`, `python`. +* Version PyTorch_ if installed. +* Version or CUDA if installed. +* Platform you're running on (OS X, Linux). diff --git a/MANIFEST.in b/MANIFEST.in index cfc0ebdc..43beee9c 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -2,7 +2,8 @@ include LICENSE include CHANGES.rst include README.rst include Makefile -graft torch_inspect +graft torch_optimizer graft tests -global-exclude *.pyc prune docs/_build +include examples/mnist.py +global-exclude *.py[cod] diff --git a/Makefile b/Makefile index f75f7ba8..00503350 100644 --- a/Makefile +++ b/Makefile @@ -1,8 +1,11 @@ # Some simple testing tasks (sorry, UNIX only). +FILES := torch_optimizer tests examples setup.py + + flake: - flake8 torch_optimizer tests examples setup.py + flake8 ${FILES} test: flake pytest -sv @@ -20,12 +23,22 @@ bandit: bandit -r ./torch_optimizer mypy: - mypy torch_optimizer --ignore-missing-imports --strict + mypy torch_optimizer --ignore-missing-imports -cov cover coverage: flake checkrst pyroma bandit +checkbuild: + python setup.py sdist bdist_wheel + twine check dist/* + +cov cover coverage: pytest -sv -vv --cov=torch_optimizer --cov-report=term --cov-report=html ./tests @echo "open file://`pwd`/htmlcov/index.html" +checkfmt: + isort --profile black --check-only --diff $(FILES) + black -l 79 --check $(FILES) + +lint: flake checkrst pyroma bandit checkfmt + clean: rm -rf `find . -name __pycache__` rm -f `find . -type f -name '*.py[co]' ` @@ -40,12 +53,18 @@ clean: rm -rf build rm -rf cover rm -rf dist + rm -rf docs/_build doc: make -C docs html @echo "open file://`pwd`/docs/_build/html/index.html" black: - black -S -l 79 setup.py torch_optimizer/ tests/ examples/ + black -l 79 setup.py torch_optimizer/ tests/ examples/ + +fmt: + isort --profile black ${FILES} + black -l 79 ${FILES} + .PHONY: all flake test vtest cov clean doc diff --git a/README.rst b/README.rst index 7a0eb522..bc7d5b17 100644 --- a/README.rst +++ b/README.rst @@ -1,17 +1,1060 @@ torch-optimizer =============== -.. image:: https://travis-ci.com/jettify/pytorch-optimizer.svg?branch=master - :target: https://travis-ci.com/jettify/pytorch-optimizer +.. image:: https://github.com/jettify/pytorch-optimizer/workflows/CI/badge.svg + :target: https://github.com/jettify/pytorch-optimizer/actions?query=workflow%3ACI + :alt: GitHub Actions status for master branch .. image:: https://codecov.io/gh/jettify/pytorch-optimizer/branch/master/graph/badge.svg :target: https://codecov.io/gh/jettify/pytorch-optimizer .. image:: https://img.shields.io/pypi/pyversions/torch-optimizer.svg :target: https://pypi.org/project/torch-optimizer +.. image:: https://readthedocs.org/projects/pytorch-optimizer/badge/?version=latest + :target: https://pytorch-optimizer.readthedocs.io/en/latest/?badge=latest + :alt: Documentation Status .. image:: https://img.shields.io/pypi/v/torch-optimizer.svg :target: https://pypi.python.org/pypi/torch-optimizer +.. image:: https://static.deepsource.io/deepsource-badge-light-mini.svg + :target: https://deepsource.io/gh/jettify/pytorch-optimizer/?ref=repository-badge -**torch-optimizer** -- collection of optimizers. +**torch-optimizer** -- collection of optimizers for PyTorch_ compatible with optim_ +module. +Simple example +-------------- + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.DiffGrad(model.parameters(), lr=0.001) + optimizer.step() + + +Installation +------------ +Installation process is simple, just:: + + $ pip install torch_optimizer + + +Documentation +------------- +https://pytorch-optimizer.rtfd.io + + +Citation +-------- +Please cite the original authors of the optimization algorithms. If you like this +package:: + + @software{Novik_torchoptimizers, + title = {{torch-optimizer -- collection of optimization algorithms for PyTorch.}}, + author = {Novik, Mykola}, + year = 2020, + month = 1, + version = {1.0.1} + } + +Or use the github feature: "cite this repository" button. + + +Supported Optimizers +==================== + ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `A2GradExp`_ | https://arxiv.org/abs/1810.00553 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `A2GradInc`_ | https://arxiv.org/abs/1810.00553 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `A2GradUni`_ | https://arxiv.org/abs/1810.00553 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `AccSGD`_ | https://arxiv.org/abs/1803.05591 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `AdaBelief`_ | https://arxiv.org/abs/2010.07468 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `AdaBound`_ | https://arxiv.org/abs/1902.09843 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `AdaMod`_ | https://arxiv.org/abs/1910.12249 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `Adafactor`_ | https://arxiv.org/abs/1804.04235 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `Adahessian`_ | https://arxiv.org/abs/2006.00719 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `AdamP`_ | https://arxiv.org/abs/2006.08217 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `AggMo`_ | https://arxiv.org/abs/1804.00325 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `Apollo`_ | https://arxiv.org/abs/2009.13586 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `DiffGrad`_ | https://arxiv.org/abs/1909.11015 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `Lamb`_ | https://arxiv.org/abs/1904.00962 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `Lookahead`_ | https://arxiv.org/abs/1907.08610 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `MADGRAD`_ | https://arxiv.org/abs/2101.11075 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `NovoGrad`_ | https://arxiv.org/abs/1905.11286 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `PID`_ | https://www4.comp.polyu.edu.hk/~cslzhang/paper/CVPR18_PID.pdf | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `QHAdam`_ | https://arxiv.org/abs/1810.06801 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `QHM`_ | https://arxiv.org/abs/1810.06801 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `RAdam`_ | https://arxiv.org/abs/1908.03265 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `Ranger`_ | https://medium.com/@lessw/new-deep-learning-optimizer-ranger-synergistic-combination-of-radam-lookahead-for-the-best-of-2dc83f79a48d | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `RangerQH`_ | https://arxiv.org/abs/1810.06801 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `RangerVA`_ | https://arxiv.org/abs/1908.00700v2 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `SGDP`_ | https://arxiv.org/abs/2006.08217 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `SGDW`_ | https://arxiv.org/abs/1608.03983 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `SWATS`_ | https://arxiv.org/abs/1712.07628 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `Shampoo`_ | https://arxiv.org/abs/1802.09568 | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ +| | | +| `Yogi`_ | https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------+ + + +Visualizations +-------------- +Visualizations help us see how different algorithms deal with simple +situations like: saddle points, local minima, valleys etc, and may provide +interesting insights into the inner workings of an algorithm. Rosenbrock_ and Rastrigin_ +benchmark_ functions were selected because: + +* Rosenbrock_ (also known as banana function), is non-convex function that has + one global minimum `(1.0. 1.0)`. The global minimum is inside a long, + narrow, parabolic shaped flat valley. Finding the valley is trivial. + Converging to the global minimum, however, is difficult. Optimization + algorithms might pay a lot of attention to one coordinate, and struggle + following the valley which is relatively flat. + + .. image:: https://upload.wikimedia.org/wikipedia/commons/3/32/Rosenbrock_function.svg + +* Rastrigin_ is a non-convex function and has one global minimum in `(0.0, 0.0)`. + Finding the minimum of this function is a fairly difficult problem due to + its large search space and its large number of local minima. + + .. image:: https://upload.wikimedia.org/wikipedia/commons/8/8b/Rastrigin_function.png + +Each optimizer performs `501` optimization steps. Learning rate is the best one found +by a hyper parameter search algorithm, the rest of the tuning parameters are default. It +is very easy to extend the script and tune other optimizer parameters. + + +.. code:: + + python examples/viz_optimizers.py + + +Warning +------- +Do not pick an optimizer based on visualizations, optimization approaches +have unique properties and may be tailored for different purposes or may +require explicit learning rate schedule etc. The best way to find out is to try +one on your particular problem and see if it improves scores. + +If you do not know which optimizer to use, start with the built in SGD/Adam. Once +the training logic is ready and baseline scores are established, swap the optimizer +and see if there is any improvement. + + +A2GradExp +--------- + ++--------------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_A2GradExp.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_A2GradExp.png | ++--------------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.A2GradExp( + model.parameters(), + kappa=1000.0, + beta=10.0, + lips=10.0, + rho=0.5, + ) + optimizer.step() + + +**Paper**: *Optimal Adaptive and Accelerated Stochastic Gradient Descent* (2018) [https://arxiv.org/abs/1810.00553] + +**Reference Code**: https://github.com/severilov/A2Grad_optimizer + + +A2GradInc +--------- + ++--------------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_A2GradInc.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_A2GradInc.png | ++--------------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.A2GradInc( + model.parameters(), + kappa=1000.0, + beta=10.0, + lips=10.0, + ) + optimizer.step() + + +**Paper**: *Optimal Adaptive and Accelerated Stochastic Gradient Descent* (2018) [https://arxiv.org/abs/1810.00553] + +**Reference Code**: https://github.com/severilov/A2Grad_optimizer + + +A2GradUni +--------- + ++--------------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_A2GradUni.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_A2GradUni.png | ++--------------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.A2GradUni( + model.parameters(), + kappa=1000.0, + beta=10.0, + lips=10.0, + ) + optimizer.step() + + +**Paper**: *Optimal Adaptive and Accelerated Stochastic Gradient Descent* (2018) [https://arxiv.org/abs/1810.00553] + +**Reference Code**: https://github.com/severilov/A2Grad_optimizer + + +AccSGD +------ + ++-----------------------------------------------------------------------------------------------------------+------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_AccSGD.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_AccSGD.png | ++-----------------------------------------------------------------------------------------------------------+------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.AccSGD( + model.parameters(), + lr=1e-3, + kappa=1000.0, + xi=10.0, + small_const=0.7, + weight_decay=0 + ) + optimizer.step() + + +**Paper**: *On the insufficiency of existing momentum schemes for Stochastic Optimization* (2019) [https://arxiv.org/abs/1803.05591] + +**Reference Code**: https://github.com/rahulkidambi/AccSGD + + +AdaBelief +--------- + ++-------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_AdaBelief.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_AdaBelief.png | ++-------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.AdaBelief( + m.parameters(), + lr= 1e-3, + betas=(0.9, 0.999), + eps=1e-3, + weight_decay=0, + amsgrad=False, + weight_decouple=False, + fixed_decay=False, + rectify=False, + ) + optimizer.step() + + +**Paper**: *AdaBelief Optimizer, adapting stepsizes by the belief in observed gradients* (2020) [https://arxiv.org/abs/2010.07468] + +**Reference Code**: https://github.com/juntang-zhuang/Adabelief-Optimizer + + +AdaBound +-------- + ++------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_AdaBound.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_AdaBound.png | ++------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.AdaBound( + m.parameters(), + lr= 1e-3, + betas= (0.9, 0.999), + final_lr = 0.1, + gamma=1e-3, + eps= 1e-8, + weight_decay=0, + amsbound=False, + ) + optimizer.step() + + +**Paper**: *Adaptive Gradient Methods with Dynamic Bound of Learning Rate* (2019) [https://arxiv.org/abs/1902.09843] + +**Reference Code**: https://github.com/Luolc/AdaBound + +AdaMod +------ +The AdaMod method restricts the adaptive learning rates with adaptive and momental +upper bounds. The dynamic learning rate bounds are based on the exponential +moving averages of the adaptive learning rates themselves, which smooth out +unexpected large learning rates and stabilize the training of deep neural networks. + ++------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_AdaMod.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_AdaMod.png | ++------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.AdaMod( + m.parameters(), + lr= 1e-3, + betas=(0.9, 0.999), + beta3=0.999, + eps=1e-8, + weight_decay=0, + ) + optimizer.step() + +**Paper**: *An Adaptive and Momental Bound Method for Stochastic Learning.* (2019) [https://arxiv.org/abs/1910.12249] + +**Reference Code**: https://github.com/lancopku/AdaMod + + +Adafactor +--------- ++------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_Adafactor.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_Adafactor.png | ++------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.Adafactor( + m.parameters(), + lr= 1e-3, + eps2= (1e-30, 1e-3), + clip_threshold=1.0, + decay_rate=-0.8, + beta1=None, + weight_decay=0.0, + scale_parameter=True, + relative_step=True, + warmup_init=False, + ) + optimizer.step() + +**Paper**: *Adafactor: Adaptive Learning Rates with Sublinear Memory Cost.* (2018) [https://arxiv.org/abs/1804.04235] + +**Reference Code**: https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py + + +Adahessian +---------- ++-------------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_Adahessian.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_Adahessian.png | ++-------------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.Adahessian( + m.parameters(), + lr= 1.0, + betas= (0.9, 0.999), + eps= 1e-4, + weight_decay=0.0, + hessian_power=1.0, + ) + loss_fn(m(input), target).backward(create_graph = True) # create_graph=True is necessary for Hessian calculation + optimizer.step() + + +**Paper**: *ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning* (2020) [https://arxiv.org/abs/2006.00719] + +**Reference Code**: https://github.com/amirgholami/adahessian + + +AdamP +------ +AdamP propose a simple and effective solution: at each iteration of the Adam optimizer +applied on scale-invariant weights (e.g., Conv weights preceding a BN layer), AdamP +removes the radial component (i.e., parallel to the weight vector) from the update vector. +Intuitively, this operation prevents the unnecessary update along the radial direction +that only increases the weight norm without contributing to the loss minimization. + ++------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_AdamP.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_AdamP.png | ++------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.AdamP( + m.parameters(), + lr= 1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=0, + delta = 0.1, + wd_ratio = 0.1 + ) + optimizer.step() + +**Paper**: *Slowing Down the Weight Norm Increase in Momentum-based Optimizers.* (2020) [https://arxiv.org/abs/2006.08217] + +**Reference Code**: https://github.com/clovaai/AdamP + + +AggMo +----- + ++------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_AggMo.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_AggMo.png | ++------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.AggMo( + m.parameters(), + lr= 1e-3, + betas=(0.0, 0.9, 0.99), + weight_decay=0, + ) + optimizer.step() + +**Paper**: *Aggregated Momentum: Stability Through Passive Damping.* (2019) [https://arxiv.org/abs/1804.00325] + +**Reference Code**: https://github.com/AtheMathmo/AggMo + + +Apollo +------ + ++------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_Apollo.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_Apollo.png | ++------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.Apollo( + m.parameters(), + lr= 1e-2, + beta=0.9, + eps=1e-4, + warmup=0, + init_lr=0.01, + weight_decay=0, + ) + optimizer.step() + +**Paper**: *Apollo: An Adaptive Parameter-wise Diagonal Quasi-Newton Method for Nonconvex Stochastic Optimization.* (2020) [https://arxiv.org/abs/2009.13586] + +**Reference Code**: https://github.com/XuezheMax/apollo + + +DiffGrad +-------- +Optimizer based on the difference between the present and the immediate past +gradient, the step size is adjusted for each parameter in such +a way that it should have a larger step size for faster gradient changing +parameters and a lower step size for lower gradient changing parameters. + ++------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_DiffGrad.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_DiffGrad.png | ++------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.DiffGrad( + m.parameters(), + lr= 1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=0, + ) + optimizer.step() + + +**Paper**: *diffGrad: An Optimization Method for Convolutional Neural Networks.* (2019) [https://arxiv.org/abs/1909.11015] + +**Reference Code**: https://github.com/shivram1987/diffGrad + +Lamb +---- + ++--------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_Lamb.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_Lamb.png | ++--------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.Lamb( + m.parameters(), + lr= 1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=0, + ) + optimizer.step() + + +**Paper**: *Large Batch Optimization for Deep Learning: Training BERT in 76 minutes* (2019) [https://arxiv.org/abs/1904.00962] + +**Reference Code**: https://github.com/cybertronai/pytorch-lamb + +Lookahead +--------- + ++-----------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_LookaheadYogi.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_LookaheadYogi.png | ++-----------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + # base optimizer, any other optimizer can be used like Adam or DiffGrad + yogi = optim.Yogi( + m.parameters(), + lr= 1e-2, + betas=(0.9, 0.999), + eps=1e-3, + initial_accumulator=1e-6, + weight_decay=0, + ) + + optimizer = optim.Lookahead(yogi, k=5, alpha=0.5) + optimizer.step() + + +**Paper**: *Lookahead Optimizer: k steps forward, 1 step back* (2019) [https://arxiv.org/abs/1907.08610] + +**Reference Code**: https://github.com/alphadl/lookahead.pytorch + + +MADGRAD +--------- + ++-----------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_MADGRAD.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_MADGRAD.png | ++-----------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.MADGRAD( + m.parameters(), + lr=1e-2, + momentum=0.9, + weight_decay=0, + eps=1e-6, + ) + optimizer.step() + + +**Paper**: *Adaptivity without Compromise: A Momentumized, Adaptive, Dual Averaged Gradient Method for Stochastic Optimization* (2021) [https://arxiv.org/abs/2101.11075] + +**Reference Code**: https://github.com/facebookresearch/madgrad + + +NovoGrad +-------- + ++------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_NovoGrad.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_NovoGrad.png | ++------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.NovoGrad( + m.parameters(), + lr= 1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=0, + grad_averaging=False, + amsgrad=False, + ) + optimizer.step() + + +**Paper**: *Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks* (2019) [https://arxiv.org/abs/1905.11286] + +**Reference Code**: https://github.com/NVIDIA/DeepLearningExamples/ + + +PID +--- + ++-------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_PID.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_PID.png | ++-------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.PID( + m.parameters(), + lr=1e-3, + momentum=0, + dampening=0, + weight_decay=1e-2, + integral=5.0, + derivative=10.0, + ) + optimizer.step() + + +**Paper**: *A PID Controller Approach for Stochastic Optimization of Deep Networks* (2018) [http://www4.comp.polyu.edu.hk/~cslzhang/paper/CVPR18_PID.pdf] + +**Reference Code**: https://github.com/tensorboy/PIDOptimizer + + +QHAdam +------ + ++----------------------------------------------------------------------------------------------------------+------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_QHAdam.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_QHAdam.png | ++----------------------------------------------------------------------------------------------------------+------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.QHAdam( + m.parameters(), + lr= 1e-3, + betas=(0.9, 0.999), + nus=(1.0, 1.0), + weight_decay=0, + decouple_weight_decay=False, + eps=1e-8, + ) + optimizer.step() + + +**Paper**: *Quasi-hyperbolic momentum and Adam for deep learning* (2019) [https://arxiv.org/abs/1810.06801] + +**Reference Code**: https://github.com/facebookresearch/qhoptim + + +QHM +--- + ++-------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_QHM.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_QHM.png | ++-------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.QHM( + m.parameters(), + lr=1e-3, + momentum=0, + nu=0.7, + weight_decay=1e-2, + weight_decay_type='grad', + ) + optimizer.step() + + +**Paper**: *Quasi-hyperbolic momentum and Adam for deep learning* (2019) [https://arxiv.org/abs/1810.06801] + +**Reference Code**: https://github.com/facebookresearch/qhoptim + + +RAdam +----- + ++---------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_RAdam.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_RAdam.png | ++---------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------------------+ + +Deprecated, please use version provided by PyTorch_. + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.RAdam( + m.parameters(), + lr= 1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=0, + ) + optimizer.step() + + +**Paper**: *On the Variance of the Adaptive Learning Rate and Beyond* (2019) [https://arxiv.org/abs/1908.03265] + +**Reference Code**: https://github.com/LiyuanLucasLiu/RAdam + + +Ranger +------ + ++----------------------------------------------------------------------------------------------------------+------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_Ranger.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_Ranger.png | ++----------------------------------------------------------------------------------------------------------+------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.Ranger( + m.parameters(), + lr=1e-3, + alpha=0.5, + k=6, + N_sma_threshhold=5, + betas=(.95, 0.999), + eps=1e-5, + weight_decay=0 + ) + optimizer.step() + + +**Paper**: *New Deep Learning Optimizer, Ranger: Synergistic combination of RAdam + LookAhead for the best of both* (2019) [https://medium.com/@lessw/new-deep-learning-optimizer-ranger-synergistic-combination-of-radam-lookahead-for-the-best-of-2dc83f79a48d] + +**Reference Code**: https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer + + +RangerQH +-------- + ++------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_RangerQH.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_RangerQH.png | ++------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.RangerQH( + m.parameters(), + lr=1e-3, + betas=(0.9, 0.999), + nus=(.7, 1.0), + weight_decay=0.0, + k=6, + alpha=.5, + decouple_weight_decay=False, + eps=1e-8, + ) + optimizer.step() + + +**Paper**: *Quasi-hyperbolic momentum and Adam for deep learning* (2018) [https://arxiv.org/abs/1810.06801] + +**Reference Code**: https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer + + +RangerVA +-------- + ++------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_RangerVA.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_RangerVA.png | ++------------------------------------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.RangerVA( + m.parameters(), + lr=1e-3, + alpha=0.5, + k=6, + n_sma_threshhold=5, + betas=(.95, 0.999), + eps=1e-5, + weight_decay=0, + amsgrad=True, + transformer='softplus', + smooth=50, + grad_transformer='square' + ) + optimizer.step() + + +**Paper**: *Calibrating the Adaptive Learning Rate to Improve Convergence of ADAM* (2019) [https://arxiv.org/abs/1908.00700v2] + +**Reference Code**: https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer + + +SGDP +---- + ++--------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_SGDP.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_SGDP.png | ++--------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.SGDP( + m.parameters(), + lr= 1e-3, + momentum=0, + dampening=0, + weight_decay=1e-2, + nesterov=False, + delta = 0.1, + wd_ratio = 0.1 + ) + optimizer.step() + + +**Paper**: *Slowing Down the Weight Norm Increase in Momentum-based Optimizers.* (2020) [https://arxiv.org/abs/2006.08217] + +**Reference Code**: https://github.com/clovaai/AdamP + + +SGDW +---- + ++--------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_SGDW.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_SGDW.png | ++--------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.SGDW( + m.parameters(), + lr= 1e-3, + momentum=0, + dampening=0, + weight_decay=1e-2, + nesterov=False, + ) + optimizer.step() + + +**Paper**: *SGDR: Stochastic Gradient Descent with Warm Restarts* (2017) [https://arxiv.org/abs/1608.03983] + +**Reference Code**: https://github.com/pytorch/pytorch/pull/22466 + + +SWATS +----- + ++---------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_SWATS.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_SWATS.png | ++---------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.SWATS( + model.parameters(), + lr=1e-1, + betas=(0.9, 0.999), + eps=1e-3, + weight_decay= 0.0, + amsgrad=False, + nesterov=False, + ) + optimizer.step() + + +**Paper**: *Improving Generalization Performance by Switching from Adam to SGD* (2017) [https://arxiv.org/abs/1712.07628] + +**Reference Code**: https://github.com/Mrpatekful/swats + + +Shampoo +------- + ++-----------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_Shampoo.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_Shampoo.png | ++-----------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.Shampoo( + m.parameters(), + lr=1e-1, + momentum=0.0, + weight_decay=0.0, + epsilon=1e-4, + update_freq=1, + ) + optimizer.step() + + +**Paper**: *Shampoo: Preconditioned Stochastic Tensor Optimization* (2018) [https://arxiv.org/abs/1802.09568] + +**Reference Code**: https://github.com/moskomule/shampoo.pytorch + + +Yogi +---- + +Yogi is optimization algorithm based on ADAM with more fine grained effective +learning rate control, and has similar theoretical guarantees on convergence as ADAM. + ++--------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_Yogi.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_Yogi.png | ++--------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------+ + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.Yogi( + m.parameters(), + lr= 1e-2, + betas=(0.9, 0.999), + eps=1e-3, + initial_accumulator=1e-6, + weight_decay=0, + ) + optimizer.step() + + +**Paper**: *Adaptive Methods for Nonconvex Optimization* (2018) [https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization] + +**Reference Code**: https://github.com/4rtemi5/Yogi-Optimizer_Keras + + +Adam (PyTorch built-in) +----------------------- + ++---------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_Adam.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_Adam.png | ++---------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------+ + +SGD (PyTorch built-in) +---------------------- + ++--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------+ +| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_SGD.png | .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_SGD.png | ++--------------------------------------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------+ + .. _Python: https://www.python.org .. _PyTorch: https://github.com/pytorch/pytorch +.. _Rastrigin: https://en.wikipedia.org/wiki/Rastrigin_function +.. _Rosenbrock: https://en.wikipedia.org/wiki/Rosenbrock_function +.. _benchmark: https://en.wikipedia.org/wiki/Test_functions_for_optimization +.. _optim: https://pytorch.org/docs/stable/optim.html diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 00000000..d4bb2cbb --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = . +BUILDDIR = _build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/api.rst b/docs/api.rst new file mode 100644 index 00000000..ce4bd5a2 --- /dev/null +++ b/docs/api.rst @@ -0,0 +1,146 @@ +Available Optimizers +==================== + +.. _AccSGD: + +AccSGD +------ + +.. autoclass:: torch_optimizer.AccSGD + :members: + +.. _AdaBound: + +AdaBound +-------- + +.. autoclass:: torch_optimizer.AdaBound + :members: + +.. _AdaMod: + +AdaMod +------ + +.. autoclass:: torch_optimizer.AdaMod + :members: + +.. _Adafactor: + +Adafactor +--------- + +.. autoclass:: torch_optimizer.Adafactor + :members: + +.. _AdamP: + +AdamP +------ + +.. autoclass:: torch_optimizer.AdamP + :members: + +.. _AggMo: + +AggMo +----- + +.. autoclass:: torch_optimizer.AggMo + :members: + +.. _DiffGrad: + +DiffGrad +-------- + +.. autoclass:: torch_optimizer.DiffGrad + :members: + +.. _Lamb: + +Lamb +---- + +.. autoclass:: torch_optimizer.Lamb + :members: + +.. _NovoGrad: + +NovoGrad +-------- + +.. autoclass:: torch_optimizer.NovoGrad + :members: + +.. _PID: + +PID +--- + +.. autoclass:: torch_optimizer.PID + :members: + +.. _QHAdam: + +QHAdam +------ + +.. autoclass:: torch_optimizer.QHAdam + :members: + +.. _QHM: + +QHM +--- + +.. autoclass:: torch_optimizer.QHM + :members: + +.. _RAdam: + +RAdam +----- + +.. autoclass:: torch_optimizer.RAdam + :members: + +.. _SGDP: + +SGDP +---- + +.. autoclass:: torch_optimizer.SGDP + :members: + +.. _SGDW: + +SGDW +---- + +.. autoclass:: torch_optimizer.SGDW + :members: + +.. _Shampoo: + +Shampoo +------- + +.. autoclass:: torch_optimizer.Shampoo + :members: + +.. _SWATS: + +SWATS +----- + +.. autoclass:: torch_optimizer.SWATS + :members: + +.. _Yogi: + +Yogi +---- + +.. autoclass:: torch_optimizer.Yogi + :members: diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 00000000..7e06a3fd --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,80 @@ +# Configuration file for the Sphinx documentation builder. +# +# This file only contains a selection of the most common options. For a full +# list see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +# -- Path setup -------------------------------------------------------------- + +# If extensions (or modules to document with autodoc) are in another directory, +# add these directories to sys.path here. If the directory is relative to the +# documentation root, use os.path.abspath to make it absolute, like shown here. +# +# import os +# import sys +# sys.path.insert(0, os.path.abspath('.')) + + +# -- Project information ----------------------------------------------------- + +project = 'pytorch-optimizer' +copyright = '2020, Nikolai Novik' +author = 'Nikolai Novik' + + +# -- General configuration --------------------------------------------------- + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom +# ones. + +# Sphinx extension modules +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.napoleon", + "sphinx_autodoc_typehints", + "sphinx.ext.doctest", + "sphinx.ext.todo", + "sphinx.ext.coverage", + "sphinx.ext.mathjax", + "sphinx.ext.ifconfig", + "sphinx.ext.viewcode", + "sphinx.ext.intersphinx", +] + +# Add any paths that contain templates here, relative to this directory. +templates_path = ['_templates'] + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +# This pattern also affects html_static_path and html_extra_path. +exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] + +# Configuration for intersphinx: refer to the Python standard library and PyTorch +intersphinx_mapping = { + "python": ("https://docs.python.org/3", None), + "pytorch": ("https://pytorch.org/docs/stable", None), +} + + +# -- Options for HTML output ------------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +# +html_theme = 'alabaster' + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +html_static_path = ['_static'] + +desc = 'collection of optimizers for PyTorch' +html_theme_options = { + 'description': desc, + 'github_user': 'jettify', + 'github_repo': 'pytorch-optimizer', + 'github_button': True, + 'github_type': 'star', + 'github_banner': True, +} diff --git a/docs/contributing.rst b/docs/contributing.rst new file mode 100644 index 00000000..e582053e --- /dev/null +++ b/docs/contributing.rst @@ -0,0 +1 @@ +.. include:: ../CONTRIBUTING.rst diff --git a/docs/examples.rst b/docs/examples.rst new file mode 100644 index 00000000..c210fb11 --- /dev/null +++ b/docs/examples.rst @@ -0,0 +1,17 @@ +Examples of pytorch-optimizer usage +=================================== + +Below is a list of examples from `pytorch-optimizer/examples +`_ + +Every example is a correct tiny python program. + +.. _pytorch-optimizer-examples-simple: + + +Basic Usage +----------- + +Simple example that shows how to use library with MNIST dataset. + +.. literalinclude:: ../examples/mnist.py diff --git a/docs/index.rst b/docs/index.rst new file mode 100644 index 00000000..53624517 --- /dev/null +++ b/docs/index.rst @@ -0,0 +1,137 @@ +.. pytorch-optimizer documentation master file, created by + sphinx-quickstart on Thu Feb 13 21:14:16 2020. + You can adapt this file completely to your liking, but it should at least + contain the root `toctree` directive. + +Welcome to pytorch-optimizer's documentation! +============================================= + +**torch-optimizer** -- collection of optimizers for PyTorch_. + +Simple example +-------------- + +.. code:: python + + import torch_optimizer as optim + + # model = ... + optimizer = optim.DiffGrad(model.parameters(), lr=0.001) + optimizer.step() + + +Installation +------------ +Installation process is simple, just:: + + $ pip install torch_optimizer + + +Citation +-------- +Please cite original authors of optimization algorithms. If you like this +package:: + + @software{Novik_torchoptimizers, + title = {{torch-optimizer -- collection of optimization algorithms for PyTorch.}}, + author = {Novik, Mykola}, + year = 2020, + month = 1, + version = {1.0.1} + } + +Or use github feature: "cite this repository" button. + + +Supported Optimizers +==================== + ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`AccSGD` | https://arxiv.org/abs/1803.05591 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`AdaBound` | https://arxiv.org/abs/1902.09843 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`AdaMod` | https://arxiv.org/abs/1910.12249 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`Adafactor`| https://arxiv.org/abs/1804.04235 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`AdamP` | https://arxiv.org/abs/1804.00325 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`AggMo` | https://arxiv.org/abs/2006.08217 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`DiffGrad` | https://arxiv.org/abs/1909.11015 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`Lamb` | https://arxiv.org/abs/1904.00962 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`NovoGrad` | https://arxiv.org/abs/1905.11286 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`PID` | https://www4.comp.polyu.edu.hk/~cslzhang/paper/CVPR18_PID.pdf | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`QHAdam` | https://arxiv.org/abs/1810.06801 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`QHM` | https://arxiv.org/abs/1810.06801 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`RAdam` | https://arxiv.org/abs/1908.03265 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`Ranger` | https://arxiv.org/abs/1908.00700v2 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`RangerQH` | https://arxiv.org/abs/1908.00700v2 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`RangerVA` | https://arxiv.org/abs/1908.00700v2 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`SGDP` | https://arxiv.org/abs/2006.08217 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`SGDW` | https://arxiv.org/abs/1608.03983 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`Shampoo` | https://arxiv.org/abs/1802.09568 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`SWATS` | https://arxiv.org/abs/1712.07628 | ++-----------------+-------------------------------------------------------------------------------+ +| | | +| :ref:`Yogi` | https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization | ++-----------------+-------------------------------------------------------------------------------+ + +.. toctree:: + :maxdepth: 2 + :caption: Contents: + +Contents +-------- + +.. toctree:: + :maxdepth: 2 + + api + examples + contributing + + +Indices and tables +================== + +* :ref:`genindex` +* :ref:`modindex` +* :ref:`search` + +.. _Python: https://www.python.org +.. _PyTorch: https://github.com/pytorch/pytorch diff --git a/docs/make.bat b/docs/make.bat new file mode 100644 index 00000000..2119f510 --- /dev/null +++ b/docs/make.bat @@ -0,0 +1,35 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=. +set BUILDDIR=_build + +if "%1" == "" goto help + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.http://sphinx-doc.org/ + exit /b 1 +) + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + +:end +popd diff --git a/docs/rastrigin_A2GradExp.png b/docs/rastrigin_A2GradExp.png new file mode 100644 index 00000000..851542e1 Binary files /dev/null and b/docs/rastrigin_A2GradExp.png differ diff --git a/docs/rastrigin_A2GradInc.png b/docs/rastrigin_A2GradInc.png new file mode 100644 index 00000000..cbee4a24 Binary files /dev/null and b/docs/rastrigin_A2GradInc.png differ diff --git a/docs/rastrigin_A2GradUni.png b/docs/rastrigin_A2GradUni.png new file mode 100644 index 00000000..d26dce8a Binary files /dev/null and b/docs/rastrigin_A2GradUni.png differ diff --git a/docs/rastrigin_AccSGD.png b/docs/rastrigin_AccSGD.png new file mode 100644 index 00000000..086497da Binary files /dev/null and b/docs/rastrigin_AccSGD.png differ diff --git a/docs/rastrigin_AdaBelief.png b/docs/rastrigin_AdaBelief.png new file mode 100644 index 00000000..dc68ec33 Binary files /dev/null and b/docs/rastrigin_AdaBelief.png differ diff --git a/docs/rastrigin_AdaBound.png b/docs/rastrigin_AdaBound.png new file mode 100644 index 00000000..f6f1f0f4 Binary files /dev/null and b/docs/rastrigin_AdaBound.png differ diff --git a/docs/rastrigin_AdaMod.png b/docs/rastrigin_AdaMod.png new file mode 100644 index 00000000..ba510e06 Binary files /dev/null and b/docs/rastrigin_AdaMod.png differ diff --git a/docs/rastrigin_Adafactor.png b/docs/rastrigin_Adafactor.png new file mode 100644 index 00000000..d34d8ffc Binary files /dev/null and b/docs/rastrigin_Adafactor.png differ diff --git a/docs/rastrigin_Adahessian.png b/docs/rastrigin_Adahessian.png new file mode 100644 index 00000000..fa4442c9 Binary files /dev/null and b/docs/rastrigin_Adahessian.png differ diff --git a/docs/rastrigin_Adam.png b/docs/rastrigin_Adam.png new file mode 100644 index 00000000..9f334c61 Binary files /dev/null and b/docs/rastrigin_Adam.png differ diff --git a/docs/rastrigin_AdamP.png b/docs/rastrigin_AdamP.png new file mode 100644 index 00000000..b98200b8 Binary files /dev/null and b/docs/rastrigin_AdamP.png differ diff --git a/docs/rastrigin_AggMo.png b/docs/rastrigin_AggMo.png new file mode 100644 index 00000000..1767de46 Binary files /dev/null and b/docs/rastrigin_AggMo.png differ diff --git a/docs/rastrigin_Apollo.png b/docs/rastrigin_Apollo.png new file mode 100644 index 00000000..4868c566 Binary files /dev/null and b/docs/rastrigin_Apollo.png differ diff --git a/docs/rastrigin_DiffGrad.png b/docs/rastrigin_DiffGrad.png new file mode 100644 index 00000000..3bce8e1b Binary files /dev/null and b/docs/rastrigin_DiffGrad.png differ diff --git a/docs/rastrigin_Lamb.png b/docs/rastrigin_Lamb.png new file mode 100644 index 00000000..ce862558 Binary files /dev/null and b/docs/rastrigin_Lamb.png differ diff --git a/docs/rastrigin_LookaheadYogi.png b/docs/rastrigin_LookaheadYogi.png new file mode 100644 index 00000000..5880a543 Binary files /dev/null and b/docs/rastrigin_LookaheadYogi.png differ diff --git a/docs/rastrigin_MADGRAD.png b/docs/rastrigin_MADGRAD.png new file mode 100644 index 00000000..9e46428f Binary files /dev/null and b/docs/rastrigin_MADGRAD.png differ diff --git a/docs/rastrigin_NovoGrad.png b/docs/rastrigin_NovoGrad.png new file mode 100644 index 00000000..eb6d0c65 Binary files /dev/null and b/docs/rastrigin_NovoGrad.png differ diff --git a/docs/rastrigin_PID.png b/docs/rastrigin_PID.png new file mode 100644 index 00000000..98fc481a Binary files /dev/null and b/docs/rastrigin_PID.png differ diff --git a/docs/rastrigin_QHAdam.png b/docs/rastrigin_QHAdam.png new file mode 100644 index 00000000..03e0c887 Binary files /dev/null and b/docs/rastrigin_QHAdam.png differ diff --git a/docs/rastrigin_QHM.png b/docs/rastrigin_QHM.png new file mode 100644 index 00000000..4331eb3f Binary files /dev/null and b/docs/rastrigin_QHM.png differ diff --git a/docs/rastrigin_RAdam.png b/docs/rastrigin_RAdam.png new file mode 100644 index 00000000..d6856e36 Binary files /dev/null and b/docs/rastrigin_RAdam.png differ diff --git a/docs/rastrigin_Ranger.png b/docs/rastrigin_Ranger.png new file mode 100644 index 00000000..f2188caf Binary files /dev/null and b/docs/rastrigin_Ranger.png differ diff --git a/docs/rastrigin_RangerQH.png b/docs/rastrigin_RangerQH.png new file mode 100644 index 00000000..16d1e247 Binary files /dev/null and b/docs/rastrigin_RangerQH.png differ diff --git a/docs/rastrigin_RangerVA.png b/docs/rastrigin_RangerVA.png new file mode 100644 index 00000000..fb2f4904 Binary files /dev/null and b/docs/rastrigin_RangerVA.png differ diff --git a/docs/rastrigin_SGD.png b/docs/rastrigin_SGD.png new file mode 100644 index 00000000..1e07df3f Binary files /dev/null and b/docs/rastrigin_SGD.png differ diff --git a/docs/rastrigin_SGDP.png b/docs/rastrigin_SGDP.png new file mode 100644 index 00000000..f1d5a299 Binary files /dev/null and b/docs/rastrigin_SGDP.png differ diff --git a/docs/rastrigin_SGDW.png b/docs/rastrigin_SGDW.png new file mode 100644 index 00000000..e45bebeb Binary files /dev/null and b/docs/rastrigin_SGDW.png differ diff --git a/docs/rastrigin_SWATS.png b/docs/rastrigin_SWATS.png new file mode 100644 index 00000000..458df351 Binary files /dev/null and b/docs/rastrigin_SWATS.png differ diff --git a/docs/rastrigin_Shampoo.png b/docs/rastrigin_Shampoo.png new file mode 100644 index 00000000..47e7ff2a Binary files /dev/null and b/docs/rastrigin_Shampoo.png differ diff --git a/docs/rastrigin_Yogi.png b/docs/rastrigin_Yogi.png new file mode 100644 index 00000000..a31f3d26 Binary files /dev/null and b/docs/rastrigin_Yogi.png differ diff --git a/docs/rosenbrock_A2GradExp.png b/docs/rosenbrock_A2GradExp.png new file mode 100644 index 00000000..e7b165c6 Binary files /dev/null and b/docs/rosenbrock_A2GradExp.png differ diff --git a/docs/rosenbrock_A2GradInc.png b/docs/rosenbrock_A2GradInc.png new file mode 100644 index 00000000..25be224b Binary files /dev/null and b/docs/rosenbrock_A2GradInc.png differ diff --git a/docs/rosenbrock_A2GradUni.png b/docs/rosenbrock_A2GradUni.png new file mode 100644 index 00000000..0314ad11 Binary files /dev/null and b/docs/rosenbrock_A2GradUni.png differ diff --git a/docs/rosenbrock_AccSGD.png b/docs/rosenbrock_AccSGD.png new file mode 100644 index 00000000..364fce3f Binary files /dev/null and b/docs/rosenbrock_AccSGD.png differ diff --git a/docs/rosenbrock_AdaBelief.png b/docs/rosenbrock_AdaBelief.png new file mode 100644 index 00000000..cfff2c87 Binary files /dev/null and b/docs/rosenbrock_AdaBelief.png differ diff --git a/docs/rosenbrock_AdaBound.png b/docs/rosenbrock_AdaBound.png new file mode 100644 index 00000000..0805c78c Binary files /dev/null and b/docs/rosenbrock_AdaBound.png differ diff --git a/docs/rosenbrock_AdaMod.png b/docs/rosenbrock_AdaMod.png new file mode 100644 index 00000000..1dc46076 Binary files /dev/null and b/docs/rosenbrock_AdaMod.png differ diff --git a/docs/rosenbrock_Adafactor.png b/docs/rosenbrock_Adafactor.png new file mode 100644 index 00000000..0b57b005 Binary files /dev/null and b/docs/rosenbrock_Adafactor.png differ diff --git a/docs/rosenbrock_Adahessian.png b/docs/rosenbrock_Adahessian.png new file mode 100644 index 00000000..2f26e2e7 Binary files /dev/null and b/docs/rosenbrock_Adahessian.png differ diff --git a/docs/rosenbrock_Adam.png b/docs/rosenbrock_Adam.png new file mode 100644 index 00000000..4695095e Binary files /dev/null and b/docs/rosenbrock_Adam.png differ diff --git a/docs/rosenbrock_AdamP.png b/docs/rosenbrock_AdamP.png new file mode 100644 index 00000000..a1b42ddc Binary files /dev/null and b/docs/rosenbrock_AdamP.png differ diff --git a/docs/rosenbrock_AggMo.png b/docs/rosenbrock_AggMo.png new file mode 100644 index 00000000..15766071 Binary files /dev/null and b/docs/rosenbrock_AggMo.png differ diff --git a/docs/rosenbrock_Apollo.png b/docs/rosenbrock_Apollo.png new file mode 100644 index 00000000..5d6190c1 Binary files /dev/null and b/docs/rosenbrock_Apollo.png differ diff --git a/docs/rosenbrock_DiffGrad.png b/docs/rosenbrock_DiffGrad.png new file mode 100644 index 00000000..33554aa5 Binary files /dev/null and b/docs/rosenbrock_DiffGrad.png differ diff --git a/docs/rosenbrock_Lamb.png b/docs/rosenbrock_Lamb.png new file mode 100644 index 00000000..9ccf1e16 Binary files /dev/null and b/docs/rosenbrock_Lamb.png differ diff --git a/docs/rosenbrock_LookaheadYogi.png b/docs/rosenbrock_LookaheadYogi.png new file mode 100644 index 00000000..8b8f3903 Binary files /dev/null and b/docs/rosenbrock_LookaheadYogi.png differ diff --git a/docs/rosenbrock_MADGRAD.png b/docs/rosenbrock_MADGRAD.png new file mode 100644 index 00000000..67ae7dd4 Binary files /dev/null and b/docs/rosenbrock_MADGRAD.png differ diff --git a/docs/rosenbrock_NovoGrad.png b/docs/rosenbrock_NovoGrad.png new file mode 100644 index 00000000..35803867 Binary files /dev/null and b/docs/rosenbrock_NovoGrad.png differ diff --git a/docs/rosenbrock_PID.png b/docs/rosenbrock_PID.png new file mode 100644 index 00000000..085c44b9 Binary files /dev/null and b/docs/rosenbrock_PID.png differ diff --git a/docs/rosenbrock_QHAdam.png b/docs/rosenbrock_QHAdam.png new file mode 100644 index 00000000..9bec4329 Binary files /dev/null and b/docs/rosenbrock_QHAdam.png differ diff --git a/docs/rosenbrock_QHM.png b/docs/rosenbrock_QHM.png new file mode 100644 index 00000000..14f258e2 Binary files /dev/null and b/docs/rosenbrock_QHM.png differ diff --git a/docs/rosenbrock_RAdam.png b/docs/rosenbrock_RAdam.png new file mode 100644 index 00000000..dc3078eb Binary files /dev/null and b/docs/rosenbrock_RAdam.png differ diff --git a/docs/rosenbrock_Ranger.png b/docs/rosenbrock_Ranger.png new file mode 100644 index 00000000..11043daa Binary files /dev/null and b/docs/rosenbrock_Ranger.png differ diff --git a/docs/rosenbrock_RangerQH.png b/docs/rosenbrock_RangerQH.png new file mode 100644 index 00000000..0f1b5a27 Binary files /dev/null and b/docs/rosenbrock_RangerQH.png differ diff --git a/docs/rosenbrock_RangerVA.png b/docs/rosenbrock_RangerVA.png new file mode 100644 index 00000000..759e492e Binary files /dev/null and b/docs/rosenbrock_RangerVA.png differ diff --git a/docs/rosenbrock_SGD.png b/docs/rosenbrock_SGD.png new file mode 100644 index 00000000..d86ac249 Binary files /dev/null and b/docs/rosenbrock_SGD.png differ diff --git a/docs/rosenbrock_SGDP.png b/docs/rosenbrock_SGDP.png new file mode 100644 index 00000000..f9bb39e3 Binary files /dev/null and b/docs/rosenbrock_SGDP.png differ diff --git a/docs/rosenbrock_SGDW.png b/docs/rosenbrock_SGDW.png new file mode 100644 index 00000000..5680fe18 Binary files /dev/null and b/docs/rosenbrock_SGDW.png differ diff --git a/docs/rosenbrock_SWATS.png b/docs/rosenbrock_SWATS.png new file mode 100644 index 00000000..b2d51fb1 Binary files /dev/null and b/docs/rosenbrock_SWATS.png differ diff --git a/docs/rosenbrock_Shampoo.png b/docs/rosenbrock_Shampoo.png new file mode 100644 index 00000000..decbb382 Binary files /dev/null and b/docs/rosenbrock_Shampoo.png differ diff --git a/docs/rosenbrock_Yogi.png b/docs/rosenbrock_Yogi.png new file mode 100644 index 00000000..fd30e3a4 Binary files /dev/null and b/docs/rosenbrock_Yogi.png differ diff --git a/examples/mnist.py b/examples/mnist.py index 0d591ef9..2d860254 100644 --- a/examples/mnist.py +++ b/examples/mnist.py @@ -1,12 +1,11 @@ import torch import torch.nn as nn import torch.nn.functional as F -import torch_optimizer as optim - -from torchvision import datasets, transforms, utils from torch.optim.lr_scheduler import StepLR from torch.utils.tensorboard import SummaryWriter -from dataclasses import dataclass +from torchvision import datasets, transforms, utils + +import torch_optimizer as optim class Net(nn.Module): @@ -46,9 +45,9 @@ def train(conf, model, device, train_loader, optimizer, epoch, writer): if batch_idx % conf.log_interval == 0: loss = loss.item() idx = batch_idx + epoch * (len(train_loader)) - writer.add_scalar('Loss/train', loss, idx) + writer.add_scalar("Loss/train", loss, idx) print( - 'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( + "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format( epoch, batch_idx * len(data), len(train_loader.dataset), @@ -67,13 +66,13 @@ def test(conf, model, device, test_loader, epoch, writer): data, target = data.to(device), target.to(device) output = model(data) # sum up batch loss - test_loss += F.nll_loss(output, target, reduction='sum').item() + test_loss += F.nll_loss(output, target, reduction="sum").item() # get the index of the max log-probability pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader.dataset) - fmt = '\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n' + fmt = "\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n" print( fmt.format( test_loss, @@ -83,15 +82,15 @@ def test(conf, model, device, test_loader, epoch, writer): ) ) - writer.add_scalar('Accuracy', correct, epoch) - writer.add_scalar('Loss/test', test_loss, epoch) + writer.add_scalar("Accuracy", correct, epoch) + writer.add_scalar("Loss/test", test_loss, epoch) def prepare_loaders(conf, use_cuda=False): - kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {} + kwargs = {"num_workers": 1, "pin_memory": True} if use_cuda else {} train_loader = torch.utils.data.DataLoader( datasets.MNIST( - '../data', + "../data", train=True, download=True, transform=transforms.Compose( @@ -103,12 +102,12 @@ def prepare_loaders(conf, use_cuda=False): ), batch_size=conf.batch_size, shuffle=True, - **kwargs + **kwargs, ) test_loader = torch.utils.data.DataLoader( datasets.MNIST( - '../data', + "../data", train=False, transform=transforms.Compose( [ @@ -119,33 +118,42 @@ def prepare_loaders(conf, use_cuda=False): ), batch_size=conf.test_batch_size, shuffle=True, - **kwargs + **kwargs, ) return train_loader, test_loader -@dataclass class Config: - batch_size: int = 64 - test_batch_size: int = 1000 - epochs: int = 15 - lr: float = 0.01 - gamma: float = 0.7 - no_cuda: bool = True - seed: int = 42 - log_interval: int = 10 - optimizer: str = 'powersign' + def __init__( + self, + batch_size: int = 64, + test_batch_size: int = 1000, + epochs: int = 15, + lr: float = 0.01, + gamma: float = 0.7, + no_cuda: bool = True, + seed: int = 42, + log_interval: int = 10, + ): + self.batch_size = batch_size + self.test_batch_size = test_batch_size + self.epochs = epochs + self.lr = lr + self.gamma = gamma + self.no_cuda = no_cuda + self.seed = seed + self.log_interval = log_interval def main(): conf = Config() - log_dir = 'runs/mnist_custom_optim' - print('Tensorboard: tensorboard --logdir={}'.format(log_dir)) + log_dir = "runs/mnist_custom_optim" + print("Tensorboard: tensorboard --logdir={}".format(log_dir)) with SummaryWriter(log_dir) as writer: use_cuda = not conf.no_cuda and torch.cuda.is_available() torch.manual_seed(conf.seed) - device = torch.device('cuda' if use_cuda else 'cpu') + device = torch.device("cuda" if use_cuda else "cpu") train_loader, test_loader = prepare_loaders(conf, use_cuda) model = Net().to(device) @@ -153,12 +161,10 @@ def main(): # create grid of images and write to tensorboard images, labels = next(iter(train_loader)) img_grid = utils.make_grid(images) - writer.add_image('mnist_images', img_grid) - # visualize NN computation graph - writer.add_graph(model, images) + writer.add_image("mnist_images", img_grid) # custom optimizer from torch_optimizer package - optimizer = optim.PowerSign(model.parameters(), lr=conf.lr) + optimizer = optim.DiffGrad(model.parameters(), lr=conf.lr) scheduler = StepLR(optimizer, step_size=1, gamma=conf.gamma) for epoch in range(1, conf.epochs + 1): @@ -167,8 +173,8 @@ def main(): scheduler.step() for name, param in model.named_parameters(): writer.add_histogram(name, param, epoch) - writer.add_histogram(f'{name}.grad', param.grad, epoch) + writer.add_histogram("{}.grad".format(name), param.grad, epoch) -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/examples/requirements-examples.txt b/examples/requirements-examples.txt new file mode 100644 index 00000000..4c1f89fb --- /dev/null +++ b/examples/requirements-examples.txt @@ -0,0 +1,4 @@ +torch==1.13.1 +hyperopt==0.2.5 +torchvision==0.11.1 +matplotlib==3.4.3 diff --git a/examples/viz_optimizers.py b/examples/viz_optimizers.py new file mode 100644 index 00000000..5c1cba10 --- /dev/null +++ b/examples/viz_optimizers.py @@ -0,0 +1,213 @@ +import math + +import matplotlib.pyplot as plt +import numpy as np +import torch +from hyperopt import fmin, hp, tpe + +import torch_optimizer as optim + +plt.style.use("seaborn-white") + + +def rosenbrock(tensor): + # https://en.wikipedia.org/wiki/Test_functions_for_optimization + x, y = tensor + return (1 - x) ** 2 + 100 * (y - x**2) ** 2 + + +def rastrigin(tensor, lib=torch): + # https://en.wikipedia.org/wiki/Test_functions_for_optimization + x, y = tensor + A = 10 + f = ( + A * 2 + + (x**2 - A * lib.cos(x * math.pi * 2)) + + (y**2 - A * lib.cos(y * math.pi * 2)) + ) + return f + + +def execute_steps( + func, initial_state, optimizer_class, optimizer_config, num_iter=500 +): + x = torch.Tensor(initial_state).requires_grad_(True) + optimizer = optimizer_class([x], **optimizer_config) + steps = [] + steps = np.zeros((2, num_iter + 1)) + steps[:, 0] = np.array(initial_state) + for i in range(1, num_iter + 1): + optimizer.zero_grad() + f = func(x) + f.backward(create_graph=True, retain_graph=True) + torch.nn.utils.clip_grad_norm_(x, 1.0) + optimizer.step() + steps[:, i] = x.detach().numpy() + return steps + + +def objective_rastrigin(params): + lr = params["lr"] + optimizer_class = params["optimizer_class"] + initial_state = (-2.0, 3.5) + minimum = (0, 0) + optimizer_config = dict(lr=lr) + num_iter = 100 + steps = execute_steps( + rastrigin, initial_state, optimizer_class, optimizer_config, num_iter + ) + return (steps[0][-1] - minimum[0]) ** 2 + (steps[1][-1] - minimum[1]) ** 2 + + +def objective_rosenbrok(params): + lr = params["lr"] + optimizer_class = params["optimizer_class"] + minimum = (1.0, 1.0) + initial_state = (-2.0, 2.0) + optimizer_config = dict(lr=lr) + num_iter = 100 + steps = execute_steps( + rosenbrock, initial_state, optimizer_class, optimizer_config, num_iter + ) + return (steps[0][-1] - minimum[0]) ** 2 + (steps[1][-1] - minimum[1]) ** 2 + + +def plot_rastrigin(grad_iter, optimizer_name, lr): + x = np.linspace(-4.5, 4.5, 250) + y = np.linspace(-4.5, 4.5, 250) + minimum = (0, 0) + + X, Y = np.meshgrid(x, y) + Z = rastrigin([X, Y], lib=np) + + iter_x, iter_y = grad_iter[0, :], grad_iter[1, :] + + fig = plt.figure(figsize=(8, 8)) + + ax = fig.add_subplot(1, 1, 1) + ax.contour(X, Y, Z, 20, cmap="jet") + ax.plot(iter_x, iter_y, color="r", marker="x") + ax.set_title( + "Rastrigin func: {} with " + "{} iterations, lr={:.6}".format(optimizer_name, len(iter_x), lr) + ) + plt.plot(*minimum, "gD") + plt.plot(iter_x[-1], iter_y[-1], "rD") + plt.savefig("docs/rastrigin_{}.png".format(optimizer_name)) + + +def plot_rosenbrok(grad_iter, optimizer_name, lr): + x = np.linspace(-2, 2, 250) + y = np.linspace(-1, 3, 250) + minimum = (1.0, 1.0) + + X, Y = np.meshgrid(x, y) + Z = rosenbrock([X, Y]) + + iter_x, iter_y = grad_iter[0, :], grad_iter[1, :] + + fig = plt.figure(figsize=(8, 8)) + + ax = fig.add_subplot(1, 1, 1) + ax.contour(X, Y, Z, 90, cmap="jet") + ax.plot(iter_x, iter_y, color="r", marker="x") + + ax.set_title( + "Rosenbrock func: {} with {} " + "iterations, lr={:.6}".format(optimizer_name, len(iter_x), lr) + ) + plt.plot(*minimum, "gD") + plt.plot(iter_x[-1], iter_y[-1], "rD") + plt.savefig("docs/rosenbrock_{}.png".format(optimizer_name)) + + +def execute_experiments( + optimizers, objective, func, plot_func, initial_state, seed=1 +): + seed = seed + for item in optimizers: + optimizer_class, lr_low, lr_hi = item + space = { + "optimizer_class": hp.choice("optimizer_class", [optimizer_class]), + "lr": hp.loguniform("lr", lr_low, lr_hi), + } + best = fmin( + fn=objective, + space=space, + algo=tpe.suggest, + max_evals=200, + rstate=np.random.RandomState(seed), + ) + print(best["lr"], optimizer_class) + + steps = execute_steps( + func, + initial_state, + optimizer_class, + {"lr": best["lr"]}, + num_iter=500, + ) + plot_func(steps, optimizer_class.__name__, best["lr"]) + + +def LookaheadYogi(*a, **kw): + base = optim.Yogi(*a, **kw) + return optim.Lookahead(base) + + +if __name__ == "__main__": + # python examples/viz_optimizers.py + + # Each optimizer has tweaked search space to produce better plots and + # help to converge on better lr faster. + optimizers = [ + # baselines + (torch.optim.Adam, -8, 0.5), + (torch.optim.SGD, -8, -1.0), + # Adam based + (optim.AdaBound, -8, 0.3), + (optim.Adahessian, -1, 8), + (optim.AdaMod, -8, 0.2), + (optim.AdamP, -8, 0.2), + (optim.DiffGrad, -8, 0.4), + (optim.Lamb, -8, -2.9), + (optim.MADGRAD, -8, 0.5), + (optim.NovoGrad, -8, -1.7), + (optim.RAdam, -8, 0.5), + (optim.Yogi, -8, 0.1), + # SGD/Momentum based + (optim.AccSGD, -8, -1.4), + (optim.SGDW, -8, -1.5), + (optim.SGDP, -8, -1.5), + (optim.PID, -8, -1.0), + (optim.QHM, -6, -0.2), + (optim.QHAdam, -8, 0.1), + (optim.Ranger, -8, 0.1), + (optim.RangerQH, -8, 0.1), + (optim.RangerVA, -8, 0.1), + (optim.Shampoo, -8, 0.1), + (LookaheadYogi, -8, 0.1), + (optim.AggMo, -8, -1.5), + (optim.SWATS, -8, -1.5), + (optim.Adafactor, -8, 0.5), + (optim.A2GradUni, -8, 0.1), + (optim.A2GradInc, -8, 0.1), + (optim.A2GradExp, -8, 0.1), + (optim.AdaBelief, -8, 0.1), + (optim.Apollo, -8, 0.1), + ] + execute_experiments( + optimizers, + objective_rastrigin, + rastrigin, + plot_rastrigin, + (-2.0, 3.5), + ) + + execute_experiments( + optimizers, + objective_rosenbrok, + rosenbrock, + plot_rosenbrok, + (-2.0, 2.0), + ) diff --git a/requirements-dev.txt b/requirements-dev.txt index 849d19d2..1b18b8f3 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,13 +1,18 @@ -e . -bandit==1.6.2 -black==19.3b0 -flake8-bugbear==19.8.0 -flake8-quotes==2.1.0 -flake8==3.7.8 -ipdb==0.12.2 -ipython==7.8.0 -mypy==0.730 -pyroma==2.5 -pytest-cov==2.8.1 -pytest==5.2.1 -torch==1.1.0 +bandit==1.7.0 +black==23.3.0 +flake8-bugbear==21.9.2 +flake8==4.0.1 +ipdb==0.13.9 +isort==5.9.3 +mypy==0.910 +numpy==1.23.2 +pyroma==3.2 +pytest-cov==3.0.0 +pytest==6.2.5 +pytorch_ranger==0.1.1 +sphinx-autodoc-typehints==1.12.0 +sphinx==4.2.0 +torch==1.13.1 +twine==3.4.2 +wheel==0.38.1 diff --git a/setup.py b/setup.py index 7b6e6203..a162df47 100644 --- a/setup.py +++ b/setup.py @@ -1,64 +1,93 @@ import os import re -import sys -from setuptools import setup, find_packages +from setuptools import find_packages, setup -install_requires = ['torch>=1.1.0'] - -PY36 = (3, 6, 0) - - -if sys.version_info < PY36: - raise RuntimeError('torch-optimizer requires Python 3.6.0+') +install_requires = [ + "torch>=1.5.0", + "pytorch_ranger>=0.1.1", +] -def read(f): - return open(os.path.join(os.path.dirname(__file__), f)).read().strip() +def _read(f): + with open(os.path.join(os.path.dirname(__file__), f)) as f_: + return f_.read().strip() -def read_version(): - regexp = re.compile(r"^__version__\W*=\W*'([\d.abrc]+)'") +def _read_version(): + regexp = re.compile(r'^__version__\W*=\W*"([\d.abrc]+)"') init_py = os.path.join( - os.path.dirname(__file__), 'torch_optimizer', '__init__.py' + os.path.dirname(__file__), "torch_optimizer", "__init__.py" ) with open(init_py) as f: for line in f: match = regexp.match(line) if match is not None: return match.group(1) - else: - raise RuntimeError( - 'Cannot find version in torch_optimizer/__init__.py' - ) + raise RuntimeError( + "Cannot find version in torch_optimizer/__init__.py" + ) classifiers = [ - 'License :: OSI Approved :: Apache Software License', - 'Intended Audience :: Developers', - 'Programming Language :: Python :: 3', - 'Programming Language :: Python :: 3.6', - 'Programming Language :: Python :: 3.7', - 'Operating System :: OS Independent', - 'Development Status :: 3 - Alpha', + "License :: OSI Approved :: Apache Software License", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.6", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Operating System :: OS Independent", + "Development Status :: 3 - Alpha", + "Topic :: Scientific/Engineering :: Artificial Intelligence", ] +keywords = [ + "torch-optimizer", + "pytorch", + # optimizers + "accsgd", + "adabound", + "adamod", + "diffgrad", + "lamb", + "lookahead", + "madgrad", + "novograd", + "pid", + "qhadam", + "qhm", + "radam", + "sgdw", + "yogi", + "ranger", +] + +project_urls = { + "Website": "https://github.com/jettify/pytorch-optimizer", + "Documentation": "https://pytorch-optimizer.readthedocs.io", + "Issues": "https://github.com/jettify/pytorch-optimizer/issues", +} + setup( - name='torch-optimizer', - version=read_version(), - description=('pytorch-optimizer'), - long_description='\n\n'.join((read('README.rst'), read('CHANGES.rst'))), + name="torch-optimizer", + version=_read_version(), + description=("pytorch-optimizer"), + long_description="\n\n".join((_read("README.rst"), _read("CHANGES.rst"))), + long_description_content_type="text/x-rst", classifiers=classifiers, - platforms=['POSIX'], - author='Nikolay Novik', - author_email='nickolainovik@gmail.com', - url='https://github.com/jettify/pytorch-optimizer', - download_url='https://pypi.org/project/torch-optimizer/', - license='Apache 2', - packages=find_packages(), + platforms=["POSIX"], + author="Nikolay Novik", + author_email="nickolainovik@gmail.com", + url="https://github.com/jettify/pytorch-optimizer", + download_url="https://pypi.org/project/torch-optimizer/", + license="Apache 2", + packages=find_packages(exclude=("tests",)), install_requires=install_requires, - keywords=['torch-optimizer', 'pytorch'], + keywords=keywords, zip_safe=True, include_package_data=True, + project_urls=project_urls, + python_requires=">=3.6.0", ) diff --git a/tests/conftest.py b/tests/conftest.py index ba05ea55..e69de29b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,15 +0,0 @@ -import torch - - -def assert_dict_equal(a, b, precision=0.000001): - if isinstance(a, dict) and isinstance(b, dict): - assert set(a.keys()) == set(b.keys()) - for k in a.keys(): - assert_dict_equal(a[k], b[k], precision) - elif isinstance(a, list) and isinstance(b, list): - assert len(a) == len(b) - for v1, v2 in zip(a, b): - assert_dict_equal(v1, v2, precision) - elif isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor): - assert torch.allclose(a, b, atol=precision) - assert a == b diff --git a/tests/test_basic.py b/tests/test_basic.py index 71c0f809..f0583ae2 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -1,35 +1,94 @@ +import pytest import torch -from torch_optimizer import Lookahead -from torch.autograd import Variable -from torch.optim import Adam + +import torch_optimizer as optim def rosenbrock(tensor): x, y = tensor - return (1 - x) ** 2 + 1 * (y - x ** 2) ** 2 + return (1 - x) ** 2 + 1 * (y - x**2) ** 2 def quadratic(tensor): x, y = tensor a = 1.0 b = 1.0 - return (x ** 2) / a + (y ** 2) / b + return (x**2) / a + (y**2) / b -def saddle(tensor): +def beale(tensor): x, y = tensor - a = 1.0 - b = 1.0 - return (x ** 2) / a - (y ** 2) / b + f = ( + (1.5 - x + x * y) ** 2 + + (2.25 - x + x * y**2) ** 2 + + (2.625 - x + x * y**3) ** 2 + ) + return f + + +cases = [ + (rosenbrock, (1.5, 1.5), (1, 1)), + (quadratic, (1.5, 1.5), (0, 0)), + (beale, (1.5, 1.5), (3, 0.5)), +] + + +def ids(v): + n = "{} {}".format(v[0].__name__, v[1:]) + return n -def test_rosenbrock(): - X = Variable(torch.Tensor([1.5, 1.5]), requires_grad=True) - optimizer = Lookahead(Adam([X], lr=0.01)) - optimizer = Adam([X], lr=0.1) - for _ in range(100): +def build_lookahead(*a, **kw): + base = optim.Yogi(*a, **kw) + return optim.Lookahead(base) + + +optimizers = [ + (optim.A2GradUni, {"lips": 40, "beta": 0.0001}, 800), + (optim.PID, {"lr": 0.002, "momentum": 0.8, "weight_decay": 0.0001}, 900), + (optim.QHM, {"lr": 0.02, "momentum": 0.95, "nu": 1}, 900), + ( + optim.NovoGrad, + {"lr": 2.9, "betas": (0.9, 0.999), "grad_averaging": True}, + 900, + ), + (optim.RAdam, {"lr": 0.01, "betas": (0.9, 0.95), "eps": 1e-3}, 800), + (optim.SGDW, {"lr": 0.002, "momentum": 0.91}, 900), + (optim.DiffGrad, {"lr": 0.5}, 500), + (optim.AdaMod, {"lr": 1.0}, 800), + (optim.AdaBound, {"lr": 1.0}, 800), + (optim.Yogi, {"lr": 1.0}, 500), + (optim.AccSGD, {"lr": 0.015}, 800), + (build_lookahead, {"lr": 1.0}, 500), + (optim.QHAdam, {"lr": 1.0}, 500), + (optim.AdamP, {"lr": 0.01, "betas": (0.9, 0.95), "eps": 1e-3}, 800), + (optim.SGDP, {"lr": 0.002, "momentum": 0.91}, 900), + (optim.AggMo, {"lr": 0.003}, 1800), + (optim.SWATS, {"lr": 0.1, "amsgrad": True, "nesterov": True}, 900), + (optim.Adafactor, {"lr": None, "decay_rate": -0.3, "beta1": 0.9}, 800), + (optim.AdaBelief, {"lr": 1.0}, 500), + (optim.Adahessian, {"lr": 0.15, "hessian_power": 0.6, "seed": 0}, 900), + (optim.MADGRAD, {"lr": 0.02}, 500), + (optim.LARS, {"lr": 0.002, "momentum": 0.91}, 900), + (optim.Lion, {"lr": 0.025}, 3600), +] + + +@pytest.mark.parametrize("case", cases, ids=ids) +@pytest.mark.parametrize("optimizer_config", optimizers, ids=ids) +def test_benchmark_function(case, optimizer_config): + func, initial_state, min_loc = case + optimizer_class, config, iterations = optimizer_config + + x = torch.Tensor(initial_state).requires_grad_(True) + x_min = torch.Tensor(min_loc) + optimizer = optimizer_class([x], **config) + for _ in range(iterations): optimizer.zero_grad() - f = saddle(X) - f.backward(retain_graph=True) + f = func(x) + f.backward(retain_graph=True, create_graph=True) optimizer.step() - print(f, X) + assert torch.allclose(x, x_min, atol=0.001) + + name = optimizer.__class__.__name__ + assert name in optimizer.__repr__() diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index fec83cab..154474e2 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -1,105 +1,127 @@ import functools from copy import deepcopy -import pytest import torch -import torch.optim as optim from torch.autograd import Variable from torch.optim.lr_scheduler import ExponentialLR, ReduceLROnPlateau, StepLR -from torch_optimizer import PowerSign, Lookahead -from tests.utils import assert_dict_equal +import torch_optimizer as optim + + +def assert_dict_equal(a, b, precision=0.000001): + if isinstance(a, dict) and isinstance(b, dict): + assert set(a.keys()) == set(b.keys()) + for k in a.keys(): + assert_dict_equal(a[k], b[k], precision) + elif isinstance(a, list) and isinstance(b, list): + assert len(a) == len(b) + for v1, v2 in zip(a, b): + assert_dict_equal(v1, v2, precision) + elif isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor): + assert torch.allclose(a, b, atol=precision) + else: + assert a == b + return True def _build_params_dict(weight, bias, **kwargs): - return [{'params': [weight]}, dict(params=[bias], **kwargs)] + return [{"params": [weight]}, dict(params=[bias], **kwargs)] def _build_params_dict_single(weight, bias, **kwargs): return [dict(params=bias, **kwargs)] -sgd_cases = [ - (lambda weight, bias: optim.SGD([weight, bias], lr=1e-3),), - ( - lambda weight, bias: optim.SGD( - _build_params_dict(weight, bias, lr=1e-2), lr=1e-3 +def make_test_params(optimizer_class): + cases = [ + (lambda weight, bias: optimizer_class([weight, bias], lr=1e-3),), + ( + lambda weight, bias: optimizer_class( + _build_params_dict(weight, bias, lr=1e-2), lr=1e-3 + ), ), - ), - ( - lambda weight, bias: optim.SGD( - _build_params_dict_single(weight, bias, lr=1e-2), lr=1e-3 + ( + lambda weight, bias: optimizer_class( + _build_params_dict_single(weight, bias, lr=1e-2), lr=1e-3 + ), ), - ), - ( - lambda weight, bias: optim.SGD( - _build_params_dict_single(weight, bias, lr=1e-2) + ( + lambda weight, bias: optimizer_class( + _build_params_dict_single(weight, bias, lr=1e-2) + ), ), - ), - ( - lambda weight, bias: optim.SGD([weight, bias], lr=1e-3), - [lambda opt: StepLR(opt, gamma=0.9, step_size=10)], - ), - ( - lambda weight, bias: optim.SGD([weight, bias], lr=1e-3), - [ - lambda opt: StepLR(opt, gamma=0.9, step_size=10), - lambda opt: ReduceLROnPlateau(opt), - ], - ), - ( - lambda weight, bias: optim.SGD([weight, bias], lr=1e-3), - [ - lambda opt: StepLR(opt, gamma=0.99, step_size=10), - lambda opt: ExponentialLR(opt, gamma=0.99), - lambda opt: ReduceLROnPlateau(opt), - ], - ), -] - - -powsersign_cases = [ - (lambda weight, bias: PowerSign([weight, bias], lr=1e-3),), - ( - lambda weight, bias: PowerSign( - _build_params_dict(weight, bias, lr=1e-2), lr=1e-3 + ( + lambda weight, bias: optimizer_class([weight, bias], lr=1e-3), + [lambda opt: StepLR(opt, gamma=0.9, step_size=10)], ), - ), - ( - lambda weight, bias: PowerSign( - _build_params_dict_single(weight, bias, lr=1e-2), lr=1e-3 + ( + lambda weight, bias: optimizer_class([weight, bias], lr=1e-3), + [ + lambda opt: StepLR(opt, gamma=0.9, step_size=10), + lambda opt: ReduceLROnPlateau(opt), + ], ), - ), - ( - lambda weight, bias: PowerSign( - _build_params_dict_single(weight, bias, lr=1e-2) + ( + lambda weight, bias: optimizer_class([weight, bias], lr=1e-3), + [ + lambda opt: StepLR(opt, gamma=0.99, step_size=10), + lambda opt: ExponentialLR(opt, gamma=0.99), + lambda opt: ReduceLROnPlateau(opt), + ], ), - ), - ( - lambda weight, bias: optim.SGD([weight, bias], lr=1e-3), - [lambda opt: StepLR(opt, gamma=0.9, step_size=10)], - ), - ( - lambda weight, bias: optim.SGD([weight, bias], lr=1e-3), - [ - lambda opt: StepLR(opt, gamma=0.9, step_size=10), - lambda opt: ReduceLROnPlateau(opt), - ], - ), - ( - lambda weight, bias: optim.SGD([weight, bias], lr=1e-3), - [ - lambda opt: StepLR(opt, gamma=0.99, step_size=10), - lambda opt: ExponentialLR(opt, gamma=0.99), - lambda opt: ReduceLROnPlateau(opt), - ], - ), + ] + ids = ["%s_%s" % (optimizer_class.__name__, i) for i in range(len(cases))] + return cases, ids + + +def build_lookahead(*a, **kw): + base = optim.Yogi(*a, **kw) + return optim.Lookahead(base) + + +optimizers = [ + build_lookahead, + optim.A2GradExp, + optim.A2GradInc, + optim.A2GradUni, + optim.AccSGD, + optim.AdaBelief, + optim.AdaBound, + optim.AdaMod, + optim.Adafactor, + optim.AdamP, + optim.AggMo, + optim.Apollo, + optim.DiffGrad, + optim.LARS, + optim.Lamb, + optim.MADGRAD, + optim.NovoGrad, + optim.PID, + optim.QHAdam, + optim.QHM, + optim.RAdam, + optim.Ranger, + optim.RangerQH, + optim.RangerVA, + optim.SGDP, + optim.SGDW, + optim.SWATS, + optim.Shampoo, + optim.Yogi, + optim.Lion, ] -lookahead_cases = [ - (lambda weight, bias: Lookahead(optim.SGD([weight, bias], lr=1e-3)),) -] +def pytest_generate_tests(metafunc): + if "optimizer_constructor" in metafunc.fixturenames: + cases = [] + ids = [] + for o in optimizers: + c, i = make_test_params(o) + cases = cases + c + ids = ids + i + metafunc.parametrize("optimizer_constructor", cases, ids=ids) class TestOptim: @@ -127,7 +149,7 @@ def fn(): ): y = y.cuda(bias.get_device()) loss = (y + bias).pow(2).sum() - loss.backward() + loss.backward(create_graph=True) return loss initial_value = fn().item() @@ -152,7 +174,7 @@ def fn_base(optimizer, weight, bias): optimizer.zero_grad() i = input_cuda if weight.is_cuda else input loss = (weight.mv(i) + bias).pow(2).sum() - loss.backward() + loss.backward(create_graph=True) return loss optimizer = constructor(weight, bias) @@ -170,12 +192,14 @@ def fn_base(optimizer, weight, bias): state_dict = deepcopy(optimizer.state_dict()) state_dict_c = deepcopy(optimizer.state_dict()) optimizer_c.load_state_dict(state_dict_c) + + precision = 0.0001 # Run both optimizations in parallel for _i in range(20): optimizer.step(fn) optimizer_c.step(fn_c) - assert torch.allclose(weight, weight_c) - assert torch.allclose(bias, bias_c) + assert torch.allclose(weight, weight_c, atol=precision) + assert torch.allclose(bias, bias_c, atol=precision) # Make sure state dict wasn't modified assert assert_dict_equal(state_dict, state_dict_c) @@ -208,7 +232,7 @@ def fn_base(optimizer, weight, bias): # validate deepcopy() copies all public attributes def getPublicAttr(obj): - return set(k for k in obj.__dict__ if not k.startswith('_')) + return set(k for k in obj.__dict__ if not k.startswith("_")) assert getPublicAttr(optimizer) == getPublicAttr(deepcopy(optimizer)) @@ -259,20 +283,5 @@ def _test_basic_cases( scheduler_constructors, ) - def test_sgd_validation(self): - with pytest.raises(ValueError) as ctx: - optim.SGD(None, lr=1e-2, momentum=-0.5) - msg = 'Invalid momentum value: -0.5' - assert msg in str(ctx.value) - - @pytest.mark.parametrize('params', sgd_cases) - def test_sgd(self, params): - self._test_basic_cases(*params) - - @pytest.mark.parametrize('params', powsersign_cases) - def test_powersign(self, params): - self._test_basic_cases(*params) - - @pytest.mark.parametrize('params', lookahead_cases) - def test_lookahead(self, params): - self._test_basic_cases(*params) + def test_optimizer(self, optimizer_constructor): + self._test_basic_cases(*optimizer_constructor) diff --git a/tests/test_optimizer_with_nn.py b/tests/test_optimizer_with_nn.py new file mode 100644 index 00000000..80829ad0 --- /dev/null +++ b/tests/test_optimizer_with_nn.py @@ -0,0 +1,114 @@ +import numpy as np +import pytest +import torch +from torch import nn + +import torch_optimizer as optim + + +def make_dataset(seed=42): + rng = np.random.RandomState(seed) + N = 100 + D = 2 + + X = rng.randn(N, D) * 2 + + # center the first N/2 points at (-2,-2) + mid = N // 2 + X[:mid, :] = X[:mid, :] - 2 * np.ones((mid, D)) + + # center the last N/2 points at (2, 2) + X[mid:, :] = X[mid:, :] + 2 * np.ones((mid, D)) + + # labels: first N/2 are 0, last N/2 are 1 + Y = np.array([0] * mid + [1] * mid).reshape(100, 1) + + x = torch.Tensor(X) + y = torch.Tensor(Y) + return x, y + + +class LogisticRegression(nn.Module): + def __init__(self): + super(LogisticRegression, self).__init__() + self.linear1 = nn.Linear(2, 4) + self.linear2 = nn.Linear(4, 1) + + def forward(self, x): + output = torch.relu(self.linear1(x)) + output = self.linear2(output) + y_pred = torch.sigmoid(output) + return y_pred + + +def ids(v): + return "{} {}".format(v[0].__name__, v[1:]) + + +def build_lookahead(*a, **kw): + base = optim.Yogi(*a, **kw) + return optim.Lookahead(base) + + +optimizers = [ + (build_lookahead, {"lr": 0.1, "weight_decay": 1e-3}, 200), + (optim.A2GradExp, {"lips": 2.0, "beta": 1e-3}, 500), + (optim.A2GradInc, {"lips": 5.0, "beta": 1e-3}, 200), + (optim.A2GradUni, {"lips": 5.0, "beta": 1e-3}, 500), + (optim.AccSGD, {"lr": 1.0, "weight_decay": 1e-3}, 200), + (optim.AdaBelief, {"lr": 0.1, "weight_decay": 1e-3}, 200), + (optim.AdaBound, {"lr": 1.5, "gamma": 0.1, "weight_decay": 1e-3}, 200), + (optim.AdaMod, {"lr": 2.0, "weight_decay": 1e-3}, 200), + (optim.Adafactor, {"lr": 0.004466, "weight_decay": 1e-3}, 1500), + (optim.AdamP, {"lr": 0.045, "weight_decay": 1e-3}, 800), + (optim.AggMo, {"lr": 0.17059, "weight_decay": 1e-3}, 1000), + (optim.Apollo, {"lr": 0.1, "weight_decay": 1e-3}, 200), + (optim.DiffGrad, {"lr": 0.5, "weight_decay": 1e-3}, 200), + ( + optim.LARS, + {"lr": 1.0, "weight_decay": 1e-3, "trust_coefficient": 0.01}, + 200, + ), + (optim.Lamb, {"lr": 0.0151, "weight_decay": 1e-3}, 1000), + (optim.MADGRAD, {"lr": 1.0, "weight_decay": 1e-3}, 200), + (optim.NovoGrad, {"lr": 0.01, "weight_decay": 1e-3}, 200), + (optim.PID, {"lr": 0.01, "weight_decay": 1e-3, "momentum": 0.1}, 200), + (optim.QHAdam, {"lr": 0.1, "weight_decay": 1e-3}, 200), + (optim.QHM, {"lr": 0.1, "weight_decay": 1e-5, "momentum": 0.2}, 200), + (optim.RAdam, {"lr": 1.0, "weight_decay": 1e-3}, 200), + (optim.Ranger, {"lr": 0.1, "weight_decay": 1e-3}, 200), + (optim.RangerQH, {"lr": 0.0124, "weight_decay": 1e-3}, 1100), + (optim.RangerVA, {"lr": 0.2214, "weight_decay": 1e-3}, 500), + (optim.SGDP, {"lr": 1.0, "weight_decay": 1e-3}, 200), + (optim.SGDW, {"lr": 1.0, "weight_decay": 1e-3}, 200), + (optim.SWATS, {"lr": 0.703, "weight_decay": 1e-3}, 600), + ( + optim.Shampoo, + {"lr": 0.279, "weight_decay": 1e-3, "momentum": 0.05}, + 1600, + ), + (optim.Yogi, {"lr": 0.1, "weight_decay": 1e-3}, 200), + (optim.Adahessian, {"lr": 0.1, "weight_decay": 1e-3}, 200), + (optim.Lion, {"lr": 0.1, "weight_decay": 1e-3}, 200), +] + + +@pytest.mark.parametrize("optimizer_config", optimizers, ids=ids) +def test_basic_nn_modeloptimizer_config(optimizer_config): + torch.manual_seed(42) + x_data, y_data = make_dataset() + model = LogisticRegression() + + loss_fn = nn.BCELoss() + optimizer_class, config, iterations = optimizer_config + optimizer = optimizer_class(model.parameters(), **config) + init_loss = None + for _ in range(iterations): + y_pred = model(x_data) + loss = loss_fn(y_pred, y_data) + if init_loss is None: + init_loss = loss + optimizer.zero_grad() + loss.backward(create_graph=True) + optimizer.step() + assert init_loss.item() > 2.0 * loss.item() diff --git a/tests/test_param_validation.py b/tests/test_param_validation.py new file mode 100644 index 00000000..c5d74390 --- /dev/null +++ b/tests/test_param_validation.py @@ -0,0 +1,162 @@ +import pytest +import torch + +import torch_optimizer as optim + + +def assert_sparse_not_supported(optimizer_class, err_msg=None): + param = torch.randn(1, 1).to_sparse().requires_grad_(True) + grad = torch.randn(1, 1).to_sparse() + param.grad = grad + optimizer = optimizer_class([param]) + optimizer.zero_grad() + with pytest.raises(RuntimeError) as ctx: + optimizer.step() + + msg = err_msg or "does not support sparse gradients" + assert msg in str(ctx.value) + + +no_sparse_optimizers = [ + optim.AdaBound, + optim.AdaMod, + optim.DiffGrad, + optim.Lamb, + optim.NovoGrad, + optim.RAdam, + optim.Yogi, +] + + +@pytest.mark.parametrize("optimizer_class", no_sparse_optimizers) +def test_sparse_not_supported(optimizer_class): + assert_sparse_not_supported(optimizer_class) + + +optimizers = [ + optim.AccSGD, + optim.AdaBelief, + optim.AdaBound, + optim.AdaMod, + optim.AdamP, + optim.AggMo, + optim.Apollo, + optim.DiffGrad, + optim.LARS, + optim.Lamb, + optim.MADGRAD, + optim.NovoGrad, + optim.PID, + optim.QHAdam, + optim.QHM, + optim.RAdam, + optim.SGDP, + optim.SGDW, + optim.SWATS, + optim.Shampoo, + optim.Yogi, + optim.Lion, +] + + +@pytest.mark.parametrize("optimizer_class", optimizers) +def test_learning_rate(optimizer_class): + lr = -0.01 + with pytest.raises(ValueError) as ctx: + optimizer_class(None, lr=-0.01) + msg = "Invalid learning rate: {}".format(lr) + assert msg in str(ctx.value) + + +eps_optimizers = [ + optim.AdaBelief, + optim.AdaBound, + optim.AdaMod, + optim.AdamP, + optim.Apollo, + optim.DiffGrad, + optim.LARS, + optim.Lamb, + optim.MADGRAD, + optim.NovoGrad, + optim.QHAdam, + optim.RAdam, + optim.SGDP, + optim.SWATS, + optim.Yogi, +] + + +@pytest.mark.parametrize("optimizer_class", eps_optimizers) +def test_eps_validation(optimizer_class): + eps = -0.1 + with pytest.raises(ValueError) as ctx: + optimizer_class(None, lr=0.1, eps=eps) + msg = "Invalid epsilon value: {}".format(eps) + assert msg in str(ctx.value) + + +weight_decay_optimizers = [ + optim.AccSGD, + optim.AdaBelief, + optim.AdaBound, + optim.AdaMod, + optim.Adafactor, + optim.AdamP, + optim.AggMo, + optim.Apollo, + optim.DiffGrad, + optim.LARS, + optim.Lamb, + optim.MADGRAD, + optim.NovoGrad, + optim.PID, + optim.QHAdam, + optim.QHM, + optim.RAdam, + optim.SGDP, + optim.SGDW, + optim.SWATS, + optim.Shampoo, + optim.Yogi, + optim.Lion, +] + + +@pytest.mark.parametrize("optimizer_class", weight_decay_optimizers) +def test_weight_decay_validation(optimizer_class): + weight_decay = -0.1 + with pytest.raises(ValueError) as ctx: + optimizer_class(None, lr=0.1, weight_decay=weight_decay) + msg = "Invalid weight_decay value: {}".format(weight_decay) + assert msg in str(ctx.value) + + +betas_optimizers = [ + optim.AdaBelief, + optim.AdaBound, + optim.AdaMod, + optim.AdamP, + optim.DiffGrad, + optim.Lamb, + optim.NovoGrad, + optim.QHAdam, + optim.RAdam, + optim.Yogi, + optim.Lion, +] + + +@pytest.mark.parametrize("optimizer_class", betas_optimizers) +def test_betas_validation(optimizer_class): + betas = (-1, 0.999) + with pytest.raises(ValueError) as ctx: + optimizer_class(None, lr=0.1, betas=(-1, 0.999)) + msg = "Invalid beta parameter at index 0: {}".format(betas[0]) + assert msg in str(ctx.value) + + betas = (0.9, -0.999) + with pytest.raises(ValueError) as ctx: + optimizer_class(None, lr=0.1, betas=betas) + msg = "Invalid beta parameter at index 1: {}".format(betas[1]) + assert msg in str(ctx.value) diff --git a/tests/utils.py b/tests/utils.py deleted file mode 100644 index 815f2613..00000000 --- a/tests/utils.py +++ /dev/null @@ -1,17 +0,0 @@ -import torch - - -def assert_dict_equal(a, b, precision=0.000001): - if isinstance(a, dict) and isinstance(b, dict): - assert set(a.keys()) == set(b.keys()) - for k in a.keys(): - assert_dict_equal(a[k], b[k], precision) - elif isinstance(a, list) and isinstance(b, list): - assert len(a) == len(b) - for v1, v2 in zip(a, b): - assert_dict_equal(v1, v2, precision) - elif isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor): - assert torch.allclose(a, b, atol=precision) - else: - assert a == b - return True diff --git a/torch_optimizer/__init__.py b/torch_optimizer/__init__.py index 8d119dd3..f0123e53 100644 --- a/torch_optimizer/__init__.py +++ b/torch_optimizer/__init__.py @@ -1,6 +1,130 @@ +"""torch-optimizer -- collection of of optimization algorithms for PyTorch. + +API and usage patterns are the same as `torch.optim`__ + +Example +------- + +>>> import torch_optimizer as optim +# model = ... +>>> optimizer = optim.DiffGrad(model.parameters(), lr=0.001) +>>> optimizer.step() + +See documentation for full list of supported optimizers. + +__ https://pytorch.org/docs/stable/optim.html#module-torch.optim +""" +from typing import Dict, List, Type + +from pytorch_ranger import Ranger, RangerQH, RangerVA +from torch.optim.optimizer import Optimizer + +from .a2grad import A2GradExp, A2GradInc, A2GradUni +from .accsgd import AccSGD +from .adabelief import AdaBelief +from .adabound import AdaBound +from .adafactor import Adafactor +from .adahessian import Adahessian +from .adamod import AdaMod +from .adamp import AdamP +from .aggmo import AggMo +from .apollo import Apollo +from .diffgrad import DiffGrad +from .lamb import Lamb +from .lars import LARS +from .lion import Lion from .lookahead import Lookahead -from .powersign import PowerSign +from .madgrad import MADGRAD +from .novograd import NovoGrad +from .pid import PID +from .qhadam import QHAdam +from .qhm import QHM +from .radam import RAdam +from .sgdp import SGDP +from .sgdw import SGDW +from .shampoo import Shampoo +from .swats import SWATS +from .yogi import Yogi + +__all__ = ( + "A2GradExp", + "A2GradInc", + "A2GradUni", + "AccSGD", + "AdaBelief", + "AdaBound", + "AdaMod", + "Adafactor", + "Adahessian", + "AdamP", + "AggMo", + "Apollo", + "DiffGrad", + "LARS", + "Lamb", + "Lookahead", + "MADGRAD", + "NovoGrad", + "PID", + "QHAdam", + "QHM", + "RAdam", + "Ranger", + "RangerQH", + "RangerVA", + "SGDP", + "SGDW", + "SWATS", + "Shampoo", + "Yogi", + "Lion", + # utils + "get", +) +__version__ = "0.3.1a0" + + +_package_opts = [ + AdaBelief, + AccSGD, + AdaBound, + AdaMod, + AdamP, + AggMo, + DiffGrad, + LARS, + Lamb, + Lookahead, + MADGRAD, + NovoGrad, + PID, + QHAdam, + QHM, + RAdam, + Ranger, + RangerQH, + RangerVA, + SGDP, + SGDW, + SWATS, + Shampoo, + Yogi, + Lion, +] # type: List[Type[Optimizer]] + + +_NAME_OPTIM_MAP = { + opt.__name__.lower(): opt for opt in _package_opts +} # type: Dict[str, Type[Optimizer]] + +def get(name: str) -> Type[Optimizer]: + r"""Returns an optimizer class from its name. Case insensitive. -__all__ = ('Lookahead', 'PowerSign') -__version__ = '0.0.1a0' + Args: + name: the optimizer name. + """ + optimizer_class = _NAME_OPTIM_MAP.get(name.lower()) + if optimizer_class is None: + raise ValueError("Optimizer {} not found".format(name)) + return optimizer_class diff --git a/torch_optimizer/a2grad.py b/torch_optimizer/a2grad.py new file mode 100644 index 00000000..b2976ae0 --- /dev/null +++ b/torch_optimizer/a2grad.py @@ -0,0 +1,309 @@ +import copy +import math +from typing import Optional + +import torch +from torch.optim.optimizer import Optimizer + +from .types import OptFloat, OptLossClosure, Params + +__all__ = ("A2GradUni", "A2GradInc", "A2GradExp") + + +class A2GradUni(Optimizer): + r"""Implements A2GradUni Optimizer Algorithm. + + It has been proposed in `Optimal Adaptive and Accelerated Stochastic + Gradient Descent`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: not used for this optimizer (default: None) + beta: (default: 10) + lips: Lipschitz constant (default: 10) + + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.A2GradUni(model.parameters(), lips=10) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1810.00553 + + Note: + Reference code: https://github.com/severilov/A2Grad_optimizer + """ + + def __init__( + self, + params: Params, + lr: Optional[float] = None, + beta: float = 10, + lips: float = 10, + ): + defaults = dict(beta=beta, lips=lips, lr=lr) + # lr is not supported for this optimizer, we need to make tests work + # and schedulers not to fail + if beta < 0.0: + raise ValueError("Invalid beta value: {}".format(beta)) + if lips < 0.0: + raise ValueError("Invalid lips value: {}".format(lips)) + + super().__init__(params, defaults) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + state = self.state[p] + + if len(state) == 0: + state["step"] = 0 + state["alpha_k"] = 1 + state["v_k"] = 0 + state["avg_grad"] = copy.deepcopy(grad) + state["x_k"] = copy.deepcopy(p.data) + + gamma_k = 2 * group["lips"] / (state["step"] + 1) + + avg_grad = state["avg_grad"] + avg_grad.mul_(state["step"]) + avg_grad.add_(grad) + avg_grad.div_(state["step"] + 1) + + delta_k = torch.add(grad, avg_grad, alpha=-1) + + state["v_k"] += torch.sum(delta_k * delta_k).item() + + h_k = math.sqrt(state["v_k"]) + alpha_k_1 = 2 / (state["step"] + 3) + coef = 1 / (gamma_k + group["beta"] * h_k) + x_k_1 = state["x_k"] + x_k_1.add_(grad, alpha=-coef) + + p.data.mul_(1 - alpha_k_1) + p.data.add_(x_k_1, alpha=alpha_k_1) + p.data.add_( + grad, alpha=-(1 - alpha_k_1) * state["alpha_k"] * coef + ) + + state["alpha_k"] = alpha_k_1 + state["step"] += 1 + + return loss + + +class A2GradInc(Optimizer): + r"""Implements A2GradInc Optimizer Algorithm. + + It has been proposed in `Optimal Adaptive and Accelerated Stochastic + Gradient Descent`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: not used for this optimizer (default: None) + beta: (default: 10) + lips: Lipschitz constant (default: 10) + + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.A2GradInc(model.parameters(), lips=10) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1810.00553 + + Note: + Reference code: https://github.com/severilov/A2Grad_optimizer + """ + + def __init__( + self, + params: Params, + lr: Optional[float] = None, + beta: float = 10, + lips: float = 10, + ): + if beta < 0.0: + raise ValueError("Invalid beta value: {}".format(beta)) + if lips < 0.0: + raise ValueError("Invalid weight_decay value: {}".format(lips)) + defaults = dict(beta=beta, lips=lips, lr=lr) + super(A2GradInc, self).__init__(params, defaults) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + state = self.state[p] + + if len(state) == 0: + state["step"] = 0 + state["alpha_k"] = 1 + state["v_k"] = 0 + state["avg_grad"] = copy.deepcopy(grad) + state["x_k"] = copy.deepcopy(p.data) + + gamma_k = 2 * group["lips"] / (state["step"] + 1) + + avg_grad = state["avg_grad"] + avg_grad.mul_(state["step"]) + avg_grad.add_(grad) + avg_grad.div_(state["step"] + 1) + + delta_k = torch.add(grad, avg_grad, alpha=-1) + + state["v_k"] *= (state["step"] / (state["step"] + 1)) ** 2 + state["v_k"] += torch.sum(delta_k * delta_k).item() + + h_k = math.sqrt(state["v_k"]) + alpha_k_1 = 2 / (state["step"] + 3) + coef = 1 / (gamma_k + group["beta"] * h_k) + x_k_1 = state["x_k"] + x_k_1.add_(grad, alpha=-coef) + + p.data.mul_(1 - alpha_k_1) + p.data.add_(x_k_1, alpha=alpha_k_1) + p.data.add_( + grad, alpha=-(1 - alpha_k_1) * state["alpha_k"] * coef + ) + + state["alpha_k"] = alpha_k_1 + state["step"] += 1 + + return loss + + +class A2GradExp(Optimizer): + r"""Implements A2GradExp Optimizer Algorithm. + + It has been proposed in `Optimal Adaptive and Accelerated Stochastic + Gradient Descent`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: not used for this optimizer (default: None) + beta: (default: 10) + lips: Lipschitz constant (default: 10) + rho: represents the degree of weighting decrease, a constant + smoothing factor between 0 and 1 (default: 0.5) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.A2GradExp(model.parameters(), lips=10) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1810.00553 + + Note: + Reference code: https://github.com/severilov/A2Grad_optimizer + """ + + def __init__( + self, + params: Params, + lr: Optional[float] = None, + beta: float = 10, + lips: float = 10, + rho: float = 0.5, + ): + defaults = dict(beta=beta, lips=lips, rho=rho, lr=lr) + super(A2GradExp, self).__init__(params, defaults) + if beta < 0.0: + raise ValueError("Invalid beta value: {}".format(beta)) + if lips < 0.0: + raise ValueError("Invalid lips value: {}".format(lips)) + if rho < 0.0 or rho > 1.0: + raise ValueError("Invalid rho value: {}".format(rho)) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + state = self.state[p] + + if len(state) == 0: + state["step"] = 0 + state["alpha_k"] = 1 + state["v_k"] = 0 + state["avg_grad"] = copy.deepcopy(grad) + state["x_k"] = copy.deepcopy(p.data) + + gamma_k = 2 * group["lips"] / (state["step"] + 1) + + avg_grad = state["avg_grad"] + avg_grad.mul_(state["step"]) + avg_grad.add_(grad) + avg_grad.div_(state["step"] + 1) + + delta_k = torch.add(grad, avg_grad, alpha=-1) + + if state["step"] == 0: + state["v_kk"] = torch.sum(delta_k * delta_k).item() + else: + state["v_kk"] *= group["rho"] + state["v_kk"] += (1 - group["rho"]) * torch.sum( + delta_k * delta_k + ).item() + state["v_k"] = max([state["v_kk"], state["v_k"]]) + + h_k = math.sqrt((state["step"] + 1) * state["v_k"]) + + alpha_k_1 = 2 / (state["step"] + 3) + + coef = -1 / (gamma_k + group["beta"] * h_k) + x_k_1 = state["x_k"] + x_k_1.add_(grad, alpha=coef) + + p.data.mul_(1 - alpha_k_1) + p.data.add_(x_k_1, alpha=alpha_k_1) + p.data.add_( + grad, alpha=(1 - alpha_k_1) * state["alpha_k"] * coef + ) + + state["alpha_k"] = alpha_k_1 + state["step"] += 1 + + return loss diff --git a/torch_optimizer/accsgd.py b/torch_optimizer/accsgd.py new file mode 100644 index 00000000..08ce4006 --- /dev/null +++ b/torch_optimizer/accsgd.py @@ -0,0 +1,102 @@ +import copy + +from torch.optim.optimizer import Optimizer + +from .types import OptFloat, OptLossClosure, Params + +__all__ = ("AccSGD",) + + +class AccSGD(Optimizer): + r"""Implements AccSGD algorithm. + + It has been proposed in `On the insufficiency of existing momentum + schemes for Stochastic Optimization`__ and `Accelerating Stochastic + Gradient Descent For Least Squares Regression`__ + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + kappa: ratio of long to short step (default: 1000) + xi: statistical advantage parameter (default: 10) + small_const: any value <=1 (default: 0.7) + weight_decay: weight decay (L2 penalty) (default: 0) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.AccSGD(model.parameters(), lr=0.1) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1704.08227 + __ https://arxiv.org/abs/1803.05591 + + Note: + Reference code: https://github.com/rahulkidambi/AccSGD + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + kappa: float = 1000.0, + xi: float = 10.0, + small_const: float = 0.7, + weight_decay: float = 0, + ) -> None: + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if weight_decay < 0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + defaults = dict( + lr=lr, + kappa=kappa, + xi=xi, + small_const=small_const, + weight_decay=weight_decay, + ) + super(AccSGD, self).__init__(params, defaults) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + weight_decay = group["weight_decay"] + large_lr = (group["lr"] * group["kappa"]) / (group["small_const"]) + alpha = 1.0 - ( + (group["small_const"] * group["small_const"] * group["xi"]) + / group["kappa"] + ) + beta = 1.0 - alpha + zeta = group["small_const"] / (group["small_const"] + beta) + for p in group["params"]: + if p.grad is None: + continue + d_p = p.grad.data + if weight_decay != 0: + d_p.add_(p.data, alpha=weight_decay) + param_state = self.state[p] + if "momentum_buffer" not in param_state: + param_state["momentum_buffer"] = copy.deepcopy(p.data) + buf = param_state["momentum_buffer"] + buf.mul_((1.0 / beta) - 1.0) + buf.add_(d_p, alpha=-large_lr) + buf.add_(p.data) + buf.mul_(beta) + + p.data.add_(d_p, alpha=-group["lr"]) + p.data.mul_(zeta) + p.data.add_(buf, alpha=1.0 - zeta) + + return loss diff --git a/torch_optimizer/adabelief.py b/torch_optimizer/adabelief.py new file mode 100644 index 00000000..7131f6f6 --- /dev/null +++ b/torch_optimizer/adabelief.py @@ -0,0 +1,220 @@ +import math + +import torch +from torch.optim.optimizer import Optimizer + +from .types import Betas2, OptFloat, OptLossClosure, Params + +__all__ = ("AdaBelief",) + + +class AdaBelief(Optimizer): + r"""Implements AdaBelief Optimizer Algorithm. + It has been proposed in `AdaBelief Optimizer, adapting stepsizes by + the belief in observed gradients`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-2) + betas: coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps: term added to the denominator to improve + numerical stability (default: 0.001) + weight_decay: weight decay (L2 penalty) (default: 0) + amsgrad: whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) + weight_decouple: If set as True, then the optimizer uses decoupled + weight decay as in AdamW (default: False) + fixed_decay : This is used when + weight_decouple is set as True. + When fixed_decay == True, the weight decay is performed as + $W_{new} = W_{old} - W_{old} \times decay$. + When fixed_decay == False, the weight decay is performed as + $W_{new} = W_{old} - W_{old} \times decay \times lr$. Note that in + this case, the weight decay ratio decreases with learning + rate (lr). (default: False) + rectify: (default: False) If set as True, then perform the rectified + update similar to RAdam + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.AdaBelief(model.parameters(), lr=0.01) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/2010.07468 + + Note: + Reference code: https://github.com/juntang-zhuang/Adabelief-Optimizer + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + betas: Betas2 = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 0, + amsgrad: bool = False, + weight_decouple: bool = False, + fixed_decay: bool = False, + rectify: bool = False, + ) -> None: + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + "Invalid beta parameter at index 0: {}".format(betas[0]) + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if weight_decay < 0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + defaults = dict( + lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + amsgrad=amsgrad, + ) + super(AdaBelief, self).__init__(params, defaults) + + self._weight_decouple = weight_decouple + self._rectify = rectify + self._fixed_decay = fixed_decay + + def __setstate__(self, state): + super(AdaBelief, self).__setstate__(state) + for group in self.param_groups: + group.setdefault("amsgrad", False) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError( + "AdaBelief does not support sparse gradients, " + "please consider SparseAdam instead" + ) + amsgrad = group["amsgrad"] + + state = self.state[p] + + beta1, beta2 = group["betas"] + + # State initialization + if len(state) == 0: + state["rho_inf"] = 2.0 / (1.0 - beta2) - 1.0 + state["step"] = 0 + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + p.data, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_avg_var"] = torch.zeros_like( + p.data, memory_format=torch.preserve_format + ) + if amsgrad: + # Maintains max of all exp. moving avg. of + # sq. grad. values + state["max_exp_avg_var"] = torch.zeros_like( + p.data, memory_format=torch.preserve_format + ) + + # get current state variable + exp_avg, exp_avg_var = state["exp_avg"], state["exp_avg_var"] + + state["step"] += 1 + bias_correction1 = 1 - beta1 ** state["step"] + bias_correction2 = 1 - beta2 ** state["step"] + + # perform weight decay, check if decoupled weight decay + if self._weight_decouple: + if not self._fixed_decay: + p.data.mul_(1.0 - group["lr"] * group["weight_decay"]) + else: + p.data.mul_(1.0 - group["weight_decay"]) + else: + if group["weight_decay"] != 0: + grad.add_(p.data, alpha=group["weight_decay"]) + + # Update first and second moment running average + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + grad_residual = grad - exp_avg + exp_avg_var.mul_(beta2).addcmul_( + grad_residual, grad_residual, value=1 - beta2 + ) + + if amsgrad: + max_exp_avg_var = state["max_exp_avg_var"] + # Maintains the maximum of all 2nd moment running + # avg. till now + torch.max( + max_exp_avg_var, exp_avg_var, out=max_exp_avg_var + ) + + # Use the max. for normalizing running avg. of gradient + denom = ( + max_exp_avg_var.add_(group["eps"]).sqrt() + / math.sqrt(bias_correction2) + ).add_(group["eps"]) + else: + denom = ( + exp_avg_var.add_(group["eps"]).sqrt() + / math.sqrt(bias_correction2) + ).add_(group["eps"]) + + if not self._rectify: + # Default update + step_size = group["lr"] / bias_correction1 + p.data.addcdiv_(exp_avg, denom, value=-step_size) + + else: # Rectified update + # calculate rho_t + state["rho_t"] = state["rho_inf"] - 2 * state[ + "step" + ] * beta2 ** state["step"] / (1.0 - beta2 ** state["step"]) + + if ( + state["rho_t"] > 4 + ): # perform Adam style update if variance is small + rho_inf, rho_t = state["rho_inf"], state["rho_t"] + rt = ( + (rho_t - 4.0) + * (rho_t - 2.0) + * rho_inf + / (rho_inf - 4.0) + / (rho_inf - 2.0) + / rho_t + ) + rt = math.sqrt(rt) + + step_size = rt * group["lr"] / bias_correction1 + + p.data.addcdiv_(-step_size, exp_avg, denom) + + else: # perform SGD style update + p.data.add_(-group["lr"], exp_avg) + + return loss diff --git a/torch_optimizer/adabound.py b/torch_optimizer/adabound.py new file mode 100644 index 00000000..587e9f18 --- /dev/null +++ b/torch_optimizer/adabound.py @@ -0,0 +1,183 @@ +import math + +import torch +from torch.optim.optimizer import Optimizer + +from .types import Betas2, OptFloat, OptLossClosure, Params, State + +__all__ = ("AdaBound",) + + +class AdaBound(Optimizer): + r"""Implements AdaBound algorithm. + + It has been proposed in `Adaptive Gradient Methods with Dynamic Bound of + Learning Rate`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + betas: coefficients used for computing running averages of gradient + and its square (default: (0.9, 0.999)) + final_lr: final (SGD) learning rate (default: 0.1) + gamma: convergence speed of the bound functions + (default: 1e-3) + eps: term added to the denominator to improve numerical stability + (default: 1e-8) + weight_decay: weight decay (L2 penalty) (default: 0) + amsbound: whether to use the AMSBound variant of this algorithm + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.AdaBound(model.parameters(), lr=0.1) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1902.09843 + + Note: + Reference code: https://github.com/Luolc/AdaBound + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + betas: Betas2 = (0.9, 0.999), + final_lr: float = 0.1, + gamma: float = 1e-3, + eps: float = 1e-8, + weight_decay: float = 0, + amsbound: bool = False, + ) -> None: + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + "Invalid beta parameter at index 0: {}".format(betas[0]) + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if final_lr < 0.0: + raise ValueError( + "Invalid final learning rate: {}".format(final_lr) + ) + if not 0.0 <= gamma < 1.0: + raise ValueError("Invalid gamma parameter: {}".format(gamma)) + if weight_decay < 0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + defaults = dict( + lr=lr, + betas=betas, + final_lr=final_lr, + gamma=gamma, + eps=eps, + weight_decay=weight_decay, + amsbound=amsbound, + ) + super(AdaBound, self).__init__(params, defaults) + self.base_lrs = [group["lr"] for group in self.param_groups] + + def __setstate__(self, state: State) -> None: + super(AdaBound, self).__setstate__(state) + for group in self.param_groups: + group.setdefault("amsbound", False) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group, base_lr in zip(self.param_groups, self.base_lrs): + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + if grad.is_sparse: + msg = ( + "AdaBound does not support sparse gradients, " + "please consider SparseAdam instead" + ) + raise RuntimeError(msg) + amsbound = group["amsbound"] + + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = 0 + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + if amsbound: + # Maintains max of all exp. moving avg. of + # sq. grad. values + state["max_exp_avg_sq"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + if amsbound: + max_exp_avg_sq = state["max_exp_avg_sq"] + beta1, beta2 = group["betas"] + + state["step"] += 1 + + if group["weight_decay"] != 0: + grad = grad.add(p.data, alpha=group["weight_decay"]) + + # Decay the first and second moment running average coefficient + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + if amsbound: + # Maintains the maximum of all 2nd moment running + # avg. till now + torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq) + # Use the max. for normalizing running avg. of gradient + denom = max_exp_avg_sq.sqrt().add_(group["eps"]) + else: + denom = exp_avg_sq.sqrt().add_(group["eps"]) + + bias_correction1 = 1 - beta1 ** state["step"] + bias_correction2 = 1 - beta2 ** state["step"] + step_size = ( + group["lr"] + * math.sqrt(bias_correction2) + / bias_correction1 + ) + + # Applies bounds on actual learning rate + # lr_scheduler cannot affect final_lr, this is a workaround + # to apply lr decay + final_lr = group["final_lr"] * group["lr"] / base_lr + lower_bound = final_lr * ( + 1 - 1 / (group["gamma"] * state["step"] + 1) + ) + upper_bound = final_lr * ( + 1 + 1 / (group["gamma"] * state["step"]) + ) + step_size = torch.full_like(denom, step_size) + step_size.div_(denom).clamp_(lower_bound, upper_bound).mul_( + exp_avg + ) + + p.data.add_(-step_size) + return loss diff --git a/torch_optimizer/adafactor.py b/torch_optimizer/adafactor.py new file mode 100644 index 00000000..ed1756b9 --- /dev/null +++ b/torch_optimizer/adafactor.py @@ -0,0 +1,218 @@ +import math +from typing import Any, Dict, Tuple + +import torch +from torch.optim.optimizer import Optimizer + +from .types import OptFloat, OptLossClosure, Params, State + +Eps2 = Tuple[float, float] +ParamGroup = Dict[str, Any] + + +class Adafactor(Optimizer): + """Implements Adafactor algorithm. + + It has been proposed in: `Adafactor: Adaptive Learning Rates with + Sublinear Memory Cost`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: external learning rate (default: None) + eps2: regularization constans for square gradient + and parameter scale respectively (default: (1e-30, 1e-3)) + clip_threshold: threshold of root mean square of + final gradient update (default: 1.0) + decay_rate: coefficient used to compute running averages of square + gradient (default: -0.8) + beta1: coefficient used for computing running averages of gradient + (default: None) + weight_decay: weight decay (L2 penalty) (default: 0) + scale_parameter: if true, learning rate is scaled by root mean square + of parameter (default: True) + relative_step: if true, time-dependent learning rate is computed + instead of external learning rate (default: True) + warmup_init: time-dependent learning rate computation depends on + whether warm-up initialization is being used (default: False) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.Adafactor(model.parameters()) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1804.04235 + + Note: + Reference code: https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py # noqa + """ + + def __init__( + self, + params: Params, + lr: OptFloat = None, + eps2: Eps2 = (1e-30, 1e-3), + clip_threshold: float = 1.0, + decay_rate: float = -0.8, + beta1: OptFloat = None, + weight_decay: float = 0.0, + scale_parameter: bool = True, + relative_step: bool = True, + warmup_init: bool = False, + ): + if lr is not None and lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if weight_decay < 0.0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + + defaults = dict( + lr=lr, + eps2=eps2, + clip_threshold=clip_threshold, + decay_rate=decay_rate, + beta1=beta1, + weight_decay=weight_decay, + scale_parameter=scale_parameter, + relative_step=relative_step, + warmup_init=warmup_init, + ) + super(Adafactor, self).__init__(params, defaults) + + def _get_lr(self, param_group: ParamGroup, param_state: State) -> float: + rel_step_sz = param_group["lr"] + if param_group["relative_step"]: + min_step = ( + 1e-6 * param_state["step"] + if param_group["warmup_init"] + else 1e-2 + ) + rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state["step"])) + param_scale = 1.0 + if param_group["scale_parameter"]: + param_scale = max(param_group["eps2"][1], param_state["RMS"]) + return param_scale * rel_step_sz + + def _get_options( + self, param_group: ParamGroup, param_shape: Tuple[int, ...] + ) -> Tuple[bool, bool]: + factored = len(param_shape) >= 2 + use_first_moment = param_group["beta1"] is not None + return factored, use_first_moment + + def _rms(self, tensor: torch.Tensor) -> float: + return tensor.norm(2) / (tensor.numel() ** 0.5) + + def _approx_sq_grad( + self, + exp_avg_sq_row: torch.Tensor, + exp_avg_sq_col: torch.Tensor, + output: torch.Tensor, + ) -> None: + r_factor = ( + (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1)) + .rsqrt_() + .unsqueeze(-1) + ) + c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt() + torch.mul(r_factor, c_factor, out=output) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError( + "Adafactor does not support sparse gradients." + ) + + state = self.state[p] + grad_shape = grad.shape + + factored, use_first_moment = self._get_options( + group, grad_shape + ) + # State Initialization + if len(state) == 0: + state["step"] = 0 + + if use_first_moment: + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + grad, memory_format=torch.preserve_format + ) + if factored: + state["exp_avg_sq_row"] = torch.zeros( + grad_shape[:-1] + ).type_as(grad) + state["exp_avg_sq_col"] = torch.zeros( + grad_shape[:-2] + grad_shape[-1:] + ).type_as(grad) + else: + state["exp_avg_sq"] = torch.zeros_like( + grad, memory_format=torch.preserve_format + ) + + state["RMS"] = 0 + + state["step"] += 1 + state["RMS"] = self._rms(p.data) + lr = self._get_lr(group, state) + + beta2t = 1.0 - math.pow(state["step"], group["decay_rate"]) + update = (grad**2) + group["eps2"][0] + if factored: + exp_avg_sq_row = state["exp_avg_sq_row"] + exp_avg_sq_col = state["exp_avg_sq_col"] + + exp_avg_sq_row.mul_(beta2t).add_( + update.mean(dim=-1), alpha=1.0 - beta2t + ) + exp_avg_sq_col.mul_(beta2t).add_( + update.mean(dim=-2), alpha=1.0 - beta2t + ) + + # Approximation of exponential moving average of square + # of gradient + self._approx_sq_grad( + exp_avg_sq_row, exp_avg_sq_col, update + ) + update.mul_(grad) + else: + exp_avg_sq = state["exp_avg_sq"] + + exp_avg_sq.mul_(beta2t).add_(update, alpha=1.0 - beta2t) + torch.rsqrt(exp_avg_sq, out=update).mul_(grad) + + update.div_( + max(1.0, self._rms(update) / group["clip_threshold"]) + ) + update.mul_(lr) + + if use_first_moment: + exp_avg = state["exp_avg"] + exp_avg.mul_(group["beta1"]).add_( + update, alpha=1 - group["beta1"] + ) + update = exp_avg + + if group["weight_decay"] != 0: + p.data.add_(p.data, alpha=-group["weight_decay"] * lr) + + p.data.add_(-update) + + return loss diff --git a/torch_optimizer/adahessian.py b/torch_optimizer/adahessian.py new file mode 100644 index 00000000..6c836481 --- /dev/null +++ b/torch_optimizer/adahessian.py @@ -0,0 +1,204 @@ +import math +from typing import List, Optional + +import torch +from torch.optim.optimizer import Optimizer + +from .types import Betas2, OptFloat, OptLossClosure, Params + +Grads = Params + +__all__ = ("Adahessian",) + + +class Adahessian(Optimizer): + r"""Implements Adahessian Algorithm. + It has been proposed in `ADAHESSIAN: An Adaptive Second Order Optimizer + for Machine Learning`. + + Arguments: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups + lr (float, optional): learning rate (default: 0.15) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-4) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + hessian_power (float, optional): Hessian power (default: 0.5) + seed (int, optional): Random number generator seed (default: None) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.Adahessian(model.parameters(), lr = 1.0) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward(create_graph=True) + >>> optimizer.step() + + __ https://arxiv.org/abs/2006.00719 + + Note: + Reference code: https://github.com/amirgholami/adahessian + """ + + def __init__( + self, + params: Params, + lr: float = 0.15, + betas: Betas2 = (0.9, 0.999), + eps: float = 1e-4, + weight_decay: float = 0, + hessian_power: float = 0.5, + seed: Optional[int] = None, + ) -> None: + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps <= 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + "Invalid beta parameter at index 0: {}".format(betas[0]) + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if not 0.0 <= hessian_power <= 1.0: + raise ValueError( + "Invalid Hessian power value: {}".format(hessian_power) + ) + if seed is not None: + torch.manual_seed(seed) + defaults = dict( + lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + hessian_power=hessian_power, + ) + super(Adahessian, self).__init__(params, defaults) + + def get_trace(self, params: Params, grads: Grads) -> List[torch.Tensor]: + """Get an estimate of Hessian Trace. + This is done by computing the Hessian vector product with a random + vector v at the current gradient point, to estimate Hessian trace by + computing the gradient of . + :param gradsH: a list of torch variables + :return: a list of torch tensors + """ + + # Check backward was called with create_graph set to True + for i, grad in enumerate(grads): + if grad.grad_fn is None: + msg = ( + "Gradient tensor {:} does not have grad_fn. When " + "calling loss.backward(), make sure the option " + "create_graph is set to True." + ) + raise RuntimeError(msg.format(i)) + + v = [ + 2 + * torch.randint_like( + p, high=2, memory_format=torch.preserve_format + ) + - 1 + for p in params + ] + + # this is for distributed setting with single node and multi-gpus, + # for multi nodes setting, we have not support it yet. + hvs = torch.autograd.grad( + grads, params, grad_outputs=v, only_inputs=True, retain_graph=True + ) + + hutchinson_trace = [] + for hv in hvs: + param_size = hv.size() + if len(param_size) <= 2: # for 0/1/2D tensor + # Hessian diagonal block size is 1 here. + # We use that torch.abs(hv * vi) = hv.abs() + tmp_output = hv.abs() + + elif len(param_size) == 4: # Conv kernel + # Hessian diagonal block size is 9 here: torch.sum() reduces + # the dim 2/3. + # We use that torch.abs(hv * vi) = hv.abs() + tmp_output = torch.mean(hv.abs(), dim=[2, 3], keepdim=True) + hutchinson_trace.append(tmp_output) + + return hutchinson_trace + + def step(self, closure: OptLossClosure = None) -> OptFloat: + """Perform a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + params = [] + groups = [] + grads = [] + + # Flatten groups into lists, so that + # hut_traces can be called with lists of parameters + # and grads + for group in self.param_groups: + for p in group["params"]: + if p.grad is not None: + params.append(p) + groups.append(group) + grads.append(p.grad) + + # get the Hessian diagonal + + hut_traces = self.get_trace(params, grads) + + for p, group, grad, hut_trace in zip( + params, groups, grads, hut_traces + ): + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = 0 + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like(p.data) + # Exponential moving average of Hessian diagonal square values + state["exp_hessian_diag_sq"] = torch.zeros_like(p.data) + + exp_avg, exp_hessian_diag_sq = ( + state["exp_avg"], + state["exp_hessian_diag_sq"], + ) + + beta1, beta2 = group["betas"] + + state["step"] += 1 + + # Decay the first and second moment running average coefficient + exp_avg.mul_(beta1).add_(grad.detach_(), alpha=1 - beta1) + exp_hessian_diag_sq.mul_(beta2).addcmul_( + hut_trace, hut_trace, value=1 - beta2 + ) + + bias_correction1 = 1 - beta1 ** state["step"] + bias_correction2 = 1 - beta2 ** state["step"] + + # make the square root, and the Hessian power + k = group["hessian_power"] + denom = ( + (exp_hessian_diag_sq.sqrt() ** k) + / math.sqrt(bias_correction2) ** k + ).add_(group["eps"]) + + # make update + p.data = p.data - group["lr"] * ( + exp_avg / bias_correction1 / denom + + group["weight_decay"] * p.data + ) + + return loss diff --git a/torch_optimizer/adamod.py b/torch_optimizer/adamod.py new file mode 100644 index 00000000..1d32d82b --- /dev/null +++ b/torch_optimizer/adamod.py @@ -0,0 +1,152 @@ +import math + +import torch +from torch.optim.optimizer import Optimizer + +from .types import Betas2, OptFloat, OptLossClosure, Params + +__all__ = ("AdaMod",) + + +class AdaMod(Optimizer): + r"""Implements AdaMod algorithm. + + It has been proposed in `Adaptive and Momental Bounds for Adaptive + Learning Rate Methods`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + betas: coefficients used for computing running averages of gradient + and its square (default: (0.9, 0.999)) + beta3: smoothing coefficient for adaptive learning rates + (default: 0.9999) + eps: term added to the denominator to improve numerical stability + (default: 1e-8) + weight_decay: weight decay (L2 penalty) (default: 0) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.AdaMod(model.parameters(), lr=0.1) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1910.12249 + + Note: + Reference code: https://github.com/lancopku/AdaMod + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + betas: Betas2 = (0.9, 0.999), + beta3: float = 0.999, + eps: float = 1e-8, + weight_decay: float = 0, + ) -> None: + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + "Invalid beta parameter at index 0: {}".format(betas[0]) + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if not 0.0 <= beta3 < 1.0: + raise ValueError("Invalid beta3 parameter: {}".format(beta3)) + if weight_decay < 0.0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + defaults = dict( + lr=lr, betas=betas, beta3=beta3, eps=eps, weight_decay=weight_decay + ) + super(AdaMod, self).__init__(params, defaults) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + """Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + if grad.is_sparse: + msg = "AdaMod does not support sparse gradients" + raise RuntimeError(msg) + + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = 0 + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + # Exponential moving average of actual learning rates + state["exp_avg_lr"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + + exp_avg, exp_avg_sq, exp_avg_lr = ( + state["exp_avg"], + state["exp_avg_sq"], + state["exp_avg_lr"], + ) + beta1, beta2 = group["betas"] + + state["step"] += 1 + + # Decay the first and second moment running average coefficient + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + + denom = exp_avg_sq.sqrt().add_(group["eps"]) + + bias_correction1 = 1 - beta1 ** state["step"] + bias_correction2 = 1 - beta2 ** state["step"] + step_size = ( + group["lr"] + * math.sqrt(bias_correction2) + / bias_correction1 + ) + + if group["weight_decay"] != 0: + p.data.add_( + p.data, alpha=-group["weight_decay"] * group["lr"] + ) + + # Applies momental bounds on actual learning rates + step_size = torch.full_like( + denom, step_size, memory_format=torch.preserve_format + ) + step_size.div_(denom) + exp_avg_lr.mul_(group["beta3"]).add_( + step_size, alpha=1 - group["beta3"] + ) + step_size = torch.min(step_size, exp_avg_lr) + step_size.mul_(exp_avg) + + p.data.add_(-step_size) + + return loss diff --git a/torch_optimizer/adamp.py b/torch_optimizer/adamp.py new file mode 100644 index 00000000..1ddf8110 --- /dev/null +++ b/torch_optimizer/adamp.py @@ -0,0 +1,199 @@ +import math + +import torch +from torch.optim.optimizer import Optimizer + +from .types import Betas2, OptFloat, OptLossClosure, Params + +__all__ = ("AdamP",) + + +class AdamP(Optimizer): + r"""Implements AdamP algorithm. + + It has been proposed in `Slowing Down the Weight Norm Increase in + Momentum-based Optimizers`__ + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + betas: coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps: term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay: weight decay (L2 penalty) (default: 0) + delta: threhold that determines whether a set of parameters is scale + invariant or not (default: 0.1) + wd_ratio: relative weight decay applied on scale-invariant parameters + compared to that applied on scale-variant parameters (default: 0.1) + nesterov: enables Nesterov momentum (default: False) + + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.AdamP(model.parameters(), lr=0.1) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/2006.08217 + + Note: + Reference code: https://github.com/clovaai/AdamP + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + betas: Betas2 = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 0, + delta: float = 0.1, + wd_ratio: float = 0.1, + nesterov: bool = False, + ) -> None: + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + "Invalid beta parameter at index 0: {}".format(betas[0]) + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if weight_decay < 0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + if delta < 0: + raise ValueError("Invalid delta value: {}".format(delta)) + if wd_ratio < 0: + raise ValueError("Invalid wd_ratio value: {}".format(wd_ratio)) + + defaults = dict( + lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + delta=delta, + wd_ratio=wd_ratio, + nesterov=nesterov, + ) + super(AdamP, self).__init__(params, defaults) + + @staticmethod + def _channel_view(x): + return x.view(x.size(0), -1) + + @staticmethod + def _layer_view(x): + return x.view(1, -1) + + @staticmethod + def _cosine_similarity(x, y, eps, view_func): + x = view_func(x) + y = view_func(y) + + x_norm = x.norm(dim=1).add_(eps) + y_norm = y.norm(dim=1).add_(eps) + dot = (x * y).sum(dim=1) + + return dot.abs() / x_norm / y_norm + + def _projection(self, p, grad, perturb, delta, wd_ratio, eps): + wd = 1 + expand_size = [-1] + [1] * (len(p.shape) - 1) + for view_func in [self._channel_view, self._layer_view]: + cosine_sim = self._cosine_similarity(grad, p.data, eps, view_func) + + if cosine_sim.max() < delta / math.sqrt(view_func(p.data).size(1)): + p_n = p.data / view_func(p.data).norm(dim=1).view( + expand_size + ).add_(eps) + perturb -= p_n * view_func(p_n * perturb).sum(dim=1).view( + expand_size + ) + wd = wd_ratio + + return perturb, wd + + return perturb, wd + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + + grad = p.grad.data + beta1, beta2 = group["betas"] + nesterov = group["nesterov"] + + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = 0 + state["exp_avg"] = torch.zeros_like( + p.data, memory_format=torch.preserve_format + ) + state["exp_avg_sq"] = torch.zeros_like( + p.data, memory_format=torch.preserve_format + ) + + # Adam + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + + state["step"] += 1 + bias_correction1 = 1 - beta1 ** state["step"] + bias_correction2 = 1 - beta2 ** state["step"] + + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + + denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_( + group["eps"] + ) + step_size = group["lr"] / bias_correction1 + + if nesterov: + perturb = (beta1 * exp_avg + (1 - beta1) * grad) / denom + else: + perturb = exp_avg / denom + + # Projection + wd_ratio = 1 + if len(p.shape) > 1: + perturb, wd_ratio = self._projection( + p, + grad, + perturb, + group["delta"], + group["wd_ratio"], + group["eps"], + ) + + # Weight decay + if group["weight_decay"] > 0: + p.data.mul_( + 1 - group["lr"] * group["weight_decay"] * wd_ratio + ) + + # Step + p.data.add_(perturb, alpha=-step_size) + + return loss diff --git a/torch_optimizer/aggmo.py b/torch_optimizer/aggmo.py new file mode 100644 index 00000000..1c2f196f --- /dev/null +++ b/torch_optimizer/aggmo.py @@ -0,0 +1,105 @@ +from typing import List, Tuple, Type, TypeVar, Union + +import torch +from torch.optim.optimizer import Optimizer + +from .types import OptFloat, OptLossClosure, Params + +__all__ = ("AggMo",) + + +T = TypeVar("T", bound="AggMo") + + +class AggMo(Optimizer): + r"""Implements Aggregated Momentum Gradient Descent. + + It has been proposed in `Aggregated Momentum: Stability Through Passive + Damping`__ + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.AggMo(model.parameters(), lr=0.1) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1804.00325 + + Note: + Reference code: https://github.com/AtheMathmo/AggMo/blob/master/aggmo.py # noqa + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + betas: Union[List[float], Tuple[float, ...]] = (0.0, 0.9, 0.99), + weight_decay: float = 0, + ) -> None: + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + + for i, beta in enumerate(betas): + if not 0.0 <= beta < 1.0: + msg = "Invalid beta parameter at index 1: {}".format(betas[i]) + raise ValueError(msg) + + if weight_decay < 0.0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + + defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay) + super(AggMo, self).__init__(params, defaults) + + @classmethod + def from_exp_form( + cls: Type[T], + params: Params, + lr: float = 1e-3, + a: float = 0.1, + k: int = 3, + weight_decay: float = 0, + ) -> T: + if lr <= 0.0: + raise ValueError("Invalid parameter k: {}".format(k)) + + betas = [1 - a**i for i in range(k)] # type: List[float] + return cls(params, lr, betas, weight_decay) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + weight_decay = group["weight_decay"] + betas = group["betas"] + total_mom = float(len(betas)) + + for p in group["params"]: + if p.grad is None: + continue + d_p = p.grad.data + if weight_decay != 0: + d_p.add_(p.data, alpha=weight_decay) + param_state = self.state[p] + if "momentum_buffer" not in param_state: + param_state["momentum_buffer"] = {} + for beta in betas: + param_state["momentum_buffer"][ + beta + ] = torch.zeros_like( + p.data, memory_format=torch.preserve_format + ) + for beta in betas: + buf = param_state["momentum_buffer"][beta] + buf.mul_(beta).add_(d_p) + p.data.sub_(buf, alpha=group["lr"] / total_mom) + return loss diff --git a/torch_optimizer/apollo.py b/torch_optimizer/apollo.py new file mode 100644 index 00000000..d2083a75 --- /dev/null +++ b/torch_optimizer/apollo.py @@ -0,0 +1,158 @@ +import torch +from torch.optim.optimizer import Optimizer + +from .types import OptFloat, OptLossClosure, Params + + +class Apollo(Optimizer): + r"""Implements Apollo Optimizer Algorithm. + + It has been proposed in `Apollo: An Adaptive Parameter-wise Diagonal + Quasi-Newton Method for Nonconvex Stochastic Optimization`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-2) + beta: coefficient used for computing + running averages of gradient (default: 0.9) + eps: term added to the denominator to improve + numerical stability (default: 1e-4) + warmup: number of warmup steps (default: 0) + init_lr: initial learning rate for warmup (default: 0.01) + weight_decay: weight decay (L2 penalty) (default: 0) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.Apollo(model.parameters(), lr=0.01) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/2009.13586 + + Note: + Reference code: https://github.com/XuezheMax/apollo + """ + + def __init__( + self, + params: Params, + lr: float = 1e-2, + beta: float = 0.9, + eps: float = 1e-4, + warmup: int = 0, + init_lr: float = 0.01, + weight_decay: float = 0, + ): + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= beta < 1.0: + raise ValueError("Invalid beta parameter: {}".format(beta)) + if not 0.0 <= weight_decay: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + if not 0.0 <= warmup: + raise ValueError("Invalid warmup updates: {}".format(warmup)) + if not 0.0 <= init_lr <= 1.0: + raise ValueError( + "Invalid initial learning rate: {}".format(init_lr) + ) + + defaults = dict( + lr=lr, + beta=beta, + eps=eps, + warmup=warmup, + init_lr=init_lr, + base_lr=lr, + weight_decay=weight_decay, + ) + super(Apollo, self).__init__(params, defaults) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = 0 + # Exponential moving average of gradient values + state["exp_avg_grad"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["approx_hessian"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + # Previous update direction + state["update"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + + # Calculate current lr + if state["step"] < group["warmup"]: + curr_lr = (group["base_lr"] - group["init_lr"]) * state[ + "step" + ] / group["warmup"] + group["init_lr"] + else: + curr_lr = group["lr"] + + # Perform optimization step + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError( + "Atom does not support sparse gradients." + ) + + # Perform step weight decay + if group["weight_decay"] != 0: + grad = grad.add(p, alpha=group["weight_decay"]) + + beta = group["beta"] + exp_avg_grad = state["exp_avg_grad"] + B = state["approx_hessian"] + d_p = state["update"] + + state["step"] += 1 + bias_correction = 1 - beta ** state["step"] + alpha = (1 - beta) / bias_correction + + # Update the running average grad + delta_grad = grad - exp_avg_grad + exp_avg_grad.add_(delta_grad, alpha=alpha) + + denom = d_p.norm(p=4).add(group["eps"]) + d_p.div_(denom) + v_sq = d_p.mul(d_p) + delta = ( + delta_grad.div_(denom).mul_(d_p).sum().mul(-alpha) + - B.mul(v_sq).sum() + ) + + # Update B + B.addcmul_(v_sq, delta) + + # calc direction of parameter updates + denom = B.abs().clamp_(min=1) + d_p.copy_(exp_avg_grad.div(denom)) + + p.data.add_(d_p, alpha=-curr_lr) + + return loss diff --git a/torch_optimizer/diffgrad.py b/torch_optimizer/diffgrad.py new file mode 100644 index 00000000..1ec0d0d8 --- /dev/null +++ b/torch_optimizer/diffgrad.py @@ -0,0 +1,145 @@ +import math + +import torch +from torch.optim.optimizer import Optimizer + +from .types import Betas2, OptFloat, OptLossClosure, Params + +__all__ = ("DiffGrad",) + + +class DiffGrad(Optimizer): + r"""Implements DiffGrad algorithm. + + It has been proposed in `DiffGrad: An Optimization Method for + Convolutional Neural Networks`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + betas: coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps: term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay: weight decay (L2 penalty) (default: 0) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.DiffGrad(model.parameters(), lr=0.1) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1909.11015 + + Note: + Reference code: https://github.com/shivram1987/diffGrad + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + betas: Betas2 = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 0.0, + ) -> None: + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + "Invalid beta parameter at index 0: {}".format(betas[0]) + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if weight_decay < 0.0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + + defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay) + super(DiffGrad, self).__init__(params, defaults) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + beta1, beta2 = group["betas"] + + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + if grad.is_sparse: + msg = ( + "DiffGrad does not support sparse gradients, " + "please consider SparseAdam instead" + ) + raise RuntimeError(msg) + + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = 0 + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + # Previous gradient + state["previous_grad"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + + exp_avg, exp_avg_sq, previous_grad = ( + state["exp_avg"], + state["exp_avg_sq"], + state["previous_grad"], + ) + + state["step"] += 1 + + if group["weight_decay"] != 0: + grad.add_(p.data, alpha=group["weight_decay"]) + + # Decay the first and second moment running average coefficient + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + denom = exp_avg_sq.sqrt().add_(group["eps"]) + + bias_correction1 = 1 - beta1 ** state["step"] + bias_correction2 = 1 - beta2 ** state["step"] + + # compute diffgrad coefficient (dfc) + diff = torch.abs(previous_grad - grad) + dfc = torch.div(1.0, (1.0 + torch.exp(-diff))) + state["previous_grad"] = grad.clone() + + # update momentum with dfc + exp_avg1 = exp_avg * dfc + + step_size = ( + group["lr"] + * math.sqrt(bias_correction2) + / bias_correction1 + ) + + p.data.addcdiv_(exp_avg1, denom, value=-step_size) + + return loss diff --git a/torch_optimizer/lamb.py b/torch_optimizer/lamb.py new file mode 100644 index 00000000..7e948b05 --- /dev/null +++ b/torch_optimizer/lamb.py @@ -0,0 +1,158 @@ +import math + +import torch +from torch.optim.optimizer import Optimizer + +from .types import Betas2, OptFloat, OptLossClosure, Params + +__all__ = ("Lamb",) + + +class Lamb(Optimizer): + r"""Implements Lamb algorithm. + + It has been proposed in `Large Batch Optimization for Deep Learning: + Training BERT in 76 minutes`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + betas: coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps: term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay: weight decay (L2 penalty) (default: 0) + clamp_value: clamp weight_norm in (0,clamp_value) (default: 10) + set to a high value to avoid it (e.g 10e3) + adam: always use trust ratio = 1, which turns this + into Adam. Useful for comparison purposes. (default: False) + debias: debias adam by (1 - beta**step) (default: False) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.Lamb(model.parameters(), lr=0.1) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1904.00962 + + Note: + Reference code: https://github.com/cybertronai/pytorch-lamb + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + betas: Betas2 = (0.9, 0.999), + eps: float = 1e-6, + weight_decay: float = 0, + clamp_value: float = 10, + adam: bool = False, + debias: bool = False, + ) -> None: + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + "Invalid beta parameter at index 0: {}".format(betas[0]) + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if weight_decay < 0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + if clamp_value < 0.0: + raise ValueError("Invalid clamp value: {}".format(clamp_value)) + + defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay) + self.clamp_value = clamp_value + self.adam = adam + self.debias = debias + + super(Lamb, self).__init__(params, defaults) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + if grad.is_sparse: + msg = ( + "Lamb does not support sparse gradients, " + "please consider SparseAdam instead" + ) + raise RuntimeError(msg) + + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = 0 + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + beta1, beta2 = group["betas"] + + state["step"] += 1 + + # Decay the first and second moment running average coefficient + # m_t + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + # v_t + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + + # Paper v3 does not use debiasing. + if self.debias: + bias_correction = math.sqrt(1 - beta2 ** state["step"]) + bias_correction /= 1 - beta1 ** state["step"] + else: + bias_correction = 1 + + # Apply bias to lr to avoid broadcast. + step_size = group["lr"] * bias_correction + + weight_norm = torch.norm(p.data).clamp(0, self.clamp_value) + + adam_step = exp_avg / exp_avg_sq.sqrt().add(group["eps"]) + if group["weight_decay"] != 0: + adam_step.add_(p.data, alpha=group["weight_decay"]) + + adam_norm = torch.norm(adam_step) + if weight_norm == 0 or adam_norm == 0: + trust_ratio = 1 + else: + trust_ratio = weight_norm / adam_norm + state["weight_norm"] = weight_norm + state["adam_norm"] = adam_norm + state["trust_ratio"] = trust_ratio + if self.adam: + trust_ratio = 1 + + p.data.add_(adam_step, alpha=-step_size * trust_ratio) + + return loss diff --git a/torch_optimizer/lars.py b/torch_optimizer/lars.py new file mode 100644 index 00000000..cf65b2d7 --- /dev/null +++ b/torch_optimizer/lars.py @@ -0,0 +1,166 @@ +import torch +from torch.optim.optimizer import Optimizer + +from .types import OptFloat, OptLossClosure, Params, State + +__all__ = ("LARS",) + + +class LARS(Optimizer): + r"""Extends SGD in PyTorch with LARS scaling from the paper + `Large batch training of Convolutional Networks`__. + + Arguments: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + momentum: momentum factor (default: 0) + dampening: dampening for momentum (default: 0) + eps: term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay: weight decay (L2 penalty) (default: 0) + nesterov: enables Nesterov momentum (default: False) + trust_coefficient: trust coefficient for computing LR (default: 0.001) + eps: eps for division denominator (default: 1e-8) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.LARS(model.parameters(), lr=0.001) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + .. note:: + The application of momentum in the SGD part is modified according to + the PyTorch standards. LARS scaling fits into the equation in the + following fashion. + + .. math:: + \begin{aligned} + g_{t+1} & = \text{lars_lr} * (\beta * p_{t} + g_{t+1}), \\ + v_{t+1} & = \\mu * v_{t} + g_{t+1}, \\ + p_{t+1} & = p_{t} - \text{lr} * v_{t+1}, + \\end{aligned} + + where :math:`p`, :math:`g`, :math:`v`, :math:`\\mu` and :math:`\beta` + denote the parameters, gradient, velocity, momentum, and weight decay + respectively. The :math:`lars_lr` is defined by Eq. 6 in the paper. + The Nesterov version is analogously modified. + + .. warning:: + Parameters with weight decay set to 0 will automatically be excluded + from layer-wise LR scaling. This is to ensure consistency with papers + like SimCLR and BYOL. + + + __ https://arxiv.org/pdf/1708.03888.pdf + + Note: + Reference code: https://github.com/PyTorchLightning/lightning-bolts/ + """ + + def __init__( + self, + params: Params, + lr: float = 1e-2, + momentum: float = 0.0, + dampening: float = 0.0, + weight_decay: float = 0.0, + nesterov: bool = False, + trust_coefficient: float = 0.01, + eps: float = 1e-8, + ): + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if momentum < 0.0: + raise ValueError("Invalid momentum value: {}".format(momentum)) + if dampening < 0.0: + raise ValueError("Invalid dampening value: {}".format(dampening)) + if weight_decay < 0.0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + if trust_coefficient < 0.0: + raise ValueError( + "Invalid trust_coefficient value: {}".format(trust_coefficient) + ) + + defaults = dict( + lr=lr, + momentum=momentum, + dampening=dampening, + weight_decay=weight_decay, + nesterov=nesterov, + trust_coefficient=trust_coefficient, + eps=eps, + ) + if nesterov and (momentum <= 0 or dampening != 0): + raise ValueError( + "Nesterov momentum requires a momentum and zero dampening" + ) + + super().__init__(params, defaults) + + def __setstate__(self, state: State) -> None: + super().__setstate__(state) + + for group in self.param_groups: + group.setdefault("nesterov", False) + + @torch.no_grad() + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + # exclude scaling for params with 0 weight decay + for group in self.param_groups: + weight_decay = group["weight_decay"] + momentum = group["momentum"] + dampening = group["dampening"] + nesterov = group["nesterov"] + + for p in group["params"]: + if p.grad is None: + continue + + d_p = p.grad + p_norm = torch.norm(p.data) + g_norm = torch.norm(p.grad.data) + + # lars scaling + weight decay part + if weight_decay != 0: + if p_norm != 0 and g_norm != 0: + lars_lr = p_norm / ( + g_norm + p_norm * weight_decay + group["eps"] + ) + lars_lr *= group["trust_coefficient"] + + d_p = d_p.add(p, alpha=weight_decay) + d_p *= lars_lr + + if momentum != 0: + param_state = self.state[p] + if "momentum_buffer" not in param_state: + buf = param_state["momentum_buffer"] = torch.clone( + d_p + ).detach() + else: + buf = param_state["momentum_buffer"] + buf.mul_(momentum).add_(d_p, alpha=1 - dampening) + if nesterov: + d_p = d_p.add(buf, alpha=momentum) + else: + d_p = buf + + p.add_(d_p, alpha=-group["lr"]) + + return loss diff --git a/torch_optimizer/lion.py b/torch_optimizer/lion.py new file mode 100644 index 00000000..3d03739c --- /dev/null +++ b/torch_optimizer/lion.py @@ -0,0 +1,100 @@ +import torch +from torch.optim.optimizer import Optimizer + +from .types import Betas2, OptFloat, OptLossClosure, Params + +__all__ = ("Lion",) + + +class Lion(Optimizer): + r"""Implements Lion algorithm. + + Addapted from https://github.com/google/automl/tree/master/lion + + The Lion - EvoLved SIgn MOmeNtum - algorithm was proposed in + https://arxiv.org/pdf/2302.06675.pdf. + Lion aims to be more memory efficient than Adam by only tracking momentum. + + Caveats: As detailed in the paper, Lion requires a smaller learning rate + lr, and larger decoupled weight decay to maintain effective weight decay + strength. Also, the gain of Lion increases with the batch size. + Furthermore, Lion was not found to outperform AdamW on some large language + and text/image datasets. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + betas: coefficients used for computing + running averages of gradient and its square (default: (0.95, 0)) + weight_decay: weight decay (L2 penalty) (default: 0) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.Lion(model.parameters(), lr=0.001) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + """ + + def __init__( + self, + params: Params, + lr: float = 1e-4, + betas: Betas2 = (0.9, 0.99), + weight_decay: float = 0.0, + ): + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + "Invalid beta parameter at index 0: {}".format(betas[0]) + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if weight_decay < 0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay) + super().__init__(params, defaults) + + @torch.no_grad() + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + + # Perform stepweight decay + p.data.mul_(1 - group["lr"] * group["weight_decay"]) + + grad = p.grad + state = self.state[p] + # State initialization + if len(state) == 0: + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like(p) + + exp_avg = state["exp_avg"] + beta1, beta2 = group["betas"] + + # Weight update + update = exp_avg * beta1 + grad * (1 - beta1) + p.add_(torch.sign(update), alpha=-group["lr"]) + # Decay the momentum running average coefficient + exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2) + + return loss diff --git a/torch_optimizer/lookahead.py b/torch_optimizer/lookahead.py index 3c3b65cd..53afe3f0 100644 --- a/torch_optimizer/lookahead.py +++ b/torch_optimizer/lookahead.py @@ -1,11 +1,50 @@ from collections import defaultdict +from typing import Any, Dict import torch -from torch.optim import Optimizer +from torch.optim.optimizer import Optimizer + +from .types import OptFloat, OptLossClosure, State + +__all__ = ("Lookahead",) class Lookahead(Optimizer): - def __init__(self, optimizer, k=5, alpha=0.5): + r"""Implements Lookahead optimization algorithm. + + It has been proposed in `Lookahead Optimizer: k steps forward, 1 + step back`__ + + Arguments: + optimizer: base inner optimizer optimize, like Yogi, DiffGrad or Adam. + k: number of lookahead steps (default: 5) + alpha: linear interpolation factor. 1.0 recovers the inner optimizer. + (default: 5) + + Example: + >>> import torch_optimizer as optim + >>> yogi = optim.Yogi(model.parameters(), lr=0.1) + >>> optimizer = optim.Lookahead(yogi, k=5, alpha=0.5) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1907.08610 + + Note: + Reference code: https://github.com/alphadl/lookahead.pytorch + """ + + def __init__( + self, optimizer: Optimizer, k: int = 5, alpha: float = 0.5 + ) -> None: + if k < 0.0: + raise ValueError("Invalid number of lookahead steps: {}".format(k)) + if alpha < 0: + raise ValueError( + "Invalid linear interpolation factor: {}".format(alpha) + ) + self.optimizer = optimizer self.k = k self.alpha = alpha @@ -13,59 +52,84 @@ def __init__(self, optimizer, k=5, alpha=0.5): self.state = defaultdict(dict) self.fast_state = self.optimizer.state for group in self.param_groups: - group['counter'] = 0 + group["counter"] = 0 + self.defaults = {"k": k, "alpha": alpha, **optimizer.defaults} - def update(self, group): - for fast in group['params']: + def _update(self, group: Dict[str, Any]) -> None: + for fast in group["params"]: + if not fast.requires_grad: + continue + param_state = self.state[fast] - if 'slow_param' not in param_state: - param_state['slow_param'] = torch.zeros_like(fast.data) - param_state['slow_param'].copy_(fast.data) - slow = param_state['slow_param'] - slow += (fast.data - slow) * self.alpha - fast.data.copy_(slow) - - def update_lookahead(self): - for group in self.param_groups: - self.update(group) + if "slow_param" not in param_state: + param_state["slow_param"] = torch.clone(fast.data).detach() + + slow = param_state["slow_param"] + fast.data.mul_(self.alpha).add_(slow, alpha=1.0 - self.alpha) + slow.data.copy_(fast) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. - def step(self, closure=None): - loss = self.optimizer.step(closure) + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = self.optimizer.step(closure=closure) for group in self.param_groups: - if group['counter'] == 0: - self.update(group) - group['counter'] += 1 - if group['counter'] >= self.k: - group['counter'] = 0 + if group["counter"] == 0: + self._update(group) + group["counter"] += 1 + group["counter"] %= self.k return loss - def state_dict(self): + def state_dict(self) -> State: + r"""Returns the state of the optimizer as a :class:`dict`. + + It contains two entries: + * state - a dict holding current optimization state. Its content + differs between optimizer classes. + * param_groups - a dict containing all parameter groups + """ + slow_state_dict = super(Lookahead, self).state_dict() fast_state_dict = self.optimizer.state_dict() - slow_state = { - (id(k) if isinstance(k, torch.Tensor) else k): v - for k, v in self.state.items() - } - fast_state = fast_state_dict['state'] - param_groups = fast_state_dict['param_groups'] + fast_state = fast_state_dict["state"] + param_groups = fast_state_dict["param_groups"] return { - 'fast_state': fast_state, - 'slow_state': slow_state, - 'param_groups': param_groups, + "fast_state": fast_state, + "slow_state": slow_state_dict["state"], + "param_groups": param_groups, } - def load_state_dict(self, state_dict): + def load_state_dict(self, state_dict: State) -> None: + r"""Loads the optimizer state. + + Arguments: + state_dict: optimizer state. Should be an object returned + from a call to :meth:`state_dict`. + """ slow_state_dict = { - 'state': state_dict['slow_state'], - 'param_groups': state_dict['param_groups'], + "state": state_dict["slow_state"], + "param_groups": state_dict["param_groups"], } fast_state_dict = { - 'state': state_dict['fast_state'], - 'param_groups': state_dict['param_groups'], + "state": state_dict["fast_state"], + "param_groups": state_dict["param_groups"], } super(Lookahead, self).load_state_dict(slow_state_dict) self.optimizer.load_state_dict(fast_state_dict) self.fast_state = self.optimizer.state - def add_param_group(self, param_group): - param_group['counter'] = 0 - self.optimizer.add_param_group(param_group) + def zero_grad(self, set_to_none: bool = False) -> None: + r"""Clears the gradients of all optimized :class:`torch.Tensor` s.""" + self.optimizer.zero_grad(set_to_none) + + def __repr__(self) -> str: + base_str = self.optimizer.__repr__() + format_string = self.__class__.__name__ + " (" + format_string += "\n" + format_string += "k: {}\n".format(self.k) + format_string += "alpha: {}\n".format(self.alpha) + format_string += base_str + format_string += "\n" + format_string += ")" + return format_string diff --git a/torch_optimizer/madgrad.py b/torch_optimizer/madgrad.py new file mode 100644 index 00000000..3dd8d6e8 --- /dev/null +++ b/torch_optimizer/madgrad.py @@ -0,0 +1,185 @@ +import math +from typing import Callable, Optional + +import torch +import torch.optim + +from .types import Params + +__all__ = ("MADGRAD",) + + +class MADGRAD(torch.optim.Optimizer): + r"""Implements MADGRAD algorithm. + + It has been proposed in `Adaptivity without Compromise: A Momentumized, + Adaptive, Dual Averaged Gradient Method for Stochastic Optimization`__ + + Arguments: + params (iterable): + Iterable of parameters to optimize + or dicts defining parameter groups. + lr (float): + Learning rate (default: 1e-2). + momentum (float): + Momentum value in the range [0,1) (default: 0.9). + weight_decay (float): + Weight decay, i.e. a L2 penalty (default: 0). + eps (float): + Term added to the denominator outside of the root operation + to improve numerical stability. (default: 1e-6). + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.MAGRAD(model.parameters(), lr=0.1) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/2101.11075 + + Note: + Reference code: https://github.com/facebookresearch/madgrad + """ + + def __init__( + self, + params: Params, + lr: float = 1e-2, + momentum: float = 0.9, + weight_decay: float = 0.0, + eps: float = 1e-6, + ): + if momentum < 0 or momentum >= 1: + raise ValueError("Invalid momentum value: {}".format(momentum)) + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if weight_decay < 0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + + defaults = dict( + lr=lr, eps=eps, momentum=momentum, weight_decay=weight_decay, k=0 + ) + super().__init__(params, defaults) + + for group in self.param_groups: + for p in group["params"]: + state = self.state[p] + + state["grad_sum_sq"] = torch.zeros_like(p.data).detach() + state["s"] = torch.zeros_like(p.data).detach() + if momentum != 0: + state["x0"] = torch.clone(p.data).detach() + + def step( + self, closure: Optional[Callable[[], float]] = None + ) -> Optional[float]: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + eps = group["eps"] + k = group["k"] + lr = group["lr"] + eps + decay = group["weight_decay"] + momentum = group["momentum"] + + ck = 1 - momentum + lamb = lr * math.pow(k + 1, 0.5) + + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + state = self.state[p] + + if momentum != 0.0 and grad.is_sparse: + raise RuntimeError( + "momentum != 0 is not compatible with " + "sparse gradients" + ) + + grad_sum_sq = state["grad_sum_sq"] + s = state["s"] + + # Apply weight decay + if decay != 0: + if grad.is_sparse: + raise RuntimeError( + "weight_decay option is not " + "compatible with sparse gradients" + ) + + grad.add_(p.data, alpha=decay) + + if grad.is_sparse: + grad = grad.coalesce() + grad_val = grad._values() + + p_masked = p.sparse_mask(grad) + grad_sum_sq_masked = grad_sum_sq.sparse_mask(grad) + s_masked = s.sparse_mask(grad) + + # Compute x_0 from other known quantities + rms_masked_vals = ( + grad_sum_sq_masked._values().pow(1 / 3).add_(eps) + ) + x0_masked_vals = p_masked._values().addcdiv( + s_masked._values(), rms_masked_vals, value=1 + ) + + # Dense + sparse op + grad_sq = grad * grad + grad_sum_sq.add_(grad_sq, alpha=lamb) + grad_sum_sq_masked.add_(grad_sq, alpha=lamb) + + rms_masked_vals = ( + grad_sum_sq_masked._values().pow_(1 / 3).add_(eps) + ) + + s.add_(grad, alpha=lamb) + s_masked._values().add_(grad_val, alpha=lamb) + + # update masked copy of p + p_kp1_masked_vals = x0_masked_vals.addcdiv( + s_masked._values(), rms_masked_vals, value=-1 + ) + # Copy updated masked p to dense p using an add operation + p_masked._values().add_(p_kp1_masked_vals, alpha=-1) + p.data.add_(p_masked, alpha=-1) + else: + if momentum == 0: + # Compute x_0 from other known quantities + rms = grad_sum_sq.pow(1 / 3).add_(eps) + x0 = p.data.addcdiv(s, rms, value=1) + else: + x0 = state["x0"] + + # Accumulate second moments + grad_sum_sq.addcmul_(grad, grad, value=lamb) + rms = grad_sum_sq.pow(1 / 3).add_(eps) + + # Update s + s.data.add_(grad, alpha=lamb) + + # Step + if momentum == 0: + p.data.copy_(x0.addcdiv(s, rms, value=-1)) + else: + z = x0.addcdiv(s, rms, value=-1) + + # p is a moving average of z + p.data.mul_(1 - ck).add_(z, alpha=ck) + + group["k"] = group["k"] + 1 + return loss diff --git a/torch_optimizer/novograd.py b/torch_optimizer/novograd.py new file mode 100644 index 00000000..6316de94 --- /dev/null +++ b/torch_optimizer/novograd.py @@ -0,0 +1,161 @@ +import torch +from torch.optim.optimizer import Optimizer + +from .types import Betas2, OptFloat, OptLossClosure, Params + +__all__ = ("NovoGrad",) + + +class NovoGrad(Optimizer): + r"""Implements Novograd optimization algorithm. + + It has been proposed in `Stochastic Gradient Methods with Layer-wise + Adaptive Moments for Training of Deep Networks`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + betas: coefficients used for computing + running averages of gradient and its square (default: (0.95, 0)) + eps: term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay: weight decay (L2 penalty) (default: 0) + grad_averaging: gradient averaging (default: False) + amsgrad: whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond` + (default: False) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.Yogi(model.parameters(), lr=0.1) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> scheduler = StepLR(optimizer, step_size=1, gamma=0.7) + >>> optimizer.step() + >>> scheduler.step() + + __ https://arxiv.org/abs/1905.11286 + + Note: + Reference code: https://github.com/NVIDIA/DeepLearningExamples + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + betas: Betas2 = (0.95, 0), + eps: float = 1e-8, + weight_decay: float = 0, + grad_averaging: bool = False, + amsgrad: bool = False, + ): + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + "Invalid beta parameter at index 0: {}".format(betas[0]) + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if weight_decay < 0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + defaults = dict( + lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + grad_averaging=grad_averaging, + amsgrad=amsgrad, + ) + + super(NovoGrad, self).__init__(params, defaults) + + def __setstate__(self, state: dict) -> None: + super(NovoGrad, self).__setstate__(state) + for group in self.param_groups: + group.setdefault("amsgrad", False) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + if grad.is_sparse: + msg = ( + "NovoGrad does not support sparse gradients, " + "please consider SparseAdam instead" + ) + raise RuntimeError(msg) + amsgrad = group["amsgrad"] + + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = 0 + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + p.data, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros([]).to( + state["exp_avg"].device + ) + if amsgrad: + # Maintains max of all exp. moving avg. of sq. + # grad. values + state["max_exp_avg_sq"] = torch.zeros([]).to( + state["exp_avg"].device + ) + + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + if amsgrad: + max_exp_avg_sq = state["max_exp_avg_sq"] + beta1, beta2 = group["betas"] + + state["step"] += 1 + + norm = torch.sum(torch.pow(grad, 2)) + + if exp_avg_sq == 0: + exp_avg_sq.copy_(norm) + else: + exp_avg_sq.mul_(beta2).add_(norm, alpha=1 - beta2) + + if amsgrad: + # Maintains the maximum of all 2nd moment running avg. + # till now + torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq) + # Use the max. for normalizing running avg. of gradient + denom = max_exp_avg_sq.sqrt().add_(group["eps"]) + else: + denom = exp_avg_sq.sqrt().add_(group["eps"]) + + grad.div_(denom) + if group["weight_decay"] != 0: + grad.add_(p.data, alpha=group["weight_decay"]) + if group["grad_averaging"]: + grad.mul_(1 - beta1) + exp_avg.mul_(beta1).add_(grad) + + p.data.add_(exp_avg, alpha=-group["lr"]) + + return loss diff --git a/torch_optimizer/pid.py b/torch_optimizer/pid.py new file mode 100644 index 00000000..44dad253 --- /dev/null +++ b/torch_optimizer/pid.py @@ -0,0 +1,125 @@ +import torch +from torch.optim.optimizer import Optimizer + +from .types import OptFloat, OptLossClosure, Params + + +class PID(Optimizer): + r"""Implements PID optimization algorithm. + + It has been proposed in `A PID Controller Approach for Stochastic + Optimization of Deep Networks`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + momentum: momentum factor (default: 0.0) + weight_decay: weight decay (L2 penalty) (default: 0.0) + dampening: dampening for momentum (default: 0.0) + derivative: D part of the PID (default: 10.0) + integral: I part of the PID (default: 5.0) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.PID(model.parameters(), lr=0.001, momentum=0.1) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ http://www4.comp.polyu.edu.hk/~cslzhang/paper/CVPR18_PID.pdf + + Note: + Reference code: https://github.com/tensorboy/PIDOptimizer + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + momentum: float = 0.0, + dampening: float = 0, + weight_decay: float = 0.0, + integral: float = 5.0, + derivative: float = 10.0, + ) -> None: + defaults = dict( + lr=lr, + momentum=momentum, + dampening=dampening, + weight_decay=weight_decay, + integral=integral, + derivative=derivative, + ) + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if momentum < 0.0: + raise ValueError("Invalid momentum value: {}".format(momentum)) + if weight_decay < 0.0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + if integral < 0.0: + raise ValueError("Invalid PID integral value: {}".format(integral)) + if derivative < 0.0: + raise ValueError( + "Invalid PID derivative value: {}".format(derivative) + ) + + super(PID, self).__init__(params, defaults) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + weight_decay = group["weight_decay"] + momentum = group["momentum"] + dampening = group["dampening"] + integral = group["integral"] + derivative = group["derivative"] + for p in group["params"]: + if p.grad is None: + continue + d_p = p.grad.data + if weight_decay != 0: + d_p.add_(p.data, alpha=weight_decay) + if momentum != 0: + param_state = self.state[p] + if "i_buffer" not in param_state: + i_buf = param_state["i_buffer"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + i_buf.mul_(momentum).add_(d_p) + else: + i_buf = param_state["i_buffer"] + i_buf.mul_(momentum).add_(d_p, alpha=1 - dampening) + if "grad_buffer" not in param_state: + g_buf = param_state["grad_buffer"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + g_buf = d_p + + d_buf = param_state["d_buffer"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + d_buf.mul_(momentum).add_(d_p - g_buf) + else: + d_buf = param_state["d_buffer"] + g_buf = param_state["grad_buffer"] + d_buf.mul_(momentum).add_( + d_p - g_buf, alpha=1 - momentum + ) + self.state[p]["grad_buffer"] = d_p.clone() + + d_p = d_p.add_(i_buf, alpha=integral).add_( + d_buf, alpha=derivative + ) + p.data.add_(d_p, alpha=-group["lr"]) + return loss diff --git a/torch_optimizer/powersign.py b/torch_optimizer/powersign.py deleted file mode 100644 index c5ab10cf..00000000 --- a/torch_optimizer/powersign.py +++ /dev/null @@ -1,85 +0,0 @@ -import torch -from torch.optim import Optimizer - - -class PowerSign(Optimizer): - def __init__( - self, - params, - lr=0.001, - momentum=0.9, - dampening=0, - weight_decay=0, - nesterov=False, - ): - defaults = dict( - lr=lr, - momentum=momentum, - dampening=dampening, - weight_decay=weight_decay, - nesterov=nesterov, - ) - if nesterov and (momentum <= 0 or dampening != 0): - raise ValueError( - 'Nesterov momentum requires a momentum and zero dampening' - ) - super(PowerSign, self).__init__(params, defaults) - - def __setstate__(self, state): - super(PowerSign, self).__setstate__(state) - for group in self.param_groups: - group.setdefault('nesterov', False) - - def step(self, closure=None): - """Performs a single optimization step. - Arguments: - closure (callable, optional): A closure that reevaluates the model - and returns the loss. - """ - loss = None - if closure is not None: - loss = closure() - - for group in self.param_groups: - weight_decay = group['weight_decay'] - momentum = group['momentum'] - dampening = group['dampening'] - nesterov = group['nesterov'] - - # print(weight_decay, momentum, dampening , nesterov) - # They use pytorch functions most likely to speed up the operations - # More research is defiently needed to understand how optimisers - # are implemented in pytorch. - - for p in group['params']: - if p.grad is None: - continue - d_p = p.grad.data - if weight_decay != 0: - d_p.add_(weight_decay, p.data) - if momentum != 0: - param_state = self.state[p] - if 'momentum_buffer' not in param_state: - buf = param_state['momentum_buffer'] = torch.zeros( - p.data.size() - ) - buf = torch.add( - buf.mul(momentum), (d_p.mul(1 - momentum)) - ) - else: - buf = param_state['momentum_buffer'] - buf = buf.mul(momentum).add(1 - dampening, d_p) - if nesterov: - d_p = d_p.add(momentum, buf) - else: - # This is the gradient update rule that was found in - # the paper Neural Optimizer search with reinfrocmement - # learning - d_p = torch.mul( - torch.exp(torch.mul(d_p.sign(), buf.sign())), d_p - ) - # print(d_p) - # d_p = buf - - p.data.add_(-group['lr'], d_p) - return loss diff --git a/torch_optimizer/py.typed b/torch_optimizer/py.typed new file mode 100644 index 00000000..fdffa2a0 --- /dev/null +++ b/torch_optimizer/py.typed @@ -0,0 +1 @@ +# placeholder diff --git a/torch_optimizer/qhadam.py b/torch_optimizer/qhadam.py new file mode 100644 index 00000000..751b80d1 --- /dev/null +++ b/torch_optimizer/qhadam.py @@ -0,0 +1,153 @@ +import torch +from torch.optim.optimizer import Optimizer + +from .types import Betas2, Nus2, OptFloat, OptLossClosure, Params + +__all__ = ("QHAdam",) + + +class QHAdam(Optimizer): + r"""Implements the QHAdam optimization algorithm. + + It has been proposed in `Adaptive methods for Nonconvex Optimization`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + betas: coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + nus: immediate discount factors used to estimate the gradient and its + square (default: (1.0, 1.0)) + eps: term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay: weight decay (L2 penalty) (default: 0) + decouple_weight_decay: whether to decouple the weight + decay from the gradient-based optimization step (default: False) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.QHAdam(model.parameters(), lr=0.1) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1810.06801 + + Note: + Reference code: https://github.com/facebookresearch/qhoptim + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + betas: Betas2 = (0.9, 0.999), + nus: Nus2 = (1.0, 1.0), + weight_decay: float = 0.0, + decouple_weight_decay: bool = False, + eps: float = 1e-8, + ): + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + "Invalid beta parameter at index 0: {}".format(betas[0]) + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if weight_decay < 0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + + defaults = { + "lr": lr, + "betas": betas, + "nus": nus, + "weight_decay": weight_decay, + "decouple_weight_decay": decouple_weight_decay, + "eps": eps, + } + super(QHAdam, self).__init__(params, defaults) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + """Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + lr = group["lr"] + beta1, beta2 = group["betas"] + nu1, nu2 = group["nus"] + weight_decay = group["weight_decay"] + decouple_weight_decay = group["decouple_weight_decay"] + eps = group["eps"] + + for p in group["params"]: + if p.grad is None: + continue + + d_p = p.grad.data + if d_p.is_sparse: + raise RuntimeError( + "QHAdam does not support sparse gradients, " + "please consider SparseAdam instead" + ) + + state = self.state[p] + + if weight_decay != 0: + if decouple_weight_decay: + p.data.mul_(1 - lr * weight_decay) + else: + d_p.add_(p.data, alpha=weight_decay) + + d_p_sq = d_p.mul(d_p) + + if len(state) == 0: + state["beta1_weight"] = 0.0 + state["beta2_weight"] = 0.0 + state["exp_avg"] = torch.zeros_like( + p.data, memory_format=torch.preserve_format + ) + state["exp_avg_sq"] = torch.zeros_like( + p.data, memory_format=torch.preserve_format + ) + + state["beta1_weight"] = 1.0 + beta1 * state["beta1_weight"] + state["beta2_weight"] = 1.0 + beta2 * state["beta2_weight"] + + beta1_weight = state["beta1_weight"] + beta2_weight = state["beta2_weight"] + exp_avg = state["exp_avg"] + exp_avg_sq = state["exp_avg_sq"] + + beta1_adj = 1.0 - (1.0 / beta1_weight) + beta2_adj = 1.0 - (1.0 / beta2_weight) + exp_avg.mul_(beta1_adj).add_(d_p, alpha=1.0 - beta1_adj) + exp_avg_sq.mul_(beta2_adj).add_(d_p_sq, alpha=1.0 - beta2_adj) + + avg_grad = exp_avg.mul(nu1) + if nu1 != 1.0: + avg_grad.add_(d_p, alpha=1.0 - nu1) + + avg_grad_rms = exp_avg_sq.mul(nu2) + if nu2 != 1.0: + avg_grad_rms.add_(d_p_sq, alpha=1.0 - nu2) + avg_grad_rms.sqrt_() + if eps != 0.0: + avg_grad_rms.add_(eps) + + p.data.addcdiv_(avg_grad, avg_grad_rms, value=-lr) + + return loss diff --git a/torch_optimizer/qhm.py b/torch_optimizer/qhm.py new file mode 100644 index 00000000..4018465c --- /dev/null +++ b/torch_optimizer/qhm.py @@ -0,0 +1,117 @@ +import torch +from torch.optim.optimizer import Optimizer + +from .types import OptFloat, OptLossClosure, Params + +__all__ = ("QHM",) + + +class QHM(Optimizer): + GRAD = "grad" + DIRECT = "direct" + + r"""Implements quasi-hyperbolic momentum (QHM) optimization algorithm. + + It has been proposed in `Quasi-hyperbolic momentum and Adam for deep + learning`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + momentum: momentum factor (:math:`\beta` from the paper) + nu: immediate discount factor (:math:`\nu` from the paper) + weight_decay: weight decay (L2 regularization coefficient, times two) + (default: 0.0) + weight_decay_type: method of applying the weight decay: + ``"grad"`` for accumulation in the gradient + (same as :class:`torch.optim.SGD`) or + ``"direct"`` for direct application to the parameters + (default: ``"grad"``) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.QHM(model.parameters(), lr=0.1, momentum=0.9) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + + __ https://arxiv.org/abs/1810.06801 + + Note: + Reference code: https://github.com/facebookresearch/qhoptim + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + momentum: float = 0.0, + nu: float = 0.7, + weight_decay: float = 0.0, + weight_decay_type: str = "grad", + ) -> None: + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if momentum < 0.0: + raise ValueError("Invalid momentum value: {}".format(momentum)) + if weight_decay < 0.0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + if weight_decay_type not in (self.GRAD, self.DIRECT): + _type = weight_decay_type + msg = "Invalid weight_decay_type value: {}".format(_type) + raise ValueError(msg) + + defaults = { + "lr": lr, + "momentum": momentum, + "nu": nu, + "weight_decay": weight_decay, + "weight_decay_type": weight_decay_type, + } + super(QHM, self).__init__(params, defaults) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + """Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + lr, nu, momentum = group["lr"], group["nu"], group["momentum"] + weight_decay, weight_decay_type = ( + group["weight_decay"], + group["weight_decay_type"], + ) + + for p in group["params"]: + if p.grad is None: + continue + d_p = p.grad.data + param_state = self.state[p] + + if weight_decay != 0: + if weight_decay_type == self.GRAD: + d_p.add_(p.data, alpha=weight_decay) + else: + p.data.mul_(1.0 - lr * weight_decay) + + if len(param_state) == 0: + param_state["momentum_buffer"] = torch.zeros_like( + p.data, memory_format=torch.preserve_format + ) + + momentum_buffer = param_state["momentum_buffer"] + momentum_buffer.mul_(momentum).add_(d_p, alpha=1.0 - momentum) + + p.data.add_(momentum_buffer, alpha=-lr * nu) + p.data.add_(d_p, alpha=-lr * (1.0 - nu)) + + return loss diff --git a/torch_optimizer/radam.py b/torch_optimizer/radam.py new file mode 100644 index 00000000..d25e158e --- /dev/null +++ b/torch_optimizer/radam.py @@ -0,0 +1,194 @@ +import math +import warnings + +import torch +from torch.optim.optimizer import Optimizer + +from .types import Betas2, OptFloat, OptLossClosure, Params + +__all__ = ("RAdam",) + + +class RAdam(Optimizer): + r"""Implements RAdam optimization algorithm. + + Note: + Deprecated, please use version provided by PyTorch_. + + It has been proposed in `On the Variance of the Adaptive Learning + Rate and Beyond`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + betas: coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps: term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay: weight decay (L2 penalty) (default: 0) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.RAdam(model.parameters(), lr=0.1) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1908.03265 + + Note: + Reference code: https://github.com/LiyuanLucasLiu/RAdam + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + betas: Betas2 = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 0, + ) -> None: + warnings.warn( + "RAdam optimizer is deprecated, since it is included " + "in pytorch natively.", + DeprecationWarning, + stacklevel=2, + ) + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + "Invalid beta parameter at index 0: {}".format(betas[0]) + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if weight_decay < 0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + + if ( + isinstance(params, (list, tuple)) + and len(params) > 0 + and isinstance(params[0], dict) + ): + for param in params: + if "betas" in param and ( + param["betas"][0] != betas[0] + or param["betas"][1] != betas[1] + ): + param["buffer"] = [[None, None, None] for _ in range(10)] + + defaults = dict( + lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + buffer=[[None, None, None] for _ in range(10)], + ) + super(RAdam, self).__init__(params, defaults) + + def __setstate__(self, state): + super(RAdam, self).__setstate__(state) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + lr = group["lr"] + weight_decay = group["weight_decay"] + beta1, beta2 = group["betas"] + eps = group["eps"] + + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data.float() + if grad.is_sparse: + msg = ( + "RAdam does not support sparse gradients, " + "please consider SparseAdam instead" + ) + raise RuntimeError(msg) + + p_data_fp32 = p.data.float() + + state = self.state[p] + + if len(state) == 0: + state["step"] = 0 + state["exp_avg"] = torch.zeros_like( + p_data_fp32, memory_format=torch.preserve_format + ) + state["exp_avg_sq"] = torch.zeros_like( + p_data_fp32, memory_format=torch.preserve_format + ) + else: + state["exp_avg"] = state["exp_avg"].type_as(p_data_fp32) + state["exp_avg_sq"] = state["exp_avg_sq"].type_as( + p_data_fp32 + ) + + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + + state["step"] += 1 + buffered = group["buffer"][int(state["step"] % 10)] + if state["step"] == buffered[0]: + N_sma, step_size = buffered[1], buffered[2] + else: + buffered[0] = state["step"] + beta2_t = beta2 ** state["step"] + N_sma_max = 2 / (1 - beta2) - 1 + N_sma = N_sma_max - 2 * state["step"] * beta2_t / ( + 1 - beta2_t + ) + buffered[1] = N_sma + + # more conservative since it's an approximated value + if N_sma >= 5: + step_size = ( + lr + * math.sqrt( + (1 - beta2_t) + * (N_sma - 4) + / (N_sma_max - 4) + * (N_sma - 2) + / N_sma + * N_sma_max + / (N_sma_max - 2) + ) + / (1 - beta1 ** state["step"]) + ) + else: + step_size = lr / (1 - beta1 ** state["step"]) + buffered[2] = step_size + + if weight_decay != 0: + p_data_fp32.add_(p_data_fp32, alpha=-weight_decay * lr) + + # more conservative since it's an approximated value + if N_sma >= 5: + denom = exp_avg_sq.sqrt().add_(eps) + p_data_fp32.addcdiv_(exp_avg, denom, value=-step_size) + else: + p_data_fp32.add_(exp_avg, alpha=-step_size) + + p.data.copy_(p_data_fp32) + + return loss diff --git a/torch_optimizer/sgdp.py b/torch_optimizer/sgdp.py new file mode 100644 index 00000000..8b933d61 --- /dev/null +++ b/torch_optimizer/sgdp.py @@ -0,0 +1,187 @@ +import math + +import torch +from torch.optim.optimizer import Optimizer + +from .types import OptFloat, OptLossClosure, Params + +__all__ = ("SGDP",) + + +class SGDP(Optimizer): + r"""Implements SGDP algorithm. + + It has been proposed in `Slowing Down the Weight Norm Increase in + Momentum-based Optimizers`__ + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + momentum: momentum factor (default: 0) + dampening: dampening for momentum (default: 0) + eps: term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay: weight decay (L2 penalty) (default: 0) + delta: threhold that determines whether a set of parameters is scale + invariant or not (default: 0.1) + wd_ratio: relative weight decay applied on scale-invariant parameters + compared to that applied on scale-variant parameters (default: 0.1) + nesterov: enables Nesterov momentum (default: False) + + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.SGDP(model.parameters(), lr=0.1) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/2006.08217 + + Note: + Reference code: https://github.com/clovaai/AdamP + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + momentum: float = 0, + dampening: float = 0, + eps: float = 1e-8, + weight_decay: float = 0, + delta: float = 0.1, + wd_ratio: float = 0.1, + nesterov: bool = False, + ) -> None: + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if momentum < 0.0: + raise ValueError("Invalid momentum value: {}".format(momentum)) + if dampening < 0.0: + raise ValueError("Invalid dampening value: {}".format(dampening)) + if weight_decay < 0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + if delta < 0: + raise ValueError("Invalid delta value: {}".format(delta)) + if wd_ratio < 0: + raise ValueError("Invalid wd_ratio value: {}".format(wd_ratio)) + + defaults = dict( + lr=lr, + momentum=momentum, + dampening=dampening, + eps=eps, + weight_decay=weight_decay, + delta=delta, + wd_ratio=wd_ratio, + nesterov=nesterov, + ) + super(SGDP, self).__init__(params, defaults) + + @staticmethod + def _channel_view(x): + return x.view(x.size(0), -1) + + @staticmethod + def _layer_view(x): + return x.view(1, -1) + + @staticmethod + def _cosine_similarity(x, y, eps, view_func): + x = view_func(x) + y = view_func(y) + + x_norm = x.norm(dim=1).add_(eps) + y_norm = y.norm(dim=1).add_(eps) + dot = (x * y).sum(dim=1) + + return dot.abs() / x_norm / y_norm + + def _projection(self, p, grad, perturb, delta, wd_ratio, eps): + wd = 1 + expand_size = [-1] + [1] * (len(p.shape) - 1) + for view_func in [self._channel_view, self._layer_view]: + cosine_sim = self._cosine_similarity(grad, p.data, eps, view_func) + + if cosine_sim.max() < delta / math.sqrt(view_func(p.data).size(1)): + p_n = p.data / view_func(p.data).norm(dim=1).view( + expand_size + ).add_(eps) + perturb -= p_n * view_func(p_n * perturb).sum(dim=1).view( + expand_size + ) + wd = wd_ratio + + return perturb, wd + + return perturb, wd + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + weight_decay = group["weight_decay"] + momentum = group["momentum"] + dampening = group["dampening"] + nesterov = group["nesterov"] + + for p in group["params"]: + if p.grad is None: + continue + + grad = p.grad.data + state = self.state[p] + + # State initialization + if len(state) == 0: + state["momentum"] = torch.zeros_like( + p.data, memory_format=torch.preserve_format + ) + + # SGD + buf = state["momentum"] + buf.mul_(momentum).add_(grad, alpha=1 - dampening) + if nesterov: + d_p = grad + momentum * buf + else: + d_p = buf + + # Projection + wd_ratio = 1 + if len(p.shape) > 1: + d_p, wd_ratio = self._projection( + p, + grad, + d_p, + group["delta"], + group["wd_ratio"], + group["eps"], + ) + + # Weight decay + if weight_decay != 0: + p.data.mul_( + 1 + - group["lr"] + * group["weight_decay"] + * wd_ratio + / (1 - momentum) + ) + + # Step + p.data.add_(d_p, alpha=-group["lr"]) + + return loss diff --git a/torch_optimizer/sgdw.py b/torch_optimizer/sgdw.py new file mode 100644 index 00000000..62e7d998 --- /dev/null +++ b/torch_optimizer/sgdw.py @@ -0,0 +1,122 @@ +import torch +from torch.optim.optimizer import Optimizer + +from .types import OptFloat, OptLossClosure, Params, State + +__all__ = ("SGDW",) + + +class SGDW(Optimizer): + r"""Implements SGDW algorithm. + + It has been proposed in `Decoupled Weight Decay Regularization`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + momentum: momentum factor (default: 0) + weight_decay: weight decay (L2 penalty) (default: 0) + dampening: dampening for momentum (default: 0) + nesterov: enables Nesterov momentum (default: False) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.SGDW(model.parameters(), lr=0.1, momentum=0.9) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1711.05101 + + Note: + Reference code: https://github.com/pytorch/pytorch/pull/22466 + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + momentum: float = 0.0, + dampening: float = 0.0, + weight_decay: float = 0.0, + nesterov: bool = False, + ) -> None: + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if momentum < 0.0: + raise ValueError("Invalid momentum value: {}".format(momentum)) + if dampening < 0.0: + raise ValueError("Invalid dampening value: {}".format(dampening)) + if weight_decay < 0.0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + + defaults = dict( + lr=lr, + momentum=momentum, + dampening=dampening, + weight_decay=weight_decay, + nesterov=nesterov, + ) + if nesterov and (momentum <= 0 or dampening != 0): + raise ValueError( + "Nesterov momentum requires a momentum and zero dampening" + ) + super(SGDW, self).__init__(params, defaults) + + def __setstate__(self, state: State) -> None: + super(SGDW, self).__setstate__(state) + for group in self.param_groups: + group.setdefault("nesterov", False) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + """Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + weight_decay = group["weight_decay"] + momentum = group["momentum"] + dampening = group["dampening"] + nesterov = group["nesterov"] + + for p in group["params"]: + if p.grad is None: + continue + d_p = p.grad.data + + if p.grad.is_sparse: + msg = ( + "SGDW does not support sparse gradients, " + "please consider SparseAdam instead" + ) + raise RuntimeError(msg) + + if momentum != 0: + param_state = self.state[p] + if "momentum_buffer" not in param_state: + buf = param_state["momentum_buffer"] = torch.clone( + d_p + ).detach() + else: + buf = param_state["momentum_buffer"] + buf.mul_(momentum).add_(d_p, alpha=1 - dampening) + if nesterov: + d_p = d_p.add(momentum, buf) + else: + d_p = buf + + # Apply momentum + p.data.add_(d_p, alpha=-group["lr"]) + + # Apply weight decay + if weight_decay != 0: + p.data.add_(weight_decay, alpha=-group["lr"]) + return loss diff --git a/torch_optimizer/shampoo.py b/torch_optimizer/shampoo.py new file mode 100644 index 00000000..b64c59bc --- /dev/null +++ b/torch_optimizer/shampoo.py @@ -0,0 +1,146 @@ +import torch +from torch.optim.optimizer import Optimizer + +from .types import OptFloat, OptLossClosure, Params + + +def _matrix_power(matrix: torch.Tensor, power: float) -> torch.Tensor: + # use CPU for svd for speed up + device = matrix.device + matrix = matrix.cpu() + u, s, v = torch.svd(matrix) + return (u @ s.pow_(power).diag() @ v.t()).to(device) + + +class Shampoo(Optimizer): + r"""Implements Shampoo Optimizer Algorithm. + + It has been proposed in `Shampoo: Preconditioned Stochastic Tensor + Optimization`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-3) + momentum: momentum factor (default: 0) + weight_decay: weight decay (L2 penalty) (default: 0) + epsilon: epsilon added to each mat_gbar_j for numerical stability + (default: 1e-4) + update_freq: update frequency to compute inverse (default: 1) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.Shampoo(model.parameters(), lr=0.01) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/abs/1802.09568 + + Note: + Reference code: https://github.com/moskomule/shampoo.pytorch + """ + + def __init__( + self, + params: Params, + lr: float = 1e-1, + momentum: float = 0.0, + weight_decay: float = 0.0, + epsilon: float = 1e-4, + update_freq: int = 1, + ): + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if momentum < 0.0: + raise ValueError("Invalid momentum value: {}".format(momentum)) + if weight_decay < 0.0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + if epsilon < 0.0: + raise ValueError("Invalid momentum value: {}".format(momentum)) + if update_freq < 1: + raise ValueError("Invalid momentum value: {}".format(momentum)) + + defaults = dict( + lr=lr, + momentum=momentum, + weight_decay=weight_decay, + epsilon=epsilon, + update_freq=update_freq, + ) + super(Shampoo, self).__init__(params, defaults) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + """Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + order = grad.ndimension() + original_size = grad.size() + state = self.state[p] + momentum = group["momentum"] + weight_decay = group["weight_decay"] + if len(state) == 0: + state["step"] = 0 + if momentum > 0: + state["momentum_buffer"] = grad.clone() + for dim_id, dim in enumerate(grad.size()): + # precondition matrices + state["precond_{}".format(dim_id)] = group[ + "epsilon" + ] * torch.eye(dim, out=grad.new(dim, dim)) + state[ + "inv_precond_{dim_id}".format(dim_id=dim_id) + ] = grad.new(dim, dim).zero_() + + if momentum > 0: + grad.mul_(1 - momentum).add_( + state["momentum_buffer"], alpha=momentum + ) + + if weight_decay > 0: + grad.add_(p.data, alpha=group["weight_decay"]) + + # See Algorithm 2 for detail + for dim_id, dim in enumerate(grad.size()): + precond = state["precond_{}".format(dim_id)] + inv_precond = state["inv_precond_{}".format(dim_id)] + + # mat_{dim_id}(grad) + grad = grad.transpose_(0, dim_id).contiguous() + transposed_size = grad.size() + grad = grad.view(dim, -1) + + grad_t = grad.t() + precond.add_(grad @ grad_t) + if state["step"] % group["update_freq"] == 0: + inv_precond.copy_(_matrix_power(precond, -1 / order)) + + if dim_id == order - 1: + # finally + grad = grad_t @ inv_precond + # grad: (-1, last_dim) + grad = grad.view(original_size) + else: + # if not final + grad = inv_precond @ grad + # grad (dim, -1) + grad = grad.view(transposed_size) + + state["step"] += 1 + state["momentum_buffer"] = grad + p.data.add_(grad, alpha=-group["lr"]) + + return loss diff --git a/torch_optimizer/swats.py b/torch_optimizer/swats.py new file mode 100644 index 00000000..d1eddef5 --- /dev/null +++ b/torch_optimizer/swats.py @@ -0,0 +1,205 @@ +import torch +from torch.optim.optimizer import Optimizer + +from .types import Betas2, OptFloat, OptLossClosure, Params, State + +__all__ = ("SWATS",) + + +class SWATS(Optimizer): + r"""Implements SWATS Optimizer Algorithm. + It has been proposed in `Improving Generalization Performance by + Switching from Adam to SGD`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-2) + betas: coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps: term added to the denominator to improve + numerical stability (default: 1e-3) + weight_decay: weight decay (L2 penalty) (default: 0) + amsgrad: whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond` + (default: False) + nesterov: enables Nesterov momentum (default: False) + + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.SWATS(model.parameters(), lr=0.01) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://arxiv.org/pdf/1712.07628.pdf + + Note: + Reference code: https://github.com/Mrpatekful/swats + """ + + def __init__( + self, + params: Params, + lr: float = 1e-3, + betas: Betas2 = (0.9, 0.999), + eps: float = 1e-3, + weight_decay: float = 0, + amsgrad: bool = False, + nesterov: bool = False, + ): + if not 0.0 <= lr: + raise ValueError("Invalid learning rate: {}".format(lr)) + if not 0.0 <= eps: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + "Invalid beta parameter at index 0: {}".format(betas[0]) + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if weight_decay < 0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + defaults = dict( + lr=lr, + betas=betas, + eps=eps, + phase="ADAM", + weight_decay=weight_decay, + amsgrad=amsgrad, + nesterov=nesterov, + ) + + super().__init__(params, defaults) + + def __setstate__(self, state: State) -> None: + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("amsgrad", False) + group.setdefault("nesterov", False) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for w in group["params"]: + if w.grad is None: + continue + grad = w.grad.data + + if grad.is_sparse: + raise RuntimeError( + "Adam does not support sparse gradients, " + "please consider SparseAdam instead" + ) + + amsgrad = group["amsgrad"] + + state = self.state[w] + + # state initialization + if len(state) == 0: + state["step"] = 0 + # exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + w.data, memory_format=torch.preserve_format + ) + # exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like( + w.data, memory_format=torch.preserve_format + ) + # moving average for the non-orthogonal projection scaling + state["exp_avg2"] = w.new(1).fill_(0) + if amsgrad: + # maintains max of all exp. moving avg. + # of sq. grad. values + state["max_exp_avg_sq"] = torch.zeros_like( + w.data, memory_format=torch.preserve_format + ) + + exp_avg, exp_avg2, exp_avg_sq = ( + state["exp_avg"], + state["exp_avg2"], + state["exp_avg_sq"], + ) + + if amsgrad: + max_exp_avg_sq = state["max_exp_avg_sq"] + beta1, beta2 = group["betas"] + + state["step"] += 1 + + if group["weight_decay"] != 0: + grad.add_(w.data, alpha=group["weight_decay"]) + + # if its SGD phase, take an SGD update and continue + if group["phase"] == "SGD": + if "momentum_buffer" not in state: + buf = state["momentum_buffer"] = torch.clone( + grad + ).detach() + else: + buf = state["momentum_buffer"] + buf.mul_(beta1).add_(grad) + grad = buf + + grad.mul_(1 - beta1) + if group["nesterov"]: + grad.add_(buf, alpha=beta1) + + w.data.add_(grad, alpha=-group["lr"]) + continue + + # decay the first and second moment running average coefficient + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + if amsgrad: + # maintains the maximum of all 2nd + # moment running avg. till now + torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq) + # use the max. for normalizing running avg. of gradient + denom = max_exp_avg_sq.sqrt().add_(group["eps"]) + else: + denom = exp_avg_sq.sqrt().add_(group["eps"]) + + bias_correction1 = 1 - beta1 ** state["step"] + bias_correction2 = 1 - beta2 ** state["step"] + step_size = ( + group["lr"] * (bias_correction2**0.5) / bias_correction1 + ) + + p = -step_size * (exp_avg / denom) + w.data.add_(p) + + p_view = p.view(-1) + pg = p_view.dot(grad.view(-1)) + + if pg != 0: + # the non-orthognal scaling estimate + scaling = p_view.dot(p_view) / -pg + exp_avg2.mul_(beta2).add_(scaling, alpha=1 - beta2) + + # bias corrected exponential average + corrected_exp_avg = exp_avg2 / bias_correction2 + + # checking criteria of switching to SGD training + if ( + state["step"] > 1 + and corrected_exp_avg.allclose(scaling, rtol=1e-6) + and corrected_exp_avg > 0 + ): + group["phase"] = "SGD" + group["lr"] = corrected_exp_avg.item() + return loss diff --git a/torch_optimizer/types.py b/torch_optimizer/types.py new file mode 100644 index 00000000..8575734e --- /dev/null +++ b/torch_optimizer/types.py @@ -0,0 +1,12 @@ +from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Union + +from torch import Tensor + +Params = Union[Iterable[Tensor], Iterable[Dict[str, Any]]] + +LossClosure = Callable[[], float] +OptLossClosure = Optional[LossClosure] +Betas2 = Tuple[float, float] +State = Dict[str, Any] +OptFloat = Optional[float] +Nus2 = Tuple[float, float] diff --git a/torch_optimizer/yogi.py b/torch_optimizer/yogi.py new file mode 100644 index 00000000..ed0ab8d3 --- /dev/null +++ b/torch_optimizer/yogi.py @@ -0,0 +1,148 @@ +import math + +import torch +import torch.nn as nn +from torch.optim.optimizer import Optimizer + +from .types import Betas2, OptFloat, OptLossClosure, Params + +__all__ = ("Yogi",) + + +class Yogi(Optimizer): + r"""Implements Yogi Optimizer Algorithm. + It has been proposed in `Adaptive methods for Nonconvex Optimization`__. + + Arguments: + params: iterable of parameters to optimize or dicts defining + parameter groups + lr: learning rate (default: 1e-2) + betas: coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps: term added to the denominator to improve + numerical stability (default: 0.001) + initial_accumulator: initial values for first and + second moments (default: 1e-6) + weight_decay: weight decay (L2 penalty) (default: 0) + + Example: + >>> import torch_optimizer as optim + >>> optimizer = optim.Yogi(model.parameters(), lr=0.01) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization # noqa + + Note: + Reference code: https://github.com/4rtemi5/Yogi-Optimizer_Keras + """ + + def __init__( + self, + params: Params, + lr: float = 1e-2, + betas: Betas2 = (0.9, 0.999), + eps: float = 1e-3, + initial_accumulator: float = 1e-6, + weight_decay: float = 0, + ) -> None: + if lr <= 0.0: + raise ValueError("Invalid learning rate: {}".format(lr)) + if eps < 0.0: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError( + "Invalid beta parameter at index 0: {}".format(betas[0]) + ) + if not 0.0 <= betas[1] < 1.0: + raise ValueError( + "Invalid beta parameter at index 1: {}".format(betas[1]) + ) + if weight_decay < 0: + raise ValueError( + "Invalid weight_decay value: {}".format(weight_decay) + ) + + defaults = dict( + lr=lr, + betas=betas, + eps=eps, + initial_accumulator=initial_accumulator, + weight_decay=weight_decay, + ) + super(Yogi, self).__init__(params, defaults) + + def step(self, closure: OptLossClosure = None) -> OptFloat: + r"""Performs a single optimization step. + + Arguments: + closure: A closure that reevaluates the model and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError( + "Yogi does not support sparse gradients, " + "please consider SparseAdam instead" + ) + + state = self.state[p] + + # State initialization + # Followed from official implementation in tensorflow addons: + # https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/yogi.py#L118 # noqa + # For more details refer to the discussion: + # https://github.com/jettify/pytorch-optimizer/issues/77 + if len(state) == 0: + state["step"] = 0 + # Exponential moving average of gradient values + state["exp_avg"] = nn.init.constant_( + torch.empty_like( + p.data, memory_format=torch.preserve_format + ), + group["initial_accumulator"], + ) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = nn.init.constant_( + torch.empty_like( + p.data, memory_format=torch.preserve_format + ), + group["initial_accumulator"], + ) + + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + beta1, beta2 = group["betas"] + + state["step"] += 1 + bias_correction1 = 1 - beta1 ** state["step"] + bias_correction2 = 1 - beta2 ** state["step"] + + if group["weight_decay"] != 0: + grad = grad.add(p.data, alpha=group["weight_decay"]) + + # Decay the first and second moment running average coefficient + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + + grad_squared = grad.mul(grad) + + exp_avg_sq.addcmul_( + torch.sign(exp_avg_sq - grad_squared), + grad_squared, + value=-(1 - beta2), + ) + + denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_( + group["eps"] + ) + step_size = group["lr"] / bias_correction1 + p.data.addcdiv_(exp_avg, denom, value=-step_size) + + return loss