Skip to content

[mixed_batch=True]: Replace SameNlocBatchSampler with #5386

@github-actions

Description

@github-actions

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
  5. 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

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions