deepmodeling / deepmodeling/deepmd-kit
Align data-stat sampling behavior between TF and PT backends
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
### Summary
The TensorFlow and PyTorch backends currently use different sampling rules when collecting data statistics (`data_stat_nbatch`). This can make descriptor/output statistics backend-dependent for the same training input.
From the current implementation:
- TensorFlow goes through the backend-agnostic helper `deepmd/utils/model_stat.py::collect_batches`, and for every system it always calls `data.get_batch(sys_idx=ii)` exactly `nbatches` times.
- PyTorch uses `deepmd/pt/utils/stat.py::make_stat_input`, builds an iterator from each per-system `DataLoader`, and collects only `min(nbatches, len(dataloaders[i]))` batches.
For systems with fewer available batches than `data_stat_nbatch`, TF may wrap/re-shuffle and sample duplicated frames to still collect `nbatches` batches, while PT stops after the actual dataloader length. PT's behavior seems more reasonable because it avoids overweighting small systems during data-stat collection.
### Why this matters
Data statistics are used to initialize/normalize model components. If TF and PT collect a different number of frames from the same system, then the computed statistics can differ even when the dataset, seed, and model configuration are intended to be equivalent.
This is especially visible for multi-system datasets where some systems have only a few batches and `data_stat_nbatch` is larger than the number of batches in those systems.
### Expected behavior
The data-stat sampling policy should be consistent across backends. A reasonable policy is to align TF with PT:
1. For each system, collect at most the number of available batches from that system, i.e. effectively `min(data_stat_nbatch, n_batches_for_system)`.
2. Avoid duplicating/overweighting small systems during the initial data-stat pass unless this is explicitly requested/documented.
3. Add a regression test that constructs systems with different frame counts and verifies that TF/PT collect the same number of batches/frames for statistics.
### Code locations
- TF/stat helper path: `deepmd/utils/model_stat.py::collect_batches`
- PT path: `deepmd/pt/utils/stat.py::make_stat_input`
- PT dataloader length comes from each per-system `DataLoader` built in `deepmd/pt/utils/dataloader.py`
### Notes
This is related in spirit to robust statistics coverage, but it is a separate backend-consistency issue from #4424.
---
Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
Contributor guide
Assessment
This issue has not been assessed yet.