Repository navigation
fix(stat): cap data-stat batches at the available batches per system - #6099
ShryukGrandhi wants to merge 1 commit into
Conversation
The backend-agnostic collect_batches/make_stat_input used by TF, JAX and pt_expt drew data_stat_nbatch batches from every system, wrapping around and resampling systems with fewer batches. The PT backend stops at the dataloader length. Cap the count at the system's batch count so small systems are not overweighted and the backends collect the same frames. Signed-off-by: ShryukGrandhi <85529990+ShryukGrandhi@users.noreply.github.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 9 functions across 3 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the main change: limiting data-stat batches to the batches available in each system. Linked Issues check Issue #5594requires consistent per-system sampling across TF and PT, with at most min(data_stat_nbatch, available batches), no implicit duplication for small systems, and a regression test. The PR up…Out of Scope Changes check The changed production code implements the sampling rule from issue #5594. The new regression test verifies backend consistency. The dist-check expectation update records the intended behavior for sma…Full details: Linked Issues check
-
Explanation
Issue
#5594requires consistent per-system sampling across TF and PT, with at most min(data_stat_nbatch, available batches), no implicit duplication for small systems, and a regression test. The PR updates _get_stat_numb_batches to use the available pass length from get_stat_numb_batches or fallback get_nbatches. Sources without either hook retain the prior behavior. TestMakeStatInputSmallSystems covers 3-frame and 7-frame systems with data_stat_nbatch=10 and compares the universal path, collect_batches, and the PT path. The updated dist-check test also verifies that a 2-frame system is not resampled.
Full details: Out of Scope Changes check
-
Explanation
The changed production code implements the sampling rule from issue
#5594. The new regression test verifies backend consistency. The dist-check expectation update records the intended behavior for small systems. No unrelated change is shown in the provided whole-PR summary.
-
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
Comment @coderabbitai help to get the list of available commands.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Agent: dot
Source LGTM at 7558f37. I checked the changed helper against DeepmdDataSystem's per-system get_nbatches() list, both collection call paths, the existing LMDB-specific hook precedence, and the pair-clearance scan behavior. The fallback now caps small systems instead of wrapping; explicit statistical-system hooks still take precedence, and sources without either count hook retain their previous behavior.
Independent verification: 31 assertions executing the exact helper/collector source with NumPy and bounded data/normalization/filter doubles passed, covering zero/small/large batch counts, requested limits, dedicated-hook precedence, legacy sources, merged/unmerged collect_batches, make_stat_input, and the filter-enabled count bound. This is not a native backend integration run. I also inspected the added 3/7-frame cross-backend regression and updated cyclic fixture. The author's reported native runs are supplementary evidence, not independently rerun here: this environment has no PyTorch, and TF/JAX/GPU execution is unverified.
Current CI is not green acceptance: seven workflows are action_required awaiting approval (including Test Python/C++/CUDA); only triage/labeler ran successfully. No source-blocking finding from this review.
This branch has not been deployed
2 participants
Summary
The backend-agnostic data-stat helpers in
deepmd/utils/model_stat.py(collect_batches, used by TF, andmake_stat_input, used by JAX/pt_expt) always drewdata_stat_nbatchbatches from every system. For aDeepmdDataSystem, a system with fewer batches than that wraps around and gets resampled, so small systems are overweighted in the statistics. The PT backend instead stops aftermin(data_stat_nbatch, len(dataloader))batches. The same input could therefore produce different descriptor/output statistics depending on the backend.Changes
_get_stat_numb_batchesnow caps the batch count at the system's available batches whenever the data source exposes them: through the existingget_stat_numb_batcheshook (LMDB) or, as a fallback,get_nbatches()(DeepmdDataSystem). This reuses_get_stat_pass_length, which already encoded that lookup for the pair-clearance filter path. Data sources that expose neither keep the old behavior.TestMakeStatInputSmallSystemsinsource/tests/consistent/test_make_stat_input.py: builds two systems with 3 and 7 frames, usesdata_stat_nbatch=10, and checks that the universal path,collect_batches, and the PT path all collect exactly 3 and 7 frames.test_disabled_filter_takes_the_next_batchesintest_dist_check.py, which asserted the old wrap-around (a 2-frame system sampled as[0, 1, 0]); it now expects[0, 1].Testing
pytest source/tests/consistent/test_make_stat_input.py: the new test fails before the fix (10 != 3 : system 0, key 'atype': nframes mismatch) and passes after it.pytest source/tests/consistent/test_make_stat_input.py source/tests/common/dpmodel/test_dist_check.py source/tests/pt_expt/test_lmdb_training.py: 71 passed, 1 skippedpytest source/tests/pt_expt/test_training.py: 45 passed, 18 skippedruff check/ruff formaton the touched files: cleanTensorFlow was not installed locally, so the TF-specific tests were not run. The TF path goes through the same
collect_batches, and the new test exercises it directly.Fixes #5594
Summary by CodeRabbit