Fix #20598 — Custom models/layers with sublayers in nested lists fail to save/load weights#22362
Fix #20598 — Custom models/layers with sublayers in nested lists fail to save/load weights#22362pctablet505 wants to merge 2 commits intokeras-team:masterfrom
Conversation
When a layer stores sub-layers in nested containers (e.g. self.blocks = [[Dense, Dense], [Dense, Dense]]), the save/load functions silently skipped inner containers because _save_container_state and _load_container_state only handled KerasSaveable items, not nested list/dict/tuple/set. Add recursive handling of nested containers with index-based namespacing (_container, _container_1, etc.) to avoid HDF5 path collisions. Fixes keras-team#20598
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a critical bug where Keras models containing sublayers within nested Python data structures (like lists of lists) would not properly save and restore the weights of those deeply embedded layers. The fix involves enhancing the saving and loading mechanisms to recursively traverse these nested containers, ensuring that all KerasSaveable objects, regardless of their nesting depth within standard Python collections, are correctly serialized and deserialized. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![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
The pull request successfully addresses the issue where custom models/layers with sublayers stored in nested lists failed to save and load their weights. The changes in saving_lib.py correctly implement recursive handling for nested list, dict, tuple, and set containers during both saving and loading processes. The addition of test_nested_list_layer_saving in keras/src/saving/saving_lib_test.py provides a good regression test, confirming the fix.
| name = "_container" | ||
| if name in used_names: | ||
| used_names[name] += 1 | ||
| name = f"{name}_{used_names[name]}" | ||
| else: | ||
| used_names[name] = 0 |
There was a problem hiding this comment.
The logic for generating a unique name for the container (_container, _container_1, etc.) is duplicated here and in _load_container_state. Consider extracting this logic into a small helper function to improve maintainability and reduce redundancy. This aligns with the principle of keeping APIs modular and automating repetitive tasks (Repository Style Guide, lines 51, 124).
name = "_container"
name = _get_unique_container_name(name, used_names)| name = "_container" | ||
| if name in used_names: | ||
| used_names[name] += 1 | ||
| name = f"{name}_{used_names[name]}" | ||
| else: | ||
| used_names[name] = 0 |
There was a problem hiding this comment.
This block duplicates the name generation logic found in _save_container_state. Extracting this into a helper function would make the code cleaner and easier to maintain. This aligns with the principle of keeping APIs modular and automating repetitive tasks (Repository Style Guide, lines 51, 124).
name = "_container"
name = _get_unique_container_name(name, used_names)
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #22362 +/- ##
=======================================
Coverage 82.95% 82.96%
=======================================
Files 595 595
Lines 66040 66054 +14
Branches 10305 10309 +4
=======================================
+ Hits 54785 54799 +14
Misses 8639 8639
Partials 2616 2616
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
3 participants
Fixes: #20598
Problem
When a custom
Modelstores sublayers in a nested list — e.g.self.blocks = [[Dense(8), Dense(8)], [Dense(8), Dense(8)]]— those inner lists are tracked by the layer tracker, but the saving logic only handledKerasSaveableobjects inside a container. Plainlist/dict/tuple/setelements were silently skipped, meaning all weights inside nested list containers were not saved or restored. Aftermodel.save()+load_model(), the sublayers would be reinitialized with fresh random weights.Root Cause
_save_container_stateand_load_container_stateiterated over a container and branched only onisinstance(saveable, KerasSaveable), with no handling for when the item itself is another nested container.Fix
Add an
elif isinstance(saveable, (list, dict, tuple, set))branch in both_save_container_stateand_load_container_statethat recurses into the nested container with a_containerpath segment (deduplicated with the same used-names counter).Files Changed
keras/src/saving/saving_lib.py— recursive nested-container handling in_save_container_stateand_load_container_statekeras/src/saving/saving_lib_test.py— regression test: model with 2×2 nestedDenseblock list, save → load → assert outputs match