search

google-research / big_vision

Official codebase used to develop Vision Transformer, SigLIP, MLP-Mixer, LiT and more.
Apache License 2.0
2.16k stars 147 forks source link

Add other custom optimizers #32

Open buttercutter opened 1 year ago

buttercutter commented 1 year ago

I noticed that there are no other choices of optimizers other than scale_by_adafactor()

This github issue would serve as a placeholder for other optimizers such asAdam#Lion or others in the future.

TODO:

buttercutter commented 1 year ago

As highlighted by @jcole75 , Lion is already implemented in deepmind's optax repo , so pip install git+https://github.com/deepmind/optax.git would do the work. See also https://optax.readthedocs.io/en/latest/api.html#lion

As for Adam#Lion, when I try to study how the optimizer function call returns, it seems that it is not that straightforward to translate the learning rate grafting to optax jax

Any advice on this ?

lucasb-eyer commented 11 months ago

I am also not sure yet how to implement lr grafting yet.