keras-team / keras-core

A multi-backend implementation of the Keras API, with support for TensorFlow, JAX, and PyTorch.
Apache License 2.0
1.27k stars 115 forks source link

Separate the metrics variables from non-trainable variables. #910

Closed qlzh727 closed 10 months ago

qlzh727 commented 10 months ago

As discussed in https://github.com/keras-team/keras-core/pull/897, we separate the metrics related variables from non-trainable variables, so that we can properly leverage the jax memory donation.

This will also allow us to skip the saving for metrics variables during checkpoint/savemodel

codecov[bot] commented 10 months ago

Codecov Report

Patch coverage: 100.00% and project coverage change: -3.58% :warning:

Comparison is base (a465816) 76.82% compared to head (c0098fd) 73.25%.

Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #910 +/- ## ========================================== - Coverage 76.82% 73.25% -3.58% ========================================== Files 329 329 Lines 31427 31434 +7 Branches 6112 6114 +2 ========================================== - Hits 24144 23027 -1117 - Misses 5719 6893 +1174 + Partials 1564 1514 -50 ``` | [Flag](https://app.codecov.io/gh/keras-team/keras-core/pull/910/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | Coverage Δ | | |---|---|---| | [keras_core](https://app.codecov.io/gh/keras-team/keras-core/pull/910/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | `73.17% <100.00%> (-3.56%)` | :arrow_down: | Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team#carryforward-flags-in-the-pull-request-comment) to find out more. | [Files Changed](https://app.codecov.io/gh/keras-team/keras-core/pull/910?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team) | Coverage Δ | | |---|---|---| | [keras\_core/layers/layer.py](https://app.codecov.io/gh/keras-team/keras-core/pull/910?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team#diff-a2VyYXNfY29yZS9sYXllcnMvbGF5ZXIucHk=) | `86.92% <100.00%> (-0.50%)` | :arrow_down: | ... and [17 files with indirect coverage changes](https://app.codecov.io/gh/keras-team/keras-core/pull/910/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team)

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.