Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Stellarpower.fix.new optimiser api for gradient monitoring #7611

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

stellarpower
Copy link

Description

Beginning of some changes needed for compatibility with Keras v3. Feel free to use as a springboard and cherrypick some lines:

Main addition is changing Gradient Callback to use non-Legacy Optimiser API
Also bundled in is a fix for !7578

The legacy optimiser API is gone in Keras v3. This version uses the current API, available in v2 and v3.

The _resource_apply_dense() etc. functions are no longer virtual in the parent class and so we need to use another method. This just overrides the base apply_gradients() function to perform the work.

Performed some basic tests using this gist, but pasted in browser, and may be far from perfect in general. Providing as something to work from at least.

Overriding config and then calling the parent did nothing also so deleted it.

To facilitate debugging, change:

    grad_acc_model.compile(
        loss=self.Model.loss,
        optimizer=_CustomOptimizer(),
        jit_compile = False,
        run_eagerly = True,
    )

to run eagerly.

self.model is a property; so needs another name

The gradient accumulator used the legacy optimisers, which are now removed.

Not tested the gradients yet, but this runs.
The legacy optimiser API is gone in Keras v3. This version uses the current API, available in v2 and v3. 

The _resource_apply_dense() etc. functions are no longer virtual in the parent class and so we need to use another method. This just overrides the base apply_gradients() function to perform the work.

Performed some basic tests using [this gist]( https://gist.github.com/stellarpower/24fd6b1cbd864a088ec2a5f3e8a9fb26/edit ), but pasted in browser, and may be far from perfect in general. Providing as something to work from at least.

Overriding config and then calling the parent did nothing also so deleted it.

To facilitate debugging, change:

        grad_acc_model.compile(
            loss=self.Model.loss,
            optimizer=_CustomOptimizer(),
            jit_compile = False,
            run_eagerly = True,
        )
to run eagerly.
Copy link

codecov bot commented May 8, 2024

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 55.66%. Comparing base (5612829) to head (2bf4a56).
Report is 11 commits behind head on main.

Additional details and impacted files

Impacted file tree graph

@@             Coverage Diff             @@
##             main    #7611       +/-   ##
===========================================
- Coverage   75.76%   55.66%   -20.10%     
===========================================
  Files         502      500        -2     
  Lines       54003    53056      -947     
===========================================
- Hits        40914    29533    -11381     
- Misses      12676    23163    +10487     
+ Partials      413      360       -53     
Flag Coverage Δ
func ?
system ?
unit 55.66% <ø> (-0.67%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

see 241 files with indirect coverage changes

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants