Repository navigation
Fix Layer.__delattr__ to untrack sub-layers (#23761) - #23781
Pushpak731 wants to merge 1 commit into
Conversation
|
❌ Approved issue check failed. This PR has been placed in draft until it links an approved issue assigned to @Pushpak731. Please review the Keras Pull Request Policy. To resolve this:
Referenced issue(s) not assigned to @Pushpak731: #23761. |
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a mechanism to untrack sub-layers when they are deleted from a parent layer or model via __delattr__. It adds the _untrack_layer helper method to manage tracker locking state during untracking and updates __delattr__ to invoke this helper when deleting a Layer attribute. A corresponding unit test is added to verify that weights, parameters, and the layer itself are correctly untracked upon deletion. The review feedback suggests expanding the unit test to verify that sub-layers nested inside container attributes (such as lists) are also correctly untracked when the container attribute is deleted.
| def test_delattr_untrack_layer(self): | ||
| inputs = Input((4,)) | ||
| model = Model(inputs, layers.Dense(4)(inputs)) | ||
| initial_num_weights = len(model.weights) | ||
| initial_params = model.count_params() | ||
|
|
||
| helper = layers.Dense(3) | ||
| helper.build((None, 4)) | ||
| model.helper = helper | ||
| self.assertGreater(len(model.weights), initial_num_weights) | ||
| self.assertGreater(model.count_params(), initial_params) | ||
|
|
||
| del model.helper | ||
| self.assertEqual(len(model.weights), initial_num_weights) | ||
| self.assertEqual(model.count_params(), initial_params) | ||
| self.assertNotIn(helper, model.layers) |
There was a problem hiding this comment.
Update the test to also verify that sub-layers nested inside container attributes (like lists) are correctly untracked when the container attribute is deleted.
def test_delattr_untrack_layer(self):
inputs = Input((4,))
model = Model(inputs, layers.Dense(4)(inputs))
initial_num_weights = len(model.weights)
initial_params = model.count_params()
helper = layers.Dense(3)
helper.build((None, 4))
model.helper = helper
self.assertGreater(len(model.weights), initial_num_weights)
self.assertGreater(model.count_params(), initial_params)
del model.helper
self.assertEqual(len(model.weights), initial_num_weights)
self.assertEqual(model.count_params(), initial_params)
self.assertNotIn(helper, model.layers)
# Test untracking layers inside a list container
helper_list = [layers.Dense(3), layers.Dense(2)]
for h in helper_list:
h.build((None, 4))
model.helper_list = helper_list
self.assertGreater(len(model.weights), initial_num_weights)
del model.helper_list
self.assertEqual(len(model.weights), initial_num_weights)
for h in helper_list:
self.assertNotIn(h, model.layers)
|
❌ Approved issue check failed. This PR has been placed in draft until it links an approved issue assigned to @Pushpak731. Please review the Keras Pull Request Policy. To resolve this:
Referenced issue(s) not assigned to @Pushpak731: #23761. |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #23781 +/- ##
==========================================
- Coverage 87.70% 84.83% -2.88%
==========================================
Files 481 481
Lines 64218 64242 +24
Branches 10802 10811 +9
==========================================
- Hits 56322 54497 -1825
- Misses 5135 7100 +1965
+ Partials 2761 2645 -116
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
@Pushpak731 could you please sign the CLA? |
|
Thanks @goyaladitya05! Signing the individual CLA now via https://cla.developers.google.com with the GitHub account. The commits on this PR are authored as |
|
@goyaladitya05 The CLA is signed — the check just hadn't re-evaluated since it last ran on Sep 29. Pushed a commit to re-trigger it, and |
keerthanakadiri
added
the
stat:awaiting keras-eng
…_delattr__ Address code review feedback on keras-team#23781: deleting a container attribute (e.g. a list of sub-layers) previously fell through to the default attribute deletion, leaving every nested sub-layer and variable tracked. Untrack tracked elements recursively before deleting the attribute, and extend test_delattr_untrack_layer with list, dict and nested-container cases. Guards against properties (e.g. `weights`, `layers`) that return containers: only instance attributes stored in `__dict__` are unwound.
|
Addressed the remaining code-review feedback in bf3c7d6: |
|
HI @Pushpak731 Could you please rebase your branch with master and fix the test failures? You can ignore the zizmor failure, as it will go away automatically once you rebase. Thanks! |
bf3c7d6 to
a9062f7
Compare
|
Thanks for the nudge! Rebased onto master and force-pushed as a single clean commit (a9062f7). The four torch attention failures were gone after the rebase — they came from the branch being ~100 commits behind, not from this change (they pass on the rebased head locally, 4/4). The format failure is fixed too; ruff and ruff-format are clean on both files now, and the delattr tests were moved into |
Really thanks for quick response, will take a look. |
| elif isinstance(obj, Layer): | ||
| self._untrack_layer(obj) | ||
| super().__delattr__(name) |
There was a problem hiding this comment.
Tracker stores each object once by id, so del of one attribute untracks the object even if another attribute still references it.
For example: self.encoder = enc; self.blocks = [enc, dec] is a normal pattern. Because the store holds enc once, del self.encoder now removes enc from _layers even though self.blocks[0] still uses it in call.
Pontential fix: only untrack when no other attribute still reaches the object. tree.flatten already understands TrackedList/TrackedDict/tuples and treats layers/variables as leaves
+ def _is_referenced_elsewhere(self, obj, skip_name):
+ others = {k: v for k, v in self.__dict__.items() if k != skip_name}
+ return any(leaf is obj for leaf in tree.flatten(others))
+
def __delattr__(self, name):
obj = getattr(self, name)
...
- elif isinstance(obj, Layer):
- self._untrack_layer(obj)
+ elif isinstance(obj, Layer):
+ if not self._is_referenced_elsewhere(obj, name):
+ self._untrack_layer(obj)
super().__delattr__(name)
unlike the container branch, the Layer branch has no name in self.dict guard; a read-only property returning a layer would be untracked and then raise AttributeError. The generic version above puts the guard on both.
| if isinstance(element, backend.Variable): | ||
| self._untrack_variable(element) | ||
| elif isinstance(element, Layer): | ||
| self._untrack_layer(element) | ||
| elif isinstance(element, (list, tuple, dict, set)): | ||
| self._untrack_tracked_elements(element) |
There was a problem hiding this comment.
Layer._initialize_tracker has five stores (trainable_variables, non_trainable_variables, metrics, layers, seed_generators). The PR's isinstance chain handles two of them; del self.my_metric / del self.seed_gen still leave the object in _metrics / _seed_generators
- def _untrack_layer(self, layer):
- if hasattr(self, "_tracker"):
- previous_lock_state = self._tracker.locked
- self._tracker.unlock()
- self._tracker.untrack(layer)
- if previous_lock_state is True:
- self._tracker.lock()
+ def _untrack(self, obj, skip_name):
+ """Untracks `obj` (any tracked type) unless another attribute
+ still references it. Recurses into containers."""
+ if not hasattr(self, "_tracker"):
+ return
+ if isinstance(obj, dict):
+ obj = list(obj.values())
+ if isinstance(obj, (list, tuple, set)):
+ for element in obj:
+ self._untrack(element, skip_name)
+ elif not self._is_referenced_elsewhere(obj, skip_name):
+ if isinstance(obj, backend.Variable):
+ self._untrack_variable(obj)
+ else:
+ self._tracker.untrack(obj)
def __delattr__(self, name):
obj = getattr(self, name)
if isinstance(obj, backend.Variable):
import gc
self._untrack_variable(obj)
super().__delattr__(name)
gc.collect()
- elif isinstance(obj, Layer):
- self._untrack_layer(obj)
- super().__delattr__(name)
- elif (
- isinstance(obj, (list, tuple, dict, set)) and name in self.__dict__
- ):
- ...
- self._untrack_tracked_elements(obj)
- super().__delattr__(name)
else:
+ if name in self.__dict__:
+ self._untrack(obj, name)
super().__delattr__(name)
-
- def _untrack_tracked_elements(self, container):
- ...
| # they are still mean 0 / variance 1. | ||
| self.assertAllClose(y, x / np.sqrt(1.0 + 1e-3), atol=1e-3) | ||
|
|
||
| def test_delattr_untrack_layer(self): |
There was a problem hiding this comment.
Fix for codecov 66%:
one 45-line test covering four scenarios with shared initial_* state; the Keras suite usually splits these (test_delattr_untrack_layer, _container, _alias, _metric) so a failure names the case. Also uses a Model as the parent; a plain layers.Layer subclass (the class under test) would exercise the locked-tracker path, since Model/Functional override _lock_state and never lock (confirmed: m._tracker.locked is False in the test fixture, so the "restore lock state" lines are effectively untested.
gaga1313
left a comment
gaga1313
left a comment
There was a problem hiding this comment.
Thanks for the fix! I have added two inline comments for avoiding deletion of object referenced multiple times and for extending the coverage to untrack metrics and seed_gen.
Deleting an attribute that held a sub-layer left the layer tracked: the parent still exposed its weights and count_params() kept counting them, so the removed layer kept influencing saving/tracking semantics. __delattr__ now untracks whatever the deleted attribute held: layers, variables, metrics and seed generators, plus the elements of container attributes (list/tuple/dict/set), which were tracked element-wise on assignment. Untracking goes through Tracker.untrack, which removes the object from every store, so all tracked types are covered. Since the tracker stores each object once by id, the object is only untracked when no other attribute still reaches it (checked with tree.flatten), so aliases of a deleted attribute keep the object tracked. Only attributes that were actually stored on assignment are untracked: they live in __dict__, or in torch.nn.Module's _modules on the torch backend. Properties like weights or layers keep the default behavior. The tracker's own store lists and the torch module registries are skipped when looking for remaining references, since they mirror what has been tracked rather than references the layer still holds.
a9062f7 to
0cdb24e
Compare
Description
Fixes #23761
This PR fixes an issue where deleting a sub-layer attribute via \del layer.attr\ did not untrack the layer, leaving its weights and parameters tracked in \layer.weights, \layer.trainable_weights, and \layer.count_params().
Changes Made
Contributor Agreement
Please review our PR Contribution Policy and AI-Assisted Contribution Policy and check all boxes below before submitting your PR for review: