pidajay

@davisyoshida Thanks for the clarification. I figured something like that needs to be done. I implemented a version where the user just has to specify a single decorator (like in...

@davisyoshida From what I see it is actually possible in eager mode but a bit more involved. You just have to recompute outside the gradient tape so nothing gets saved....

@davisyoshida and @Kokonut133 . I have created my pull request over here - https://github.com/tensorflow/addons/pull/1600. @davisyoshida maybe this answers the questions you had for me. Let me know. Note - I...

@davisyoshida the current implementation does not do checkpointing yet i.e save the outputs for specific layers. It just recomputes for every layer during backward pass. I am working on the...

The idea is to support various flavors of Keras (sequential, functional, etc) provided there is enough traction. I don't see a point moving outside Keras though when TF as a...

@paulter I have a version working with Keras but sequential models only. I have create a pull request as part of TF addons github repo - https://github.com/tensorflow/addons/pull/1600. You can find...

@paulter I have posted a small tutorial here https://github.com/pidajay/tf2_gradient_checkpointing/blob/master/tf_recompute_grad_tutorial.ipynb For this to work you need to replace (or just copy the delta) the custom_gradient.py file with this version in my...