RandomSampler(replacement=False) + padding collate_fn.
Changes needed:
- _collate_lmdb_batch: pad coord/force/atype to max_nloc,
add "atom_mask" bool tensor (nframes, max_nloc)
- Use RandomSampler(_data, replacement=False) as sampler
- Use fixed batch_size in DataLoader (not batch_sampler)
- Model forward: apply atom_mask to descriptor/fitting
- Loss: mask out padded atoms in force loss
Line: 240
|
_data: LmdbDataset, |
|
) -> tuple[DataLoader, Generator[Any, None, None]]: |
|
if _data.mixed_batch: |
|
# TODO [mixed_batch=True]: Replace SameNlocBatchSampler with |
|
# RandomSampler(replacement=False) + padding collate_fn. |
|
# Changes needed: |
|
# 1. _collate_lmdb_batch: pad coord/force/atype to max_nloc, |
|
# add "atom_mask" bool tensor (nframes, max_nloc) |
|
# 2. Use RandomSampler(_data, replacement=False) as sampler |
|
# 3. Use fixed batch_size in DataLoader (not batch_sampler) |
|
# 4. Model forward: apply atom_mask to descriptor/fitting |
RandomSampler(replacement=False) + padding collate_fn.
Changes needed:
add "atom_mask" bool tensor (nframes, max_nloc)
Line: 240
deepmd-kit/deepmd/pt/train/training.py
Lines 237 to 247 in 565f4be