Add RAdam and LookAhead optimizers #506

CyberZHG · 2019-09-13T13:19:05Z

Related to #422:

The name of the optimizer is RectifiedAdam in case of conflicting with other abbreviations. (The name appears in the first line of the second page in the paper)
Add Lookahead optimizer wrapper.
Unit tests for RAdam are compared with the results of the official implementation.
Unit tests for Ranger (RAdam + Lookahead) are compared with the results of the implementation from the proposer.

Squadrick · 2019-09-13T16:20:27Z

I think a more useful use of LookAhead would be to make a generic wrapper that can take a tf.keras.optimizers.Optimizer and apply lookahead. (Similar to tfa.optimizers.MovingAverage).

So the final API for using Ranger would look like this:

radam = tfa.optimizers.RectifiedAdam(lr=1e-3)
ranger = tfa.optimizers.LookAhead(radam, step=4, ratio=0.5)

This will allow for other optimizers to be used with LookAhead.

Ping @WindQAQ @facaiy for their thoughts on the matter.

CyberZHG · 2019-09-14T12:45:42Z

I've added a standalone Lookahead wrapper.

Squadrick

@CyberZHG Thank you so much for the contributions. Did a prelim review, and so far the code looks great. I've felt a few minor comments, I'll do a more thorough review later where I can verify the logic with the paper.

tensorflow_addons/optimizers/BUILD

tensorflow_addons/optimizers/README.md

tensorflow_addons/optimizers/__init__.py

tensorflow_addons/optimizers/lookahead.py

tensorflow_addons/optimizers/rectified_adam.py

googlebot · 2019-09-16T03:33:38Z

We found a Contributor License Agreement for you (the sender of this pull request), but were unable to find agreements for all the commit author(s) or Co-authors. If you authored these, maybe you used a different email address in the git commits than was used to sign the CLA (login here to double check)? If these were authored by someone else, then they will need to sign a CLA as well, and confirm that they're okay with these being contributed to Google.
In order to pass this check, please resolve this problem and then comment @googlebot I fixed it.. If the bot doesn't comment, it means it doesn't think anything has changed.

ℹ️ Googlers: Go here for more info.

googlebot · 2019-09-16T03:34:14Z

CLAs look good, thanks!

ℹ️ Googlers: Go here for more info.

CyberZHG · 2019-09-17T07:23:06Z

It seems that the mod operation is not working properly on GPU. I've changed the implementation to use floordiv in Lookahead. I don't know if there will be any numerical stability issue with the modification.

WindQAQ

Generally LGTM. Thanks for the contribution 😃

tensorflow_addons/optimizers/lookahead.py

seanpmorgan · 2019-09-17T18:28:28Z

It seems that the mod operation is not working properly on GPU. I've changed the implementation to use floordiv in Lookahead. I don't know if there will be any numerical stability issue with the modification.

I'm not sure why TF2 isn't using a soft device placement for this. There is no GPU implementation for the overloaded __mod__:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/cwise_op_gpu_mod.cu.cc#L22

But I would have assumed it'd fall back to CPU with dynamic kernels implemented. cc @karmel @martinwicke

seanpmorgan

LGTM.. thank you very much for this great contribution. Also thank you to everyone for the reviews and suggestions.

If there are no other objections I'd vote we merge and can possibly handle the __mod__ issue if some information becomes available.

Squadrick

Just a small nit.

tensorflow_addons/optimizers/lookahead.py

facaiy · 2019-09-18T05:37:55Z

tensorflow_addons/optimizers/lookahead.py

        slow_step_size = self._get_hyper('slow_step_size', var_dtype)
        step_back = slow_var + slow_step_size * (var - slow_var)
-        sync_cond = tf.equal(local_step % sync_period, 0)


I would have assumed it'd fall back to CPU with dynamic kernels implemented.

I believe so. Have you ever tried tf.math.floormod?

I've tried this one and it's not working.

https://github.com/tensorflow/tensorflow/blob/d30ba8443710e94d6a1ab1c71a2c61c437e24790/tensorflow/core/kernels/cwise_op_gpu_floor_mod.cu.cc#L22

Do you use tf.device to force all ops place on GPU? Can you give a minimal reproducing example? Thanks

The outputs are the same as the results of this CI. It fails the fit_simple_linear_model test.

Thanks for the information, really useful! Could you file an issue for it and we might resolve it later?

Squadrick · 2019-09-18T07:55:57Z

LGTM. @CyberZHG Thank you so much for this contribution, I'm sure the community will find the optimizers very useful.

CyberZHG added 4 commits September 13, 2019 14:43

Add Rectified Adam optimizer

01262fa

Add tests for amsgrad and weight decay

22bd673

Add tests of warmup for RAdam optimizer

8a8a217

Add lookahead for RAdam

be08a4a

CyberZHG requested review from facaiy, Squadrick and WindQAQ as code owners September 13, 2019 13:19

googlebot added the cla: yes label Sep 13, 2019

Squadrick added the optimizers label Sep 13, 2019

CyberZHG added 4 commits September 14, 2019 15:08

Decouple lookahead optimizer

289f0d8

Fix compatibility of Lookahead

cba8e4a

Add test case for training a simple linear model with Lookahead

3c3a0f5

Fix Lookahead when executing eagerly

24176af

Squadrick changed the title ~~Add RAdam optimizer~~ Add RAdam and LookAhead optimizers Sep 15, 2019

Squadrick suggested changes Sep 15, 2019

View reviewed changes

Fix orders and use public TensorFlow API

e627d75

googlebot added cla: no and removed cla: yes labels Sep 16, 2019

googlebot added cla: yes and removed cla: no labels Sep 16, 2019

seanpmorgan added the kokoro:force-run label Sep 16, 2019

kokoro-team removed the kokoro:force-run label Sep 16, 2019

CyberZHG added 2 commits September 16, 2019 22:25

Apply the formatting tool for RAdam and Lookahead

bd74c53

Use floordiv instead of mod in Lookahead

f2da716

seanpmorgan added the kokoro:force-run label Sep 17, 2019

kokoro-team removed the kokoro:force-run label Sep 17, 2019

CyberZHG requested a review from Squadrick September 17, 2019 16:23

WindQAQ reviewed Sep 17, 2019

View reviewed changes

tensorflow_addons/optimizers/lookahead.py Outdated Show resolved Hide resolved

Fix docstring for Lookahead

8149ac9

seanpmorgan added the kokoro:force-run label Sep 17, 2019

kokoro-team removed the kokoro:force-run label Sep 17, 2019

seanpmorgan previously approved these changes Sep 17, 2019

View reviewed changes

Squadrick suggested changes Sep 17, 2019

View reviewed changes

tensorflow_addons/optimizers/lookahead.py Outdated Show resolved Hide resolved

Fix docstring for Lookahead

b4cea81

CyberZHG dismissed seanpmorgan’s stale review via b4cea81 September 17, 2019 23:52

facaiy reviewed Sep 18, 2019

View reviewed changes

Squadrick added kokoro:force-run awaiting testing (then merge) labels Sep 18, 2019

kokoro-team removed the kokoro:force-run label Sep 18, 2019

Squadrick approved these changes Sep 18, 2019

View reviewed changes

Squadrick merged commit 220fad1 into tensorflow:master Sep 18, 2019

seanpmorgan removed the awaiting testing (then merge) label Sep 19, 2019

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add RAdam and LookAhead optimizers #506

Add RAdam and LookAhead optimizers #506

CyberZHG commented Sep 13, 2019 •

edited

Loading

Squadrick commented Sep 13, 2019 •

edited

Loading

CyberZHG commented Sep 14, 2019

Squadrick left a comment

googlebot commented Sep 16, 2019

googlebot commented Sep 16, 2019

CyberZHG commented Sep 17, 2019

WindQAQ left a comment

seanpmorgan commented Sep 17, 2019

seanpmorgan left a comment

Squadrick left a comment

facaiy Sep 18, 2019

CyberZHG Sep 18, 2019

CyberZHG Sep 18, 2019 •

edited

Loading

facaiy Sep 18, 2019

CyberZHG Sep 18, 2019

facaiy Sep 18, 2019

Squadrick commented Sep 18, 2019

Add RAdam and LookAhead optimizers #506

Add RAdam and LookAhead optimizers #506

Conversation

CyberZHG commented Sep 13, 2019 • edited Loading

Squadrick commented Sep 13, 2019 • edited Loading

CyberZHG commented Sep 14, 2019

Squadrick left a comment

Choose a reason for hiding this comment

googlebot commented Sep 16, 2019

googlebot commented Sep 16, 2019

CyberZHG commented Sep 17, 2019

WindQAQ left a comment

Choose a reason for hiding this comment

seanpmorgan commented Sep 17, 2019

seanpmorgan left a comment

Choose a reason for hiding this comment

Squadrick left a comment

Choose a reason for hiding this comment

facaiy Sep 18, 2019

Choose a reason for hiding this comment

CyberZHG Sep 18, 2019

Choose a reason for hiding this comment

CyberZHG Sep 18, 2019 • edited Loading

Choose a reason for hiding this comment

facaiy Sep 18, 2019

Choose a reason for hiding this comment

CyberZHG Sep 18, 2019

Choose a reason for hiding this comment

facaiy Sep 18, 2019

Choose a reason for hiding this comment

Squadrick commented Sep 18, 2019

CyberZHG commented Sep 13, 2019 •

edited

Loading

Squadrick commented Sep 13, 2019 •

edited

Loading

CyberZHG Sep 18, 2019 •

edited

Loading