Gradient accumulate optimizer

[dathudeptrai]

Describe the feature and the current behavior/state.

Hi, I think it's good if someone can support Gradient Accumulate optimizer for this repo, this feature is really helpful for those who train the large model with a low resource such as Bert, etc. The usage should be similar with tfa.optimizer.SWA:

opt = ...
accumulate_opt = tfa.optimizer.AccumulationOptimizer(opt, accumulate_steps=5)

There is an implementation of gradient accumulator but for custom training loop rather than Keras model fit here link.

Relevant information

• Are you willing to contribute it (yes/no): no
• Are you willing to maintain it going forward? (yes/no): no
• Is there a relevant academic paper? (if so, where):
• Is there already an implementation in another framework? (if so, where): here but for custom training loop.
• Was it part of tf.contrib? (if so, where): no

Which API type would this fall under (layer, metric, optimizer, etc.)
optimizer
Who will benefit with this feature?
all tensorflow users.
Any other info.

[seanpmorgan]

@tomerk please bring this up in ecosystem review, though I don't expect any conflicts. For the future do you need us to tag anyone or is the ecosystem-review label sufficient?

[bhack]

@seanpmorgan tensorflow/tensorflow#32176

[dathudeptrai]

@tomerk @bhack @seanpmorgan any update :D.

[dathudeptrai]

@bhack just remind :D

[bhack]

Check #2196 (comment)

[tomerk]

Notes from ecosystem review:

Rather than an optimizer that wraps another optimizer, we think this might actually make sense as an object to use as a gradient_transformer in optimizers (a new feature after an optimizer refactoring earlier this year):

https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Optimizer

Looping in @omalleyt12 who might have more insight / suggestions on how to do this.

Depending on how it comes out it might make sense in core, but addons seems like a good initial spot.

[dathudeptrai]

@omalleyt12 can you take a look ?

[dathudeptrai]

@tomerk @bhack @omalleyt12 a gentle ping in case you missed my previous comment :D

[bhack]

We have a quite duplicated ticket at tensorflow/tensorflow#32176

[dathudeptrai]

@bhack it's a long time ago, any plan for this feature for TF2 ?

[chn-lee-yumi]

It would be great to have such a param:

From a user perspective, having a parameter (e.g. in .fit) like gradient_accumulate_batch_frequency=1 that updates gradients every N batches would be perfect. (1 as default to update every batch, as it is at the moment)