From b76f84167549d84006517af291cd28f271595a96 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 21 Oct 2026 23:55:36 +0800 Subject: [PATCH 1/3] feat(data): support recursive LMDB inputs and weighted training groups --- deepmd/common.py | 50 +- deepmd/dpmodel/train/data.py | 37 +- deepmd/dpmodel/utils/lmdb_data.py | 8 +- deepmd/jax/entrypoints/train.py | 4 +- deepmd/pd/entrypoints/main.py | 34 +- deepmd/pd/train/training.py | 86 +-- deepmd/pd/utils/dataloader.py | 95 +++- deepmd/pt/entrypoints/main.py | 120 +--- deepmd/pt/train/training.py | 257 +++------ deepmd/pt/utils/dataloader.py | 168 +++++- deepmd/pt/utils/stat.py | 28 +- deepmd/pt_expt/entrypoints/main.py | 100 +--- deepmd/pt_expt/train/training.py | 67 +-- deepmd/pt_expt/train/validation.py | 62 ++- deepmd/pt_expt/utils/lmdb_dataset.py | 19 + deepmd/tf/entrypoints/change_bias.py | 22 +- deepmd/tf2/entrypoints/train.py | 4 +- deepmd/utils/argcheck.py | 79 ++- deepmd/utils/data_config.py | 190 +++++++ deepmd/utils/data_group.py | 525 ++++++++++++++++++ deepmd/utils/data_path.py | 111 ++++ deepmd/utils/data_sampling.py | 128 +++++ deepmd/utils/data_system.py | 292 +++++----- deepmd/utils/json_schema.py | 20 +- deepmd/utils/model_stat.py | 3 +- deepmd/utils/path.py | 5 +- doc/train/training-advanced.md | 52 +- source/tests/common/test_argcheck_training.py | 70 +++ source/tests/common/test_data_config.py | 107 ++++ source/tests/common/test_json_schema.py | 21 + source/tests/common/test_path.py | 56 ++ source/tests/pt/test_lmdb_dataloader.py | 149 ++++- source/tests/pt/test_sampler.py | 26 + source/tests/pt_expt/test_lmdb_training.py | 128 ++++- source/tests/pt_expt/test_training.py | 2 +- 35 files changed, 2337 insertions(+), 788 deletions(-) create mode 100644 deepmd/utils/data_config.py create mode 100644 deepmd/utils/data_group.py create mode 100644 deepmd/utils/data_path.py create mode 100644 deepmd/utils/data_sampling.py create mode 100644 source/tests/common/test_data_config.py diff --git a/deepmd/common.py b/deepmd/common.py index fea637328a..94d07e3ab0 100644 --- a/deepmd/common.py +++ b/deepmd/common.py @@ -29,8 +29,9 @@ from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) -from deepmd.utils.path import ( - DPPath, +from deepmd.utils.data_path import ( + expand_sys_str, + rglob_sys_str, ) __all__ = [ @@ -41,6 +42,7 @@ "get_np_precision", "j_loader", "make_default_mesh", + "rglob_sys_str", "select_idx_map", ] @@ -188,50 +190,6 @@ def j_loader(filename: str | Path) -> dict[str, Any]: raise TypeError("config file must be json, or yaml/yml") -def expand_sys_str(root_dir: str | Path) -> list[str]: - """Recursively iterate over directories taking those that contain `type.raw` file. - - Parameters - ---------- - root_dir : Union[str, Path] - starting directory - - Returns - ------- - list[str] - list of string pointing to system directories - """ - root_dir = DPPath(root_dir) - matches = [str(d) for d in root_dir.rglob("*") if (d / "type.raw").is_file()] - if (root_dir / "type.raw").is_file(): - matches.append(str(root_dir)) - return matches - - -def rglob_sys_str(root_dir: str, patterns: list[str]) -> list[str]: - """Recursively iterate over directories taking those that contain `type.raw` file. - - Parameters - ---------- - root_dir : str, Path - starting directory - patterns : list[str] - list of glob patterns to match directories - - Returns - ------- - list[str] - list of string pointing to system directories - """ - root_dir = Path(root_dir) - matches = [] - for pattern in patterns: - matches.extend( - [str(d) for d in root_dir.rglob(pattern) if (d / "type.raw").is_file()] - ) - return list(set(matches)) # remove duplicates - - def get_np_precision(precision: "_PRECISION") -> np.dtype: """Get numpy precision constant from string. diff --git a/deepmd/dpmodel/train/data.py b/deepmd/dpmodel/train/data.py index 3e91dcfc53..86fde438f2 100644 --- a/deepmd/dpmodel/train/data.py +++ b/deepmd/dpmodel/train/data.py @@ -14,6 +14,14 @@ Any, ) +from deepmd.utils.data import ( + DataRequirementItem, + min_pair_dist_requirement, +) +from deepmd.utils.data_config import ( + DataConfig, + data_group_configs, +) from deepmd.utils.stat_file import ( StatFileMode, StatFileSpec, @@ -37,7 +45,7 @@ class TrainingTaskConfig: key: str model_params: Mapping[str, Any] - training_data_params: Mapping[str, Any] + training_data_params: DataConfig validation_data_params: Mapping[str, Any] | None stat_file: str | None valid_numb_batch: int @@ -110,6 +118,33 @@ def make_task_maps( return training_data, validation_data, stat_data +def add_training_data_requirements( + data: Any, + config: DataConfig, + model_params: Mapping[str, Any], + requirements: list[DataRequirementItem], + register: Callable[[Any, list[DataRequirementItem]], None], +) -> bool: + """Register each group's labels and its own derived pair-clearance margin. + + The returned flag enables margin filtering in the trainer and statistics + collector. Distances belong to individual groups' requirements, not to + that task-wide enable flag. + """ + groups = data_group_configs(config) + datasets = data.children if len(groups) > 1 else (data,) + filter_enabled = [] + for dataset, group in zip(datasets, groups, strict=True): + pair_filter = min_pair_dist_requirement( + model_params, float(group.get("min_pair_dist", 0.0)) + ) + register(dataset, requirements + ([] if pair_filter is None else [pair_filter])) + filter_enabled.append(pair_filter is not None) + if len(groups) > 1: + data.set_filter_enabled(filter_enabled) + return any(filter_enabled) + + def print_data_summaries( training_data: Mapping[str, Any], validation_data: Mapping[str, Any | None], diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 0d7ea42a89..e299b458bb 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -60,6 +60,7 @@ DataRequirementItem, DataRequirementSourcePolicy, ) +from deepmd.utils.data_path import is_lmdb as is_lmdb log = logging.getLogger(__name__) @@ -1928,11 +1929,6 @@ def close(self) -> None: _release_lmdb_executor(self._num_workers) -def is_lmdb(systems: str) -> bool: - """Check if systems points to an LMDB dataset.""" - return systems.endswith(".lmdb") or Path(systems, "data.mdb").is_file() - - def _parse_metadata(meta: dict) -> tuple[int, str, list[int]]: """Parse LMDB metadata into (nframes, frame_fmt, natoms_per_type). @@ -3015,7 +3011,7 @@ def compute_block_targets( Each element is ``(system_indices_in_block, target_frame_count)``. Returns empty list if no expansion is needed (all targets == actual). """ - from deepmd.utils.data_system import ( + from deepmd.utils.data_sampling import ( prob_sys_size_ext, ) diff --git a/deepmd/jax/entrypoints/train.py b/deepmd/jax/entrypoints/train.py index 61d415a340..e86e5d4c8c 100644 --- a/deepmd/jax/entrypoints/train.py +++ b/deepmd/jax/entrypoints/train.py @@ -183,7 +183,7 @@ def factory( type_map = task_model.get_type_map() ipt_type_map = type_map if len(type_map) > 0 else None train_data = get_data( - dict(task_config.training_data_params), + task_config.training_data_params, task_model.get_rcut(), ipt_type_map, None, @@ -291,7 +291,7 @@ def update_sel( for task_config in iter_training_task_configs(jdata): type_map = task_config.model_params.get("type_map") train_data = get_data( - dict(task_config.training_data_params), + task_config.training_data_params, 0, # not used type_map, None, # not used diff --git a/deepmd/pd/entrypoints/main.py b/deepmd/pd/entrypoints/main.py index cf15e99192..250a218574 100644 --- a/deepmd/pd/entrypoints/main.py +++ b/deepmd/pd/entrypoints/main.py @@ -43,6 +43,7 @@ ) from deepmd.pd.utils.dataloader import ( DpLoaderSet, + create_dataset, ) from deepmd.pd.utils.env import ( DEVICE, @@ -66,9 +67,14 @@ from deepmd.utils.compat import ( update_deepmd_input, ) +from deepmd.utils.data_group import ( + GroupedDataset, +) +from deepmd.utils.data_path import ( + process_systems, +) from deepmd.utils.data_system import ( get_data, - process_systems, ) from deepmd.utils.model_preset import ( expand_model_preset, @@ -104,19 +110,11 @@ def prepare_trainer_input_single( data_dict_single: dict[str, Any], rank: int = 0, seed: int | None = None, - ) -> tuple[DpLoaderSet, DpLoaderSet | None, DPPath | None]: + ) -> tuple[ + DpLoaderSet | GroupedDataset, DpLoaderSet | GroupedDataset | None, DPPath | None + ]: training_dataset_params = data_dict_single["training_data"] validation_dataset_params = data_dict_single.get("validation_data") - validation_systems = ( - validation_dataset_params["systems"] if validation_dataset_params else None - ) - training_systems = training_dataset_params["systems"] - trn_patterns = training_dataset_params.get("rglob_patterns", None) - training_systems = process_systems(training_systems, patterns=trn_patterns) - if validation_systems is not None: - val_patterns = validation_dataset_params.get("rglob_patterns", None) - validation_systems = process_systems(validation_systems, val_patterns) - # stat files stat_file_path_single = data_dict_single.get("stat_file") if rank != 0: @@ -134,18 +132,16 @@ def prepare_trainer_input_single( # avoid the same batch sequence among devices rank_seed = [rank, seed % (2**32)] if seed is not None else None validation_data_single = ( - DpLoaderSet( - validation_systems, - validation_dataset_params["batch_size"], + create_dataset( + validation_dataset_params, model_params_single["type_map"], seed=rank_seed, ) - if validation_systems + if validation_dataset_params is not None else None ) - train_data_single = DpLoaderSet( - training_systems, - training_dataset_params["batch_size"], + train_data_single = create_dataset( + training_dataset_params, model_params_single["type_map"], seed=rank_seed, ) diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py index a33908c632..8210ba8022 100644 --- a/deepmd/pd/train/training.py +++ b/deepmd/pd/train/training.py @@ -24,15 +24,11 @@ from paddle.framework import ( core, ) -from paddle.io import ( - DataLoader, -) from deepmd.common import ( symlink_prefix_files, ) from deepmd.dpmodel.utils import ( - compute_total_numb_batch, resolve_model_prob, resolve_model_prob_from_epochs, ) @@ -59,7 +55,7 @@ ) from deepmd.pd.utils.dataloader import ( BufferedIterator, - get_sampler_from_params, + create_dataloader, ) from deepmd.pd.utils.env import ( CINN, @@ -67,7 +63,6 @@ DEFAULT_PRECISION, DEVICE, JIT, - NUM_WORKERS, SAMPLER_RECORD, enable_prim, ) @@ -185,23 +180,10 @@ def get_data_loader( def get_dataloader_and_buffer( _data: Any, _params: dict[str, Any] ) -> tuple[Any, Any]: - _sampler = get_sampler_from_params(_data, _params) - if _sampler is None: - log.warning( - "Sampler not specified!" - ) # None sampler will lead to a premature stop iteration. Replacement should be True in attribute of the sampler to produce expected number of items in one iteration. - _dataloader = DataLoader( - _data, - batch_sampler=paddle.io.BatchSampler( - sampler=_sampler, - drop_last=False, - ), - num_workers=NUM_WORKERS - if dist.is_available() - else 0, # setting to 0 diverges the behavior of its iterator; should be >=1 - collate_fn=lambda batch: batch[0], # prevent extra conversion + _dataloader = create_dataloader( + _data, _params, seed=_training_params.get("seed") ) - _data_buffered = BufferedIterator(iter(_dataloader)) + _data_buffered = BufferedIterator(_dataloader) return _dataloader, _data_buffered training_dataloader, training_data_buffered = get_dataloader_and_buffer( @@ -336,17 +318,9 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: self.validation_data, self.valid_numb_batch, ) = get_data_loader(training_data, validation_data, training_params) - training_data.print_summary( - "training", - to_numpy_array(self.training_dataloader.batch_sampler.sampler.weights), - ) + self.training_dataloader.print_summary("training") if validation_data is not None: - validation_data.print_summary( - "validation", - to_numpy_array( - self.validation_dataloader.batch_sampler.sampler.weights - ), - ) + self.validation_dataloader.print_summary("validation") else: ( self.training_dataloader, @@ -385,25 +359,15 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: training_params["data_dict"][model_key], ) - training_data[model_key].print_summary( - f"training in {model_key}", - to_numpy_array( - self.training_dataloader[ - model_key - ].batch_sampler.sampler.weights - ), + self.training_dataloader[model_key].print_summary( + f"training in {model_key}" ) if ( validation_data is not None and validation_data[model_key] is not None ): - validation_data[model_key].print_summary( - f"validation in {model_key}", - to_numpy_array( - self.validation_dataloader[ - model_key - ].batch_sampler.sampler.weights - ), + self.validation_dataloader[model_key].print_summary( + f"validation in {model_key}" ) per_task_total = [] @@ -415,13 +379,7 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: ) if self.num_epoch <= 0: raise ValueError("training.num_epoch must be positive.") - sampler_weights = to_numpy_array( - self.training_dataloader.batch_sampler.sampler.weights - ) - total_numb_batch = compute_total_numb_batch( - training_data.index, - sampler_weights, - ) + total_numb_batch = self.training_dataloader.epoch_length if total_numb_batch <= 0: raise ValueError( "Total number of training batches must be positive." @@ -441,16 +399,8 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: "are mutually exclusive." ) for model_key in self.model_keys: - sampler_weights = to_numpy_array( - self.training_dataloader[ - model_key - ].batch_sampler.sampler.weights - ) per_task_total.append( - compute_total_numb_batch( - training_data[model_key].index, - sampler_weights, - ) + self.training_dataloader[model_key].epoch_length ) ( self.model_prob, @@ -1184,9 +1134,7 @@ def get_data( batch_data = next(iter(self.training_data)) except StopIteration: # Refresh the status of the dataloader to start from a new epoch - self.training_data = BufferedIterator( - iter(self.training_dataloader) - ) + self.training_data = BufferedIterator(self.training_dataloader) batch_data = next(iter(self.training_data)) else: if self.validation_data is None: @@ -1194,9 +1142,7 @@ def get_data( try: batch_data = next(iter(self.validation_data)) except StopIteration: - self.validation_data = BufferedIterator( - iter(self.validation_dataloader) - ) + self.validation_data = BufferedIterator(self.validation_dataloader) batch_data = next(iter(self.validation_data)) else: if is_train: @@ -1205,7 +1151,7 @@ def get_data( except StopIteration: # Refresh the status of the dataloader to start from a new epoch self.training_data[task_key] = BufferedIterator( - iter(self.training_dataloader[task_key]) + self.training_dataloader[task_key] ) batch_data = next(iter(self.training_data[task_key])) else: @@ -1215,7 +1161,7 @@ def get_data( batch_data = next(iter(self.validation_data[task_key])) except StopIteration: self.validation_data[task_key] = BufferedIterator( - iter(self.validation_dataloader[task_key]) + self.validation_dataloader[task_key] ) batch_data = next(iter(self.validation_data[task_key])) diff --git a/deepmd/pd/utils/dataloader.py b/deepmd/pd/utils/dataloader.py index 090d1713f5..09e92ae7e4 100644 --- a/deepmd/pd/utils/dataloader.py +++ b/deepmd/pd/utils/dataloader.py @@ -42,15 +42,29 @@ ) from deepmd.pd.utils.utils import ( mix_entropy, + to_numpy_array, ) from deepmd.utils import random as dp_random from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_config import ( + DataConfig, + data_group_configs, + resolve_data_groups, +) +from deepmd.utils.data_group import ( + GroupedDataset, + GroupedTrainingDataLoader, + TrainingDataLoader, +) +from deepmd.utils.data_sampling import ( + data_child_seeds, + global_batch_counts, + system_probabilities, +) from deepmd.utils.data_system import ( print_summary, - prob_sys_size_ext, - process_sys_probs, ) log = logging.getLogger(__name__) @@ -289,7 +303,7 @@ def run(self) -> None: class BufferedIterator: def __init__(self, iterable: Iterable) -> None: self._queue = queue.Queue(QUEUESIZE) - self._iterable = iterable + self._iterable = iter(iterable) self._consumer = None self.start_time = time.time() @@ -353,20 +367,13 @@ def collate_batch(batch: list[dict]) -> dict: def get_weighted_sampler( - training_data: DpLoaderSet, prob_style: str, sys_prob: bool = False + training_data: DpLoaderSet, prob_style: str | list[float], sys_prob: bool = False ) -> WeightedRandomSampler: - if sys_prob is False: - if prob_style == "prob_uniform": - prob_v = 1.0 / float(training_data.__len__()) - probs = [prob_v for ii in range(training_data.__len__())] - else: # prob_sys_size;A:B:p1;C:D:p2 or prob_sys_size = prob_sys_size;0:nsys:1.0 - if prob_style == "prob_sys_size": - style = f"prob_sys_size;0:{len(training_data)}:1.0" - else: - style = prob_style - probs = prob_sys_size_ext(style, len(training_data), training_data.index) - else: - probs = process_sys_probs(prob_style, training_data.index) + probs = system_probabilities( + training_data.index, + sys_probs=prob_style if sys_prob else None, + auto_prob="prob_sys_size" if sys_prob else prob_style, + ) log.debug("Generated weighted sampler with prob array: " + str(probs)) # training_data.total_batch is the size of one epoch, you can increase it to avoid too many rebuilding of iteraters len_sampler = training_data.total_batch * max(env.NUM_WORKERS, 1) @@ -390,3 +397,59 @@ def get_sampler_from_params( else: _sampler = get_weighted_sampler(_data, "prob_sys_size") return _sampler + + +def create_dataset( + config: DataConfig, type_map: list[str], *, seed: int | list[int] | None = None +) -> DpLoaderSet | GroupedDataset: + """Construct NumPy loader groups through the shared path resolver.""" + groups = resolve_data_groups(config, formats=("deepmd",)) + datasets = [ + DpLoaderSet(list(group.systems), group.batch_size, type_map, seed=seed) + for group in groups + ] + return ( + datasets[0] + if len(datasets) == 1 + else GroupedDataset( + datasets, data_group_configs(config), [group.prob for group in groups] + ) + ) + + +def create_dataloader( + data: DpLoaderSet | GroupedDataset, + config: DataConfig, + *, + seed: int | None = None, +) -> TrainingDataLoader | GroupedTrainingDataLoader: + """Bind native Paddle loaders to the shared batch-composition interface.""" + if isinstance(data, GroupedDataset): + seeds = data_child_seeds( + 0 if seed is None and dist.get_world_size() > 1 else seed, + len(data.children) + 1, + ) + children = [ + create_dataloader(child, params, seed=child_seed) + for child, params, child_seed in zip( + data.children, data.configs, seeds[1:], strict=True + ) + ] + return GroupedTrainingDataLoader(data, children, seeds[0]) + params = config[0] if isinstance(config, list) else config + sampler = get_sampler_from_params(data, params) + loader = DataLoader( + data, + batch_sampler=BatchSampler(sampler=sampler, drop_last=False), + num_workers=env.NUM_WORKERS if dist.is_available() else 0, + collate_fn=lambda batch: batch[0], + ) + return TrainingDataLoader( + data, + loader, + to_numpy_array(sampler.weights), + lambda: data.index, + sampling_sizes=lambda: global_batch_counts( + [len(system) for system in data.systems], data.batch_sizes + ), + ) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 9d54519e9d..f5d306b259 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -61,6 +61,7 @@ ) from deepmd.pt.utils.dataloader import ( DpLoaderSet, + create_dataset, ) from deepmd.pt.utils.env import ( DEVICE, @@ -70,7 +71,6 @@ ) from deepmd.pt.utils.lmdb_dataset import ( LmdbDataset, - is_lmdb, ) from deepmd.pt.utils.multi_task import ( preprocess_shared_params, @@ -90,10 +90,15 @@ from deepmd.utils.data import ( min_pair_dist_requirement, ) -from deepmd.utils.data_system import ( - get_data, +from deepmd.utils.data_group import ( + GroupedDataset, +) +from deepmd.utils.data_path import ( process_systems, ) +from deepmd.utils.data_system import ( + get_neighbor_stat_data, +) from deepmd.utils.model_preset import ( expand_model_preset, ) @@ -155,8 +160,8 @@ def prepare_trainer_input_single( rank: int = 0, seed: int | None = None, ) -> tuple[ - DpLoaderSet | LmdbDataset, - DpLoaderSet | LmdbDataset | None, + DpLoaderSet | LmdbDataset | GroupedDataset, + DpLoaderSet | LmdbDataset | GroupedDataset | None, StatFileSpec, ]: # get data modifier @@ -167,11 +172,6 @@ def prepare_trainer_input_single( training_dataset_params = data_dict_single["training_data"] validation_dataset_params = data_dict_single.get("validation_data") - validation_systems = ( - validation_dataset_params["systems"] if validation_dataset_params else None - ) - training_systems = training_dataset_params["systems"] - stat_file_spec = StatFileSpec( data_dict_single.get("stat_file"), data_dict_single.get("stat_file_mode", "update"), @@ -179,56 +179,22 @@ def prepare_trainer_input_single( rank_seed = [rank, seed % (2**32)] if seed is not None else None - def _make_dp_loader_set( - systems: str | list[str], - dataset_params: dict[str, Any], - ) -> DpLoaderSet: - """Create a DpLoaderSet from systems with pattern expansion.""" - patterns = dataset_params.get("rglob_patterns") - systems = process_systems(systems, patterns=patterns) - return DpLoaderSet( - systems, - dataset_params["batch_size"], + train_data_single = create_dataset( + training_dataset_params, + model_params_single["type_map"], + seed=rank_seed, + modifier=modifier, + ) + validation_data_single = ( + create_dataset( + validation_dataset_params, model_params_single["type_map"], seed=rank_seed, modifier=modifier, ) - - # LMDB path: single string → LmdbDataset - if isinstance(training_systems, str) and is_lmdb(training_systems): - auto_prob = training_dataset_params.get("auto_prob", None) - train_data_single = LmdbDataset( - training_systems, - model_params_single["type_map"], - training_dataset_params["batch_size"], - auto_prob_style=auto_prob, - ) - if ( - validation_systems is not None - and isinstance(validation_systems, str) - and is_lmdb(validation_systems) - ): - validation_data_single = LmdbDataset( - validation_systems, - model_params_single["type_map"], - validation_dataset_params["batch_size"], - ) - elif validation_systems is not None: - validation_data_single = _make_dp_loader_set( - validation_systems, validation_dataset_params - ) - else: - validation_data_single = None - else: - # Standard npy path - train_data_single = _make_dp_loader_set( - training_systems, training_dataset_params - ) - validation_data_single = ( - _make_dp_loader_set(validation_systems, validation_dataset_params) - if validation_systems - else None - ) + if validation_dataset_params is not None + else None + ) return ( train_data_single, validation_data_single, @@ -399,21 +365,9 @@ def train( if not multi_task: type_map = config["model"].get("type_map") - training_systems = config["training"]["training_data"].get("systems") - if ( - training_systems is not None - and isinstance(training_systems, str) - and is_lmdb(training_systems) - ): - from deepmd.dpmodel.utils.lmdb_data import ( - make_neighbor_stat_data, - ) - - train_data = make_neighbor_stat_data(training_systems, type_map) - else: - train_data = get_data( - config["training"]["training_data"], 0, type_map, None - ) + train_data = get_neighbor_stat_data( + config["training"]["training_data"], type_map + ) config["model"], min_nbor_dist = BaseModel.update_sel( train_data, type_map, config["model"] ) @@ -421,26 +375,10 @@ def train( min_nbor_dist = {} for model_item in config["model"]["model_dict"]: type_map = config["model"]["model_dict"][model_item].get("type_map") - training_systems = config["training"]["data_dict"][model_item][ - "training_data" - ].get("systems") - if ( - training_systems is not None - and isinstance(training_systems, str) - and is_lmdb(training_systems) - ): - from deepmd.dpmodel.utils.lmdb_data import ( - make_neighbor_stat_data, - ) - - train_data = make_neighbor_stat_data(training_systems, type_map) - else: - train_data = get_data( - config["training"]["data_dict"][model_item]["training_data"], - 0, - type_map, - None, - ) + train_data = get_neighbor_stat_data( + config["training"]["data_dict"][model_item]["training_data"], + type_map, + ) config["model"]["model_dict"][model_item], min_nbor_dist[model_item] = ( BaseModel.update_sel( train_data, type_map, config["model"]["model_dict"][model_item] diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 7efc76e4d9..adf26c4383 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -35,8 +35,8 @@ change_model_out_bias, resolve_step_schedule, ) -from deepmd.dpmodel.utils import ( - compute_total_numb_batch, +from deepmd.dpmodel.train.data import ( + add_training_data_requirements, ) from deepmd.loggers import ( is_node_main_process, @@ -85,31 +85,23 @@ ) from deepmd.pt.utils.dataloader import ( DpLoaderSet, - get_sampler_from_params, + create_dataloader, ) from deepmd.pt.utils.env import ( DEVICE, JIT, LOCAL_RANK, - NUM_WORKERS, SAMPLER_RECORD, ) from deepmd.pt.utils.learning_rate import ( BaseLR, ) -from deepmd.pt.utils.lmdb_dataset import ( - LmdbBatchDataLoader, - LmdbDataset, -) from deepmd.pt.utils.stat import ( make_stat_input, pair_margin_frame_mask, scan_redu_stats, select_batch_frames, ) -from deepmd.pt.utils.utils import ( - to_numpy_array, -) from deepmd.pt_expt.train.ema import ( EMA_CHECKPOINT_KEY, ModelEMA, @@ -132,7 +124,11 @@ from deepmd.utils.data import ( DataRequirementItem, has_data_requirement, - min_pair_dist_requirement, +) +from deepmd.utils.data_group import ( + GroupedTrainingDataLoader, + TrainingDataLoader, + iter_data_sources, ) from deepmd.utils.finetune import ( warn_configuration_mismatch_during_finetune, @@ -162,9 +158,6 @@ ZeroRedundancyOptimizer, ) from torch.nn.parallel import DistributedDataParallel as DDP -from torch.utils.data import ( - DataLoader, -) from deepmd.utils.stat_file import ( StatFileSpec, @@ -289,103 +282,36 @@ def cycle_iterator(iterable: Iterable) -> Generator[Any, None, None]: yield from it def get_data_loader( - _training_data: DpLoaderSet | LmdbDataset, - _validation_data: DpLoaderSet | LmdbDataset | None, + _training_data: Any, + _validation_data: Any | None, _training_params: dict[str, Any], ) -> tuple[ - DataLoader | LmdbBatchDataLoader, + TrainingDataLoader | GroupedTrainingDataLoader, Generator[Any, None, None], - DataLoader | LmdbBatchDataLoader | None, + TrainingDataLoader | GroupedTrainingDataLoader | None, Generator[Any, None, None] | None, int, ]: - def get_dataloader_and_iter_lmdb( - _data: LmdbDataset, - ) -> tuple[LmdbBatchDataLoader, Generator[Any, None, None]]: - # The sampler yields one list[int] per batch. Every batch is - # homogeneous in label availability; whether it is also - # homogeneous in atom count follows from the dataset's - # batch_size rule, which the sampler reads off the reader. - from deepmd.dpmodel.utils.lmdb_data import ( - LmdbBatchSampler, - ) - - _block_targets = getattr(_data, "_block_targets", None) - _sampler_kwargs = { - "shuffle": True, - "seed": _training_params.get("seed"), - "block_targets": _block_targets, - } - - if self.world_size > 1: - from deepmd.dpmodel.utils.lmdb_data import ( - DistributedLmdbBatchSampler, - ) - - _inner_sampler = DistributedLmdbBatchSampler( - _data._reader, - rank=self.rank, - world_size=self.world_size, - **_sampler_kwargs, - ) - else: - _inner_sampler = LmdbBatchSampler( - _data._reader, - **_sampler_kwargs, - ) - - _dataloader = LmdbBatchDataLoader( - _data, - _inner_sampler, - pin_memory=DEVICE.type != "cpu", - ) - _data_iter = cycle_iterator(_dataloader) - return _dataloader, _data_iter - def get_dataloader_and_iter( - _data: DpLoaderSet, _params: dict[str, Any] - ) -> tuple[DataLoader, Generator[Any, None, None]]: - _sampler = get_sampler_from_params(_data, _params) - if _sampler is None: - log.warning( - "Sampler not specified!" - ) # None sampler will lead to a premature stop iteration. Replacement should be True in attribute of the sampler to produce expected number of items in one iteration. - _dataloader = DataLoader( + _data: Any, _params: Any + ) -> tuple[Any, Generator[Any, None, None]]: + _dataloader = create_dataloader( _data, - sampler=_sampler, - batch_size=None, - num_workers=NUM_WORKERS - if dist.is_available() - else 0, # setting to 0 diverges the behavior of its iterator; should be >=1 - drop_last=False, - collate_fn=lambda batch: batch, # prevent extra conversion - pin_memory=(DEVICE != "cpu"), # pin memory only if not on CPU + _params, + seed=_training_params.get("seed"), + rank=self.rank, + world_size=self.world_size, ) - _data_iter = cycle_iterator(_dataloader) - return _dataloader, _data_iter + return _dataloader, cycle_iterator(_dataloader) - if isinstance(_training_data, LmdbDataset): - training_dataloader, training_data_iter = get_dataloader_and_iter_lmdb( - _training_data - ) - else: - training_dataloader, training_data_iter = get_dataloader_and_iter( - _training_data, _training_params["training_data"] - ) + training_dataloader, training_data_iter = get_dataloader_and_iter( + _training_data, _training_params["training_data"] + ) if _validation_data is not None: - if isinstance(_validation_data, LmdbDataset): - ( - validation_dataloader, - validation_data_iter, - ) = get_dataloader_and_iter_lmdb(_validation_data) - else: - ( - validation_dataloader, - validation_data_iter, - ) = get_dataloader_and_iter( - _validation_data, _training_params["validation_data"] - ) + validation_dataloader, validation_data_iter = get_dataloader_and_iter( + _validation_data, _training_params["validation_data"] + ) valid_numb_batch = _training_params["validation_data"].get( "numb_btch", 1 ) @@ -404,9 +330,9 @@ def get_dataloader_and_iter( def single_model_stat( _model: Any, _data_stat_nbatch: int, - _training_data: DpLoaderSet, + _training_data: Any, _stat_file_spec: StatFileSpec, - _min_pair_dist: float = 0.0, + _filter_enabled: bool = False, _data_stat_full: bool = False, finetune_has_new_type: bool = False, preset_observed_type: list[str] | None = None, @@ -417,7 +343,10 @@ def get_sample() -> Any: _training_data.systems, _training_data.dataloaders, _data_stat_nbatch, - min_pair_dist=_min_pair_dist, + min_pair_dist=_filter_enabled, + filter_enabled_by_system=getattr( + _training_data, "stat_filter_enabled", None + ), ) return sampled @@ -430,7 +359,10 @@ def get_sample() -> Any: ntypes, keys, intensive=intensive, - min_pair_dist=_min_pair_dist, + min_pair_dist=_filter_enabled, + filter_enabled_by_system=getattr( + _training_data, "stat_filter_enabled", None + ), ) ) @@ -520,38 +452,37 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: ) # Data - # Radius of the pair-clearance filter, per task. - self.min_pair_dist_by_task: dict[str, float] = {} + # Per-group requirements own the filter radii; tasks only enable filtering. + self.pair_filter_enabled_by_task: dict[str, bool] = {} if not self.multi_task: # add data requirement for labels data_requirement = self.loss.label_requirement data_requirement += get_additional_data_requirement(self.model) - pair_filter = min_pair_dist_requirement( + filter_enabled = add_training_data_requirements( + training_data, + training_params["training_data"], model_params, - float( - training_params.get("training_data", {}).get("min_pair_dist", 0.0) - ), + data_requirement, + lambda data, requirements: data.add_data_requirement(requirements), ) - min_pair_dist = 0.0 if pair_filter is None else float(pair_filter.default) - self.min_pair_dist_by_task["Default"] = min_pair_dist - if pair_filter is not None: - data_requirement.append(pair_filter) - training_data.add_data_requirement(data_requirement) + self.pair_filter_enabled_by_task["Default"] = filter_enabled if validation_data is not None: validation_data.add_data_requirement( self.loss.label_requirement + get_additional_data_requirement(self.model) ) # Preload and apply modifiers to all data before computing statistics - training_data.preload_and_modify_all_data_torch() + for source in iter_data_sources(training_data): + source.preload_and_modify_all_data_torch() if validation_data is not None: - validation_data.preload_and_modify_all_data_torch() + for source in iter_data_sources(validation_data): + source.preload_and_modify_all_data_torch() self.get_sample_func = single_model_stat( self.model, model_params.get("data_stat_nbatch", 10), training_data, self.stat_file_specs["Default"], - _min_pair_dist=min_pair_dist, + _filter_enabled=filter_enabled, _data_stat_full=model_params.get("data_stat_full", False), finetune_has_new_type=self.finetune_links["Default"].get_has_new_type() if self.finetune_links is not None @@ -571,19 +502,9 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: self.validation_data, self.valid_numb_batch, ) = get_data_loader(training_data, validation_data, training_params) - training_data.print_summary( - "training", - to_numpy_array(self.training_dataloader.sampler.weights) - if not isinstance(training_data, LmdbDataset) - else [1.0], - ) + self.training_dataloader.print_summary("training") if validation_data is not None: - validation_data.print_summary( - "validation", - to_numpy_array(self.validation_dataloader.sampler.weights) - if not isinstance(validation_data, LmdbDataset) - else [1.0], - ) + self.validation_dataloader.print_summary("validation") else: ( self.training_dataloader, @@ -599,30 +520,25 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: data_requirement += get_additional_data_requirement( self.model[model_key] ) - pair_filter = min_pair_dist_requirement( + filter_enabled = add_training_data_requirements( + training_data[model_key], + training_params["data_dict"][model_key]["training_data"], model_params["model_dict"][model_key], - float( - training_params["data_dict"][model_key] - .get("training_data", {}) - .get("min_pair_dist", 0.0) - ), - ) - min_pair_dist = ( - 0.0 if pair_filter is None else float(pair_filter.default) + data_requirement, + lambda data, requirements: data.add_data_requirement(requirements), ) - self.min_pair_dist_by_task[model_key] = min_pair_dist - if pair_filter is not None: - data_requirement.append(pair_filter) - training_data[model_key].add_data_requirement(data_requirement) + self.pair_filter_enabled_by_task[model_key] = filter_enabled if validation_data[model_key] is not None: validation_data[model_key].add_data_requirement( self.loss[model_key].label_requirement + get_additional_data_requirement(self.model[model_key]) ) # Preload and apply modifiers to all data before computing statistics - training_data[model_key].preload_and_modify_all_data_torch() + for source in iter_data_sources(training_data[model_key]): + source.preload_and_modify_all_data_torch() if validation_data[model_key] is not None: - validation_data[model_key].preload_and_modify_all_data_torch() + for source in iter_data_sources(validation_data[model_key]): + source.preload_and_modify_all_data_torch() _mt_user_observed = ( model_params["model_dict"][model_key] .get("info", {}) @@ -633,7 +549,7 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: model_params["model_dict"][model_key].get("data_stat_nbatch", 10), training_data[model_key], self.stat_file_specs[model_key], - _min_pair_dist=min_pair_dist, + _filter_enabled=filter_enabled, _data_stat_full=model_params["model_dict"][model_key].get( "data_stat_full", False ), @@ -667,40 +583,26 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: training_params["data_dict"][model_key], ) - training_data[model_key].print_summary( - f"training in {model_key}", - to_numpy_array(self.training_dataloader[model_key].sampler.weights) - if not isinstance(training_data[model_key], LmdbDataset) - else [1.0], + self.training_dataloader[model_key].print_summary( + f"training in {model_key}" ) if ( validation_data is not None and validation_data[model_key] is not None ): - validation_data[model_key].print_summary( - f"validation in {model_key}", - to_numpy_array( - self.validation_dataloader[model_key].sampler.weights - ) - if not isinstance(validation_data[model_key], LmdbDataset) - else [1.0], + self.validation_dataloader[model_key].print_summary( + f"validation in {model_key}" ) # Resolve training steps def epoch_length(model_key: str) -> int: """Return the batches this rank consumes in one epoch of a task.""" - _data = training_data[model_key] if self.multi_task else training_data _dataloader = ( self.training_dataloader[model_key] if self.multi_task else self.training_dataloader ) - if isinstance(_data, LmdbDataset): - return len(_dataloader) - return compute_total_numb_batch( - _data.index, - to_numpy_array(_dataloader.sampler.weights), - ) + return _dataloader.epoch_length schedule = resolve_step_schedule( training_params, @@ -733,13 +635,7 @@ def epoch_length(model_key: str) -> int: self.nonfinite_grad_guard = NonFiniteGradGuard() self.lr_schedule = get_lr(config["learning_rate"]) - # Radius the pair-clearance filter was registered with. The - # filter reads it as an enable flag, the frames carrying a margin that - # is already measured against that radius. It is resolved once with the - # data requirement, so the two cannot disagree. - self.has_min_pair_filter = any( - value > 0.0 for value in self.min_pair_dist_by_task.values() - ) + self.has_min_pair_filter = any(self.pair_filter_enabled_by_task.values()) if self.has_min_pair_filter: if self.multi_task: local_training_batch_attempts = max( @@ -1409,7 +1305,7 @@ def run(self) -> None: try: self._run() finally: - self._close_lmdb_loaders() + self._close_data_loaders() def _training_results( self, more_loss: dict[str, Any], task_key: str = "Default" @@ -1933,10 +1829,9 @@ def log_loss_valid(_task_key: str = "Default") -> dict: f"The profiling trace has been saved to: {self.profiling_file}" ) - def _close_lmdb_loaders(self) -> None: - """Release LMDB pipelines owned by training and validation loaders.""" + def _close_data_loaders(self) -> None: + """Release native pipelines and readers owned by both dataset splits.""" closed: set[int] = set() - datasets: dict[int, LmdbDataset] = {} for loaders in (self.training_dataloader, self.validation_dataloader): values = loaders.values() if isinstance(loaders, dict) else (loaders,) for loader in values: @@ -1946,10 +1841,6 @@ def _close_lmdb_loaders(self) -> None: close = getattr(loader, "close", None) if close is not None: close() - if isinstance(loader, LmdbBatchDataLoader): - datasets[id(loader.dataset)] = loader.dataset - for dataset in datasets.values(): - dataset.close() def _collect_checkpoint_states( self, @@ -2205,7 +2096,11 @@ def get_data( return {}, {}, {} batch_data = next(iterator) # === Filter frames with atoms too close (training only) === - filter_enabled = is_train and self.min_pair_dist_by_task[task_key] > 0.0 + filter_enabled = is_train and getattr( + batch_data, + "pair_filter_enabled", + self.pair_filter_enabled_by_task[task_key], + ) valid_mask = pair_margin_frame_mask(batch_data, filter_enabled) local_has_valid = valid_mask is None or bool(torch.any(valid_mask)) if not local_has_valid: diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py index 90a3a3ca0c..1267df0dfc 100644 --- a/deepmd/pt/utils/dataloader.py +++ b/deepmd/pt/utils/dataloader.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging import os +from collections.abc import ( + Iterator, +) from multiprocessing.dummy import ( Pool, ) @@ -35,16 +38,38 @@ from deepmd.pt.utils.dataset import ( DeepmdDataSetForLoader, ) +from deepmd.pt.utils.env import ( + DEVICE, +) +from deepmd.pt.utils.lmdb_dataset import ( + LmdbBatchDataLoader, + LmdbDataset, +) from deepmd.pt.utils.utils import ( mix_entropy, + to_numpy_array, ) from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_config import ( + DataConfig, + data_group_configs, + resolve_data_groups, + resolve_data_sources, +) +from deepmd.utils.data_group import ( + GroupedDataset, + GroupedTrainingDataLoader, + TrainingDataLoader, +) +from deepmd.utils.data_sampling import ( + data_child_seeds, + global_batch_counts, + system_probabilities, +) from deepmd.utils.data_system import ( print_summary, - prob_sys_size_ext, - process_sys_probs, ) log = logging.getLogger(__name__) @@ -254,20 +279,13 @@ def collate_batch(batch: list[dict[str, Any]]) -> dict[str, Any]: def get_weighted_sampler( - training_data: Any, prob_style: str, sys_prob: bool = False + training_data: Any, prob_style: str | list[float], sys_prob: bool = False ) -> WeightedRandomSampler: - if sys_prob is False: - if prob_style == "prob_uniform": - prob_v = 1.0 / float(training_data.__len__()) - probs = [prob_v for ii in range(training_data.__len__())] - else: # prob_sys_size;A:B:p1;C:D:p2 or prob_sys_size = prob_sys_size;0:nsys:1.0 - if prob_style == "prob_sys_size": - style = f"prob_sys_size;0:{len(training_data)}:1.0" - else: - style = prob_style - probs = prob_sys_size_ext(style, len(training_data), training_data.index) - else: - probs = process_sys_probs(prob_style, training_data.index) + probs = system_probabilities( + training_data.index, + sys_probs=prob_style if sys_prob else None, + auto_prob="prob_sys_size" if sys_prob else prob_style, + ) log.debug("Generated weighted sampler with prob array: " + str(probs)) # training_data.total_batch is the size of one epoch, you can increase it to avoid too many rebuilding of iterators len_sampler = training_data.total_batch * max(env.NUM_WORKERS, 1) @@ -294,3 +312,123 @@ def get_sampler_from_params(_data: Any, _params: dict[str, Any]) -> Any: else: _sampler = get_weighted_sampler(_data, "prob_sys_size") return _sampler + + +def create_dataset( + config: DataConfig, + type_map: list[str], + *, + seed: int | list[int] | None = None, + modifier: BaseModifier | None = None, +) -> DpLoaderSet | LmdbDataset | GroupedDataset: + """Construct native readers from the shared recursive group configuration.""" + groups = resolve_data_groups(config) + configs = data_group_configs(config) + datasets = [] + for group, params in zip(groups, configs, strict=True): + sources = resolve_data_sources(group) + has_lmdb = any(source.format == "lmdb" for source in sources) + children = [] + child_configs = [] + for source in sources: + if source.format == "deepmd": + child = DpLoaderSet( + list(source.systems), + source.batch_size, + type_map, + seed=seed, + modifier=modifier, + ) + else: + child = LmdbDataset( + source.systems[0], + type_map, + source.batch_size, + auto_prob_style=group.auto_prob, + ) + children.append(child) + child_configs.append( + {**params, "sys_probs": group.sys_probs if not has_lmdb else None} + ) + datasets.append( + children[0] + if len(children) == 1 + else GroupedDataset(children, child_configs, [None] * len(children)) + ) + return ( + datasets[0] + if len(datasets) == 1 + else GroupedDataset(datasets, configs, [group.prob for group in groups]) + ) + + +def create_dataloader( + data: DpLoaderSet | LmdbDataset | GroupedDataset, + config: DataConfig, + *, + seed: int | None = None, + rank: int = 0, + world_size: int = 1, +) -> TrainingDataLoader | GroupedTrainingDataLoader: + """Bind the native batch pipelines and their local sampling metadata.""" + if isinstance(data, GroupedDataset): + seeds = data_child_seeds( + 0 if seed is None and world_size > 1 else seed, len(data.children) + 1 + ) + children = [ + create_dataloader( + child, params, seed=child_seed, rank=rank, world_size=world_size + ) + for child, params, child_seed in zip( + data.children, data.configs, seeds[1:], strict=True + ) + ] + return GroupedTrainingDataLoader(data, children, seeds[0]) + if isinstance(data, LmdbDataset): + from deepmd.dpmodel.utils.lmdb_data import ( + DistributedLmdbBatchSampler, + LmdbBatchSampler, + ) + + options = {"shuffle": True, "seed": seed, "block_targets": data._block_targets} + sampler = ( + DistributedLmdbBatchSampler( + data._reader, rank=rank, world_size=world_size, **options + ) + if world_size > 1 + else LmdbBatchSampler(data._reader, **options) + ) + loader = LmdbBatchDataLoader(data, sampler, pin_memory=DEVICE.type != "cpu") + return TrainingDataLoader( + data, + loader, + np.ones(1, dtype=np.float64), + lambda: [len(loader)], + sampling_sizes=lambda: [sampler.total_batches], + ) + params = config[0] if isinstance(config, list) else config + sampler = get_sampler_from_params(data, params) + loader = DataLoader( + data, + sampler=sampler, + batch_size=None, + num_workers=env.NUM_WORKERS if dist.is_available() else 0, + drop_last=False, + collate_fn=lambda batch: batch, + pin_memory=DEVICE.type != "cpu", + ) + + def iterate() -> Iterator[dict[str, Any]]: + with torch.device("cpu"): + return iter(loader) + + return TrainingDataLoader( + data, + loader, + to_numpy_array(sampler.weights), + lambda: data.index, + sampling_sizes=lambda: global_batch_counts( + [len(system) for system in data.systems], data.batch_sizes + ), + iterator_factory=iterate, + ) diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 49ffd995cf..b787738d76 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -156,6 +156,8 @@ def make_stat_input( dataloaders: list[Any], nbatches: int, min_pair_dist: float = 0.0, + *, + filter_enabled_by_system: Sequence[bool] | None = None, ) -> list[dict[str, Any]]: """Pack data for statistics. @@ -172,6 +174,9 @@ def make_stat_input( positive, frames holding a pair inside its window are dropped before the statistics are accumulated, so they see the frames the optimizer sees. + filter_enabled_by_system + Optional source-specific filter policies for grouped datasets. The + registered requirements retain each source's filter radius. Returns ------- @@ -182,12 +187,17 @@ def make_stat_input( lst = [] log.info("Packing data for statistics from %d systems", len(datasets)) for system_index in range(len(datasets)): + filter_enabled = ( + min_pair_dist > 0.0 + if filter_enabled_by_system is None + else filter_enabled_by_system[system_index] + ) sys_stat = {} with torch.device("cpu"): dataloader = dataloaders[system_index] dataloader_size = len(dataloader) target_batches = min(nbatches, dataloader_size) - scan_limit = dataloader_size if min_pair_dist > 0.0 else target_batches + scan_limit = dataloader_size if filter_enabled else target_batches iterator = iter(dataloader) accepted_batches = 0 scanned_batches = 0 @@ -201,7 +211,7 @@ def make_stat_input( except StopIteration: break scanned_batches += 1 - frame_mask = pair_margin_frame_mask(stat_data, min_pair_dist > 0.0) + frame_mask = pair_margin_frame_mask(stat_data, filter_enabled) if frame_mask is not None and not torch.any(frame_mask): continue if frame_mask is not None: @@ -279,6 +289,8 @@ def scan_redu_stats( keys: Sequence[str], intensive: bool = False, min_pair_dist: float = 0.0, + *, + filter_enabled_by_system: Sequence[bool] | None = None, ) -> ReduScanResult: """Accumulate exact reduced-label statistics over every training frame. @@ -301,6 +313,9 @@ def scan_redu_stats( min_pair_dist Radius the training-frame filter was registered with. When positive, frames holding a pair inside its window are excluded. + filter_enabled_by_system + Optional source-specific filter policies for grouped datasets. The + registered requirements retain each source's filter radius. Returns ------- @@ -319,9 +334,14 @@ def scan_redu_stats( "Scanning all frames of %d systems for output statistics", len(dataloaders) ) with torch.device("cpu"): - for dataloader in dataloaders: + for system_index, dataloader in enumerate(dataloaders): + filter_enabled = ( + min_pair_dist > 0.0 + if filter_enabled_by_system is None + else filter_enabled_by_system[system_index] + ) for batch in _full_pass_loader(dataloader): - frame_mask = pair_margin_frame_mask(batch, min_pair_dist > 0.0) + frame_mask = pair_margin_frame_mask(batch, filter_enabled) if frame_mask is not None: if not torch.any(frame_mask): continue diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index e4ff2f813d..4e6af3a7cb 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -25,9 +25,6 @@ make_task_maps, print_data_summaries, ) -from deepmd.dpmodel.utils.lmdb_data import ( - is_lmdb, -) from deepmd.loggers import ( is_node_main_process, ) @@ -40,10 +37,16 @@ from deepmd.utils.data import ( min_pair_dist_requirement, ) +from deepmd.utils.data_config import ( + DataConfig, +) +from deepmd.utils.data_path import ( + process_systems, +) from deepmd.utils.data_system import ( DeepmdDataSystem, get_data, - process_systems, + get_neighbor_stat_data, ) from deepmd.utils.stat_file import ( StatFileSpec, @@ -93,87 +96,30 @@ def _update_changed_model_tensors( target_state_dict[key] = source_value -def _detect_lmdb_path(systems_raw: Any) -> str | None: - """Return the LMDB path when ``systems_raw`` is a scalar LMDB string. - - Returns ``None`` for non-LMDB inputs. Raises ``ValueError`` if - ``systems_raw`` is a list containing any LMDB path, so both - ``_get_neighbor_stat_data`` and ``_build_data_system`` fail with the - same clear message instead of the opaque error from - :func:`process_systems` / :class:`DeepmdData`. - """ - if isinstance(systems_raw, str) and is_lmdb(systems_raw): - return systems_raw - if isinstance(systems_raw, list) and any( - isinstance(s, str) and is_lmdb(s) for s in systems_raw - ): - raise ValueError( - "LMDB datasets must be passed as a scalar 'systems' string " - "(e.g. 'systems': '/path/to/data.lmdb'); list-form systems " - "with LMDB paths are not supported." - ) - return None - - -def _get_neighbor_stat_data( - dataset_params: dict[str, Any], - type_map: list[str] | None, -) -> Any: - """Return a data proxy suitable for ``BaseModel.update_sel`` (neighbor stat). - - Routes a scalar LMDB ``systems`` path through dpmodel's - ``make_neighbor_stat_data``; falls back to the legacy ``get_data`` for - npy/HDF5 directories. - """ - lmdb_path = _detect_lmdb_path(dataset_params.get("systems")) - if lmdb_path is not None: - from deepmd.dpmodel.utils.lmdb_data import ( - make_neighbor_stat_data, - ) - - return make_neighbor_stat_data(lmdb_path, type_map) - return get_data(dataset_params, 0, type_map, None) - - def _build_data_system( - dataset_params: dict[str, Any], + dataset_params: DataConfig, type_map: list[str], seed: int | None = None, rank: int = 0, world_size: int = 1, -) -> DeepmdDataSystem | LmdbDataSystem: - """Build a data system from dataset config, routing LMDB paths to LmdbDataSystem. - - A scalar ``systems`` value pointing at an LMDB directory triggers the - LMDB adapter; otherwise we fall through to the legacy - :class:`DeepmdDataSystem` path with system expansion. ``rank`` and - ``world_size`` shard LMDB training batches without changing legacy data - systems. - """ - systems_raw = dataset_params["systems"] - lmdb_path = _detect_lmdb_path(systems_raw) - if lmdb_path is not None: +) -> Any: + """Bind native LMDB batching to the common grouped-data factory.""" + + def build_lmdb( + path: str, batch_size: int | str, auto_prob: str | None + ) -> LmdbDataSystem: return LmdbDataSystem( - lmdb_path=lmdb_path, + lmdb_path=path, type_map=type_map, - batch_size=dataset_params["batch_size"], - auto_prob_style=dataset_params.get("auto_prob"), + batch_size=batch_size, + auto_prob_style=auto_prob, seed=seed, rank=rank, world_size=world_size, ) - systems = process_systems( - systems_raw, - patterns=dataset_params.get("rglob_patterns"), - ) - return DeepmdDataSystem( - systems=systems, - batch_size=dataset_params["batch_size"], - test_size=1, - type_map=type_map, - trn_all_set=True, - sys_probs=dataset_params.get("sys_probs"), - auto_prob_style=dataset_params.get("auto_prob", "prob_sys_size"), + + return get_data( + dataset_params, 0, type_map, None, lmdb_factory=build_lmdb, seed=seed ) @@ -201,7 +147,7 @@ def factory( ) -> tuple[DeepmdDataSystem | LmdbDataSystem, Any | None, StatFileSpec]: type_map = list(task_config.model_params["type_map"]) train_data = _build_data_system( - dict(task_config.training_data_params), + task_config.training_data_params, type_map, seed=data_seed, rank=rank, @@ -379,8 +325,8 @@ def update_neighbor_stat( min_nbor_dist: dict[str, float | None] = {} for task_config in iter_training_task_configs(config): type_map = task_config.model_params.get("type_map") - train_data = _get_neighbor_stat_data( - dict(task_config.training_data_params), type_map + train_data = get_neighbor_stat_data( + task_config.training_data_params, type_map ) updated_model_params, task_min_nbor_dist = BaseModel.update_sel( train_data, type_map, dict(task_config.model_params) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f1e7f27c58..f93f47c55e 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -61,6 +61,9 @@ change_model_out_bias_by_task, resolve_step_schedule, ) +from deepmd.dpmodel.train.data import ( + add_training_data_requirements, +) from deepmd.dpmodel.utils.batch import ( normalize_batch, split_batch, @@ -148,7 +151,9 @@ ) from deepmd.utils.data import ( DataRequirementItem, - min_pair_dist_requirement, +) +from deepmd.utils.data_group import ( + epoch_batch_counts, ) from deepmd.utils.data_system import ( DeepmdDataSystem, @@ -1909,7 +1914,7 @@ def __init__( reset_epoch_streams() self.valid_numb_batch_by_task: dict[str, int] = {} - self.min_pair_dist_by_task: dict[str, float] = {} + self.pair_filter_enabled_by_task: dict[str, bool] = {} for model_key in self.model_keys: data_requirement = list(self.losses[model_key].label_requirement) data_requirement += get_additional_data_requirement(self.models[model_key]) @@ -1918,19 +1923,14 @@ def __init__( if self.multi_task else training_params ) - pair_filter = min_pair_dist_requirement( - self.model_params_by_task[model_key], - float( - (task_params.get("training_data") or {}).get("min_pair_dist", 0.0) - ), - ) - self.min_pair_dist_by_task[model_key] = ( - 0.0 if pair_filter is None else float(pair_filter.default) - ) - # Only training frames are filtered, so only the training data - # pays for the derived margin. - self.training_data_by_task[model_key].add_data_requirements( - data_requirement + ([] if pair_filter is None else [pair_filter]) + self.pair_filter_enabled_by_task[model_key] = ( + add_training_data_requirements( + self.training_data_by_task[model_key], + task_params["training_data"], + self.model_params_by_task[model_key], + data_requirement, + lambda data, requirements: data.add_data_requirements(requirements), + ) ) if self.validation_data_by_task[model_key] is not None: self.validation_data_by_task[model_key].add_data_requirements( @@ -2007,9 +2007,9 @@ def __init__( def _make_sample( _d: DeepmdDataSystem = _data, _n: int = _nbatch, - _min_pair_dist: float = self.min_pair_dist_by_task[model_key], + _filter_enabled: bool = self.pair_filter_enabled_by_task[model_key], ) -> list[dict[str, np.ndarray]]: - return make_stat_input(_d, _n, _min_pair_dist) + return make_stat_input(_d, _n, _filter_enabled) self._sample_funcs[model_key] = _make_sample @@ -2794,11 +2794,15 @@ def _next_training_batch(self, data_sys: Any, task_key: str) -> dict[str, Any]: RuntimeError If one local reader pass of consecutive batches holds no valid frame. """ - filter_enabled = self.min_pair_dist_by_task[task_key] > 0.0 + filter_enabled = self.pair_filter_enabled_by_task[task_key] discarded = 0 while True: - batch = normalize_batch(data_sys.get_batch()) - frame_mask = pair_margin_frame_mask(batch, filter_enabled) + raw_batch = data_sys.get_batch() + batch = normalize_batch(raw_batch) + batch_filter_enabled = getattr( + raw_batch, "pair_filter_enabled", filter_enabled + ) + frame_mask = pair_margin_frame_mask(batch, batch_filter_enabled) if frame_mask is None or frame_mask.all(): return batch if frame_mask.any(): @@ -2826,23 +2830,22 @@ def _epoch_length(self, model_key: str) -> int: Returns ------- int - Number of steps covering one pass over the task's training data. + Nominal number of optimizer steps in one training epoch. Notes ----- - A data system reports ``nbatches[i]``, the batch count of system ``i``, - and ``sys_probs[i]``, the probability of drawing from that system, from - which ``compute_total_numb_batch`` derives the dataset-wide epoch - length ``ceil(max_i(nbatches[i] / sys_probs[i]))``. LMDB data reports - that global count while its sampler shards batches evenly across ranks; - legacy data systems remain replicated. In both cases one rank takes - ``ceil(total / world_size)`` optimizer steps per scheduled epoch. This - differs from the local reader pass length for replicated data, which - the data system exposes through ``get_batch_pass_length``. + The epoch length is ``ceil(max_i(local_batches[i] / sys_probs[i]))``. + Each LMDB contributes its sampler's rank-local workload, including + its own distributed padding. Replicated NumPy readers contribute + their global batch counts divided by world size. Sampling probabilities + use global source sizes and do not change with per-rank padding. + This scheduled workload differs from a replicated reader's pass length, + which is exposed separately through ``get_batch_pass_length``. """ data = self.training_data_by_task[model_key] - total = compute_total_numb_batch(data.nbatches, data.sys_probs) - return int(np.ceil(total / self.world_size)) + return compute_total_numb_batch( + epoch_batch_counts(data, self.world_size), data.sys_probs + ) # ------------------------------------------------------------------ # Distribution helpers diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py index 746b6df91a..377d18bcba 100644 --- a/deepmd/pt_expt/train/validation.py +++ b/deepmd/pt_expt/train/validation.py @@ -56,6 +56,9 @@ normalize_full_validation_metric, resolve_full_validation_start_step, ) +from deepmd.utils.data_group import ( + iter_data_sources, +) from deepmd.utils.eval_metrics import ( ENERGY_FULL_VALIDATION_PROFILE, SPIN_FULL_VALIDATION_PROFILE, @@ -294,11 +297,11 @@ def __init__( self.checkpoint_dir.mkdir(parents=True, exist_ok=True) self._initialize_best_checkpoints(restart_training=restart_training) - # Lazily-populated full test snapshot for LMDB validation. Frames are + # Lazily-populated full test snapshots for LMDB validation. Frames are # evaluated in atom-count and label-availability groups so each stack # has consistent shapes and scalar find_* flags. Reuse the decoded - # snapshot on subsequent validation calls. - self._lmdb_test_data: LmdbTestData | None = None + # snapshots on subsequent validation calls, independently for each source. + self._lmdb_snapshots: dict[int, LmdbTestData] = {} def should_run(self, display_step: int) -> bool: """Check whether the current step should trigger full validation.""" @@ -457,28 +460,27 @@ def _iter_validation_data_systems(self) -> Iterator[Any]: - A loader set owns ``systems``, each wrapping a ``DeepmdData`` in ``data_system``. """ - validation_data = self.validation_data - if hasattr(validation_data, "_reader"): - lmdb_test_data = self._get_lmdb_test_data_snapshot(validation_data) - for nloc in sorted(lmdb_test_data.nloc_groups): - for indices in lmdb_test_data.availability_groups( - lmdb_test_data.nloc_groups[nloc] - ): - yield LmdbTestDataNlocView(lmdb_test_data, nloc, indices) - return - - if hasattr(validation_data, "data_systems"): - yield from validation_data.data_systems - return - - for dataset in validation_data.systems: - data_system = getattr(dataset, "data_system", None) - if data_system is None: - raise TypeError( - "Full validation expects each dataset in validation_data.systems " - f"to expose a `data_system`, got {type(dataset)!r}." + for source in iter_data_sources(self.validation_data): + if hasattr(source, "_reader"): + snapshot = self._get_lmdb_test_data_snapshot(source) + yield from ( + LmdbTestDataNlocView(snapshot, nloc, indices) + for nloc in sorted(snapshot.nloc_groups) + for indices in snapshot.availability_groups( + snapshot.nloc_groups[nloc] + ) ) - yield data_system + elif hasattr(source, "data_systems"): + yield from source.data_systems + else: + for dataset in source.systems: + data_system = getattr(dataset, "data_system", None) + if data_system is None: + raise TypeError( + "Full validation expects each dataset in validation_data.systems " + f"to expose a `data_system`, got {type(dataset)!r}." + ) + yield data_system def _get_lmdb_test_data_snapshot(self, lmdb_dataset: Any) -> LmdbTestData: """Build (once) and return the cached LMDB test snapshot. @@ -488,8 +490,9 @@ def _get_lmdb_test_data_snapshot(self, lmdb_dataset: Any) -> LmdbTestData: the full-validation snapshot sees exactly the same fields and dtypes as training batches. """ - if self._lmdb_test_data is not None: - return self._lmdb_test_data + key = id(lmdb_dataset) + if key in self._lmdb_snapshots: + return self._lmdb_snapshots[key] reader = getattr(lmdb_dataset, "_reader", None) lmdb_path = getattr(lmdb_dataset, "lmdb_path", None) @@ -512,14 +515,15 @@ def _get_lmdb_test_data_snapshot(self, lmdb_dataset: Any) -> LmdbTestData: f"{type(lmdb_dataset)!r}." ) - self._lmdb_test_data = LmdbTestData( + snapshot = LmdbTestData( lmdb_path, type_map=list(type_map), shuffle_test=False, ) if data_requirements: - self._lmdb_test_data.add_data_requirement(data_requirements) - return self._lmdb_test_data + snapshot.add_data_requirement(data_requirements) + self._lmdb_snapshots[key] = snapshot + return snapshot def _evaluate_system( self, data_system: DeepmdData diff --git a/deepmd/pt_expt/utils/lmdb_dataset.py b/deepmd/pt_expt/utils/lmdb_dataset.py index 9483b3ef66..72d407f2b3 100644 --- a/deepmd/pt_expt/utils/lmdb_dataset.py +++ b/deepmd/pt_expt/utils/lmdb_dataset.py @@ -13,6 +13,8 @@ Any, ) +import numpy as np + from deepmd.dpmodel.utils.lmdb_data import ( DistributedLmdbBatchSampler, LmdbBatchIterator, @@ -236,6 +238,23 @@ def get_batch_pass_length(self) -> int: """Return the local pass length belonging to the last returned batch.""" return self._batch_iterator.batch_pass_length + def get_epoch_batch_counts(self, world_size: int) -> np.ndarray: + """Return the actual local workload after distributed batch padding.""" + return np.asarray([len(self._sampler)], dtype=np.float64) + + def get_ntypes(self) -> int: + return len(self.type_map) + + def get_type_map(self) -> list[str]: + return self.type_map + + def get_batch_size(self) -> list[int]: + return [self._reader.batch_size] + + @property + def system_dirs(self) -> list[str]: + return [self._reader.lmdb_path] + @property def nbatches(self) -> list[int]: """Return the global batch count of one full pass.""" diff --git a/deepmd/tf/entrypoints/change_bias.py b/deepmd/tf/entrypoints/change_bias.py index b333fc470e..78e2ebd36c 100644 --- a/deepmd/tf/entrypoints/change_bias.py +++ b/deepmd/tf/entrypoints/change_bias.py @@ -47,6 +47,7 @@ ) from deepmd.utils.data_system import ( DeepmdDataSystem, + get_data, ) __all__ = ["change_bias"] @@ -218,13 +219,22 @@ def _change_bias_checkpoint_file( if stop_batch is None and num_epoch is not None: if num_epoch <= 0: raise ValueError("training.num_epoch must be positive.") - # Apply sys_probs and auto_prob from original training config - # to ensure stop_batch calculation matches the original training training_data_config = training_params.get("training_data", {}) - sys_probs = training_data_config.get("sys_probs", None) - auto_prob = training_data_config.get("auto_prob", "prob_sys_size") - data.set_sys_probs(sys_probs=sys_probs, auto_prob_style=auto_prob) - total_numb_batch = compute_total_numb_batch(data.nbatches, data.sys_probs) + if isinstance(training_data_config, list): + # Group probabilities and batch sizes define the training schedule; + # the independent bias dataset has no corresponding group indices. + schedule_data = get_data( + training_data_config, 0, jdata["model"].get("type_map"), None + ) + else: + data.set_sys_probs( + sys_probs=training_data_config.get("sys_probs"), + auto_prob_style=training_data_config.get("auto_prob", "prob_sys_size"), + ) + schedule_data = data + total_numb_batch = compute_total_numb_batch( + schedule_data.nbatches, schedule_data.sys_probs + ) if total_numb_batch <= 0: raise ValueError("Total number of training batches must be positive.") stop_batch = int(np.ceil(num_epoch * total_numb_batch)) diff --git a/deepmd/tf2/entrypoints/train.py b/deepmd/tf2/entrypoints/train.py index a026bcc03f..3350e98c1c 100644 --- a/deepmd/tf2/entrypoints/train.py +++ b/deepmd/tf2/entrypoints/train.py @@ -159,7 +159,7 @@ def factory( type_map = list(task_config.model_params.get("type_map", [])) ipt_type_map = type_map if type_map else None train_data = get_data( - dict(task_config.training_data_params), + task_config.training_data_params, None, ipt_type_map, None, @@ -244,7 +244,7 @@ def update_sel( for task_config in iter_training_task_configs(jdata): type_map = task_config.model_params.get("type_map") train_data = get_data( - dict(task_config.training_data_params), + task_config.training_data_params, 0, type_map, None, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 9f9c98c26c..caffb5de52 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -36,6 +36,9 @@ DEFAULT_FRACTION_INNER, DEFAULT_FRACTION_OUTER, ) +from deepmd.utils.data_config import ( + data_group_configs, +) from deepmd.utils.eval_metrics import ( ENERGY_FULL_VALIDATION_PROFILE, FULL_VALIDATION_PROFILES, @@ -5773,15 +5776,15 @@ def loss_args() -> list[Argument]: # --- Training configurations: --- # -def training_data_args() -> list[ - Argument -]: # ! added by Ziyao: new specification style for data systems. +def training_data_args() -> Argument: link_sys = make_link("systems", "training/training_data/systems") doc_systems = ( "The data systems for training. " "This key can be a list or a str. " "When provided as a string, it can be a system directory path (containing 'type.raw') or a parent directory path to recursively search for all system subdirectories. " - "When provided as a list, each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories." + "When provided as a list, each string item is expanded recursively in input order. " + "The PyTorch and PyTorch Exportable backends also recognize LMDB datasets, " + "either directly or beneath a parent directory." ) doc_patterns = ( "The customized patterns used in `rglob` to collect all training systems. " @@ -5800,7 +5803,12 @@ def training_data_args() -> list[ - "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems()\n\n\ - "prob_sys_size" : the probability of a system is proportional to the number of batches in the system\n\n\ - "prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;..." : the list of systems is divided into blocks. A block is specified by `stt_idx:end_idx:weight`, where `stt_idx` is the starting index of the system, `end_idx` is then ending (not including) index of the system, the probabilities of the systems in this block sums up to `weight`, and the relatively probabilities within this block is proportional to the number of batches in the system.' - doc_sys_probs = "A list of float if specified. Should be of the same length as `systems`, specifying the probability of each system." + doc_sys_probs = ( + "A list of float if specified. Should be of the same length as `systems`, " + "specifying the probability of each system. A group expanding to multiple " + "sources including LMDB must omit this option; use separate training_data " + "groups with prob to assign source probabilities." + ) doc_min_pair_dist = ( "Minimum pairwise atomic distance threshold in Å. " "Frames containing any atom pair closer than this distance are excluded " @@ -5842,7 +5850,10 @@ def training_data_args() -> list[ str, optional=True, default="prob_sys_size", - doc=doc_auto_prob_style, + doc=doc_auto_prob_style + + " When a group expands to multiple sources including LMDB, only " + "prob_sys_size is supported. Use separate groups to configure " + "individual LMDB system probabilities.", alias=[ "auto_prob_style", ], @@ -5862,12 +5873,28 @@ def training_data_args() -> list[ default=0.0, doc=supported_backends("pt", "pt_expt") + doc_min_pair_dist, ), + Argument( + "prob", + [float, None], + optional=True, + default=None, + doc="The probability of selecting a training batch from this entire " + "group. Applies when training_data is a list of group dictionaries. " + "Assigned probabilities must be non-negative and sum to at most one. " + "Groups without prob share the remaining probability in proportion " + "to their batch counts. If every group specifies prob, their sum " + "must be one. Within-group auto_prob and sys_probs remain independent.", + ), ] - doc_training_data = "Configurations of training data." + doc_training_data = ( + "Training data configured as one dictionary or a nonempty list of group " + "dictionaries. Each group has its own systems, batch_size and within-group " + "sampling rules; prob optionally controls the group's total probability." + ) return Argument( "training_data", - dict, + [dict, list[dict]], optional=True, sub_fields=args, sub_variants=[], @@ -5884,7 +5911,9 @@ def validation_data_args() -> list[ "The data systems for validation. " "This key can be a list or a str. " "When provided as a string, it can be a system directory path (containing 'type.raw') or a parent directory path to recursively search for all system subdirectories. " - "When provided as a list, each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories." + "When provided as a list, each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories. " + "The PyTorch and PyTorch Exportable backends also recognize LMDB datasets, " + "including multiple LMDB paths and parent directories containing LMDB datasets." ) doc_patterns = ( "The customized patterns used in `rglob` to collect all validation systems. " @@ -5901,7 +5930,11 @@ def validation_data_args() -> list[ - "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems()\n\n\ - "prob_sys_size" : the probability of a system is proportional to the number of batches in the system\n\n\ - "prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;..." : the list of systems is divided into blocks. A block is specified by `stt_idx:end_idx:weight`, where `stt_idx` is the starting index of the system, `end_idx` is then ending (not including) index of the system, the probabilities of the systems in this block sums up to `weight`, and the relatively probabilities within this block is proportional to the number of batches in the system.' - doc_sys_probs = "A list of float if specified. Should be of the same length as `systems`, specifying the probability of each system." + doc_sys_probs = ( + "A list of float if specified. Should be of the same length as `systems`, " + "specifying the probability of each system. Must be omitted when the " + "validation data expands to multiple sources including LMDB." + ) doc_numb_btch = "An integer that specifies the number of batches to be sampled for each validation period." args = [ @@ -5928,7 +5961,9 @@ def validation_data_args() -> list[ str, optional=True, default="prob_sys_size", - doc=doc_auto_prob_style, + doc=doc_auto_prob_style + + " Only prob_sys_size is supported when systems expands to " + "multiple sources including LMDB.", alias=[ "auto_prob_style", ], @@ -6883,6 +6918,7 @@ def gen_json_schema(multi_task: bool = False) -> str: JSON schema. """ from deepmd.utils.json_schema import ( + with_data_groups, with_model_presets, ) @@ -6894,7 +6930,9 @@ def gen_json_schema(multi_task: bool = False) -> str: ) return json.dumps( with_model_presets( - generate_json_schema(arg), generate_json_schema(model_args()), multi_task + with_data_groups(generate_json_schema(arg), multi_task), + generate_json_schema(model_args()), + multi_task, ) ) @@ -6973,6 +7011,23 @@ def normalize( """Normalize config values and optionally run strict schema checks.""" base = Argument("base", dict, gen_args(multi_task=multi_task)) data = base.normalize_value(data, trim_pattern="_*") + training = data.get("training", {}) + task_configs = training.get("data_dict", {}).values() if multi_task else (training,) + group_argument = training_data_args() + for task in task_configs: + groups = task.get("training_data") + if groups is None: + continue + data_group_configs(groups) + if isinstance(groups, list): + normalized = [ + group_argument.normalize_value(group, trim_pattern="_*") + for group in groups + ] + if check: + for group in normalized: + group_argument.check_value(group, strict=True) + task["training_data"] = normalized if check: base.check_value(data, strict=True) validate_full_validation_config(data, multi_task=multi_task) diff --git a/deepmd/utils/data_config.py b/deepmd/utils/data_config.py new file mode 100644 index 0000000000..20b6b1fb45 --- /dev/null +++ b/deepmd/utils/data_config.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Format-independent configuration of training data groups.""" + +import logging +from collections.abc import ( + Mapping, + Sequence, +) +from dataclasses import ( + dataclass, +) +from typing import ( + Any, + TypeAlias, +) + +from deepmd.utils.data_path import ( + DatasetFormat, + is_lmdb, + process_systems, +) +from deepmd.utils.data_sampling import ( + group_probabilities, +) +from deepmd.utils.path import ( + DPPath, +) + +DataConfig: TypeAlias = Mapping[str, Any] | list[Mapping[str, Any]] +BatchSize: TypeAlias = int | str | list[int] +log = logging.getLogger(__name__) + + +def check_type_map_consistency(type_maps: Sequence[list[str] | None]) -> list[str]: + """Return the longest mutually consistent type map.""" + result: list[str] = [] + for type_map in type_maps: + if type_map is None: + continue + common = min(len(result), len(type_map)) + if result[:common] != type_map[:common]: + raise RuntimeError( + f"Inconsistent type map: {result!s} {type_map!s} in different data systems. " + "If you didn't set model/type_map, please set it, " + "since the type map of the model cannot be decided by data." + ) + if len(type_map) > len(result): + result = list(type_map) + return result + + +def data_group_configs(config: DataConfig) -> list[Mapping[str, Any]]: + """View a legacy configuration as one group without mutating its input.""" + groups = [config] if isinstance(config, Mapping) else config + if not isinstance(groups, list) or not groups: + raise ValueError( + "Data configuration must be a dictionary or a nonempty list of dictionaries." + ) + if any(not isinstance(group, Mapping) for group in groups): + raise ValueError("Every data group must be a dictionary.") + group_probabilities([1] * len(groups), [group.get("prob") for group in groups]) + return groups + + +@dataclass(frozen=True) +class DataGroupConfig: + """One configured group after recursive dataset discovery. + + ``prob`` controls the whole group's batch-selection probability. + ``auto_prob`` and ``sys_probs`` retain their native within-group + semantics. An LMDB's internal system numbering is local to that file; + custom ``auto_prob`` rules therefore require a single LMDB source. + """ + + systems: tuple[str, ...] + batch_size: BatchSize + auto_prob: str | None + sys_probs: list[float] | None + prob: float | None + + +@dataclass(frozen=True) +class DataSourceConfig: + """One native batch stream within a configured group.""" + + format: DatasetFormat + systems: tuple[str, ...] + batch_size: BatchSize + + +def resolve_data_sources(group: DataGroupConfig) -> list[DataSourceConfig]: + """Partition a group into native readers after whole-system filtering. + + NumPy systems share a native reader; each LMDB owns one. NumPy atom-count + filtering precedes this partition so an empty NumPy subset does not reject + a group with retained LMDB data. LMDB filtering remains frame-local. + """ + numpy_indices = [ + index for index, path in enumerate(group.systems) if not is_lmdb(path) + ] + numpy_batch_size = group.batch_size + if isinstance(numpy_batch_size, str) and numpy_batch_size.startswith("filter:"): + limit = int(numpy_batch_size.split(":")[1]) + retained = [ + index + for index in numpy_indices + if (DPPath(group.systems[index]) / "type.raw").load_txt(ndmin=1).size + <= limit + ] + if len(retained) != len(numpy_indices): + log.warning( + "Remove %d systems with more than %d atoms", + len(numpy_indices) - len(retained), + limit, + ) + numpy_indices = retained + numpy_batch_size = f"max:{limit}" + elif isinstance(numpy_batch_size, list): + numpy_batch_size = [numpy_batch_size[index] for index in numpy_indices] + sources = [] + if numpy_indices: + sources.append( + DataSourceConfig( + "deepmd", + tuple(group.systems[index] for index in numpy_indices), + numpy_batch_size, + ) + ) + for index, path in enumerate(group.systems): + if is_lmdb(path): + batch_size = group.batch_size + if isinstance(batch_size, list): + batch_size = batch_size[index] + sources.append(DataSourceConfig("lmdb", (path,), batch_size)) + if not sources: + raise ValueError( + f"No system left in data group after applying {group.batch_size!r}." + ) + return sources + + +def resolve_data_groups( + config: DataConfig, + *, + formats: tuple[DatasetFormat, ...] = ("deepmd", "lmdb"), +) -> list[DataGroupConfig]: + """Resolve every group's paths and validate format-dependent options.""" + result = [] + for index, group in enumerate(data_group_configs(config)): + roots = ( + [group["systems"]] + if isinstance(group["systems"], str) + else group["systems"] + ) + if "lmdb" not in formats and any(is_lmdb(path) for path in roots): + raise ValueError("LMDB training data requires the pt or pt_expt backend.") + systems = process_systems( + group["systems"], group.get("rglob_patterns"), formats=formats + ) + if not systems: + raise ValueError(f"Data group {index} contains no supported data systems.") + if len(systems) > 1 and any(is_lmdb(path) for path in systems): + if group.get("auto_prob", "prob_sys_size") != "prob_sys_size": + raise ValueError( + f"Data group {index} expands to multiple sources including LMDB, " + "so auto_prob must be prob_sys_size. Use separate training_data groups " + "to configure each LMDB's internal system probabilities." + ) + if group.get("sys_probs") is not None: + raise ValueError( + f"Data group {index} expands to multiple sources including LMDB, " + "so sys_probs is not supported. Use training_data groups with prob " + "to assign source probabilities." + ) + batch_size = group.get("batch_size", "auto") + if isinstance(batch_size, list) and len(batch_size) != len(systems): + raise ValueError( + f"Data group {index} has {len(systems)} expanded systems but " + f"{len(batch_size)} batch_size entries." + ) + result.append( + DataGroupConfig( + systems=tuple(systems), + batch_size=batch_size, + auto_prob=group.get("auto_prob"), + sys_probs=group.get("sys_probs"), + prob=group.get("prob"), + ) + ) + return result diff --git a/deepmd/utils/data_group.py b/deepmd/utils/data_group.py new file mode 100644 index 0000000000..eb0f5c6fc7 --- /dev/null +++ b/deepmd/utils/data_group.py @@ -0,0 +1,525 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Batch-level composition of independently configured data systems.""" + +from __future__ import ( + annotations, +) + +import logging +from typing import ( + TYPE_CHECKING, + Any, +) + +import numpy as np + +from deepmd.dpmodel.utils.training_utils import ( + compute_total_numb_batch, +) +from deepmd.utils.data_config import ( + check_type_map_consistency, +) +from deepmd.utils.data_sampling import ( + data_child_seeds, + group_probabilities, +) +from deepmd.utils.model_stat import ( + _get_stat_batch, + _get_stat_nsystems, + _get_stat_pass_length, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + Iterator, + Mapping, + Sequence, + ) + + from deepmd.utils.data import ( + DataRequirementItem, + ) + +log = logging.getLogger(__name__) + + +class DataBatch(dict[str, Any]): + """Carry source-specific execution context without adding model input fields.""" + + def __init__(self, values: Mapping[str, Any], pair_filter_enabled: bool) -> None: + super().__init__(values) + self.pair_filter_enabled = pair_filter_enabled + + +def iter_data_sources(data: Any) -> Iterator[Any]: + """Visit each native data source of a composed dataset in order.""" + children = getattr(data, "children", None) + if children is not None: + for child in children: + yield from iter_data_sources(child) + else: + yield data + + +def epoch_batch_counts(data: Any, world_size: int) -> np.ndarray: + """Resolve local workloads while preserving the replicated-reader contract.""" + get_counts = getattr(data, "get_epoch_batch_counts", None) + if get_counts is not None: + return get_counts(world_size) + return np.asarray(data.nbatches, dtype=np.float64) / world_size + + +class GroupedDataSystem: + """Select whole batches while preserving every child's batching rules. + + The child owns its within-system probabilities, batching, and I/O. This + node owns only the probability of selecting that child. Training and + statistical system indices are separate flattened views: an LMDB is one + training source but may expose several homogeneous statistical systems. + """ + + def __init__( + self, + children: Sequence[Any], + probabilities: Sequence[float | None], + *, + seed: int | None = None, + ) -> None: + if not children or len(children) != len(probabilities): + raise ValueError("Data groups require one probability per nonempty child.") + self.children = tuple(children) + self._assigned_probabilities = tuple(probabilities) + self._rng = np.random.default_rng(data_child_seeds(seed, 1)[0]) + self._probabilities: np.ndarray | None = None + self._cdf: np.ndarray | None = None + self._filter_enabled: tuple[bool, ...] | None = None + self._remaining = 0 + self._pass_length = 0 + self._align_metadata() + + def _align_metadata(self) -> None: + """Preserve dataset-wide NumPy type-map and atom-count contracts.""" + from deepmd.utils.data_system import ( + DeepmdDataSystem, + ) + + sources = list(iter_data_sources(self)) + self.type_map = check_type_map_consistency( + [source.get_type_map() for source in sources] + ) + self._ntypes = max(source.get_ntypes() for source in sources) + numpy_sources = [ + source for source in sources if isinstance(source, DeepmdDataSystem) + ] + if len({source.mixed_type for source in numpy_sources}) > 1: + raise ValueError("All NumPy data systems must agree on mixed_type.") + for source in numpy_sources: + source.sys_ntypes = self._ntypes + source.natoms_vec = [ + system.get_natoms_vec(self._ntypes).astype(int) + for system in source.data_systems + ] + self.mixed_type = any(source.mixed_type for source in sources) + + @property + def component_probabilities(self) -> np.ndarray: + """Return this sampling pass's probabilities over child datasets.""" + if self._probabilities is None: + self._probabilities = group_probabilities( + [sum(child.nbatches) for child in self.children], + self._assigned_probabilities, + ) + return self._probabilities + + @property + def nbatches(self) -> list[int]: + """Return global batch counts in the flattened training-system view.""" + return [count for child in self.children for count in child.nbatches] + + @property + def sys_probs(self) -> np.ndarray: + """Return global probabilities without duplicating child sampling state.""" + return np.concatenate( + [ + probability * np.asarray(child.sys_probs) + for child, probability in zip( + self.children, self.component_probabilities, strict=True + ) + ] + ) + + def set_filter_enabled(self, enabled: Sequence[bool]) -> None: + """Bind the configured frame-filter policy to each child stream.""" + if len(enabled) != len(self.children): + raise ValueError("Each data group requires one frame-filter policy.") + self._filter_enabled = tuple(enabled) + + def _batch_context(self, batch: dict[str, Any], child: int) -> dict[str, Any]: + if self._filter_enabled is None: + return batch + return DataBatch(batch, self._filter_enabled[child]) + + def _invalidate_sampling(self) -> None: + self._probabilities = None + self._cdf = None + self._remaining = 0 + + def _start_pass(self) -> None: + self._probabilities = None + probabilities = self.component_probabilities + self._cdf = np.cumsum(probabilities) + self._cdf[-1] = 1.0 + self._pass_length = self._reader_pass_length() + self._remaining = self._pass_length + + @staticmethod + def _locate(index: int, sizes: Sequence[int]) -> tuple[int, int]: + if index < 0: + raise IndexError(f"System index {index} must be non-negative.") + for child, size in enumerate(sizes): + if index < size: + return child, index + index -= size + raise IndexError("System index is outside the grouped dataset.") + + def get_batch(self, sys_idx: int | None = None) -> dict[str, Any]: + """Draw a child batch, or address one flattened training system.""" + if sys_idx is not None: + child, local = self._locate( + sys_idx, [source.get_nsystems() for source in self.children] + ) + return self._batch_context( + self.children[child].get_batch(sys_idx=local), child + ) + if self._remaining == 0: + self._start_pass() + assert self._cdf is not None + index = int(np.searchsorted(self._cdf, self._rng.random(), side="right")) + self._remaining -= 1 + return self._batch_context(self.children[index].get_batch(), index) + + def _reader_pass_length(self) -> int: + return compute_total_numb_batch( + [child.get_batch_pass_length() for child in self.children], + self.component_probabilities, + ) + + def get_batch_pass_length(self) -> int: + """Return the local reader-pass length belonging to the current draw.""" + return self._pass_length or self._reader_pass_length() + + def get_epoch_batch_counts(self, world_size: int) -> np.ndarray: + """Expose local optimizer workloads, including each leaf's sharding.""" + return np.concatenate( + [epoch_batch_counts(child, world_size) for child in self.children] + ) + + def get_stat_nsystems(self) -> int: + return sum(_get_stat_nsystems(child) for child in self.children) + + def get_stat_numb_batches(self, sys_idx: int) -> int: + child, local = self._locate( + sys_idx, [_get_stat_nsystems(source) for source in self.children] + ) + return _get_stat_pass_length(self.children[child], local) + + def get_stat_batch(self, sys_idx: int) -> dict[str, Any]: + child, local = self._locate( + sys_idx, [_get_stat_nsystems(source) for source in self.children] + ) + return self._batch_context(_get_stat_batch(self.children[child], local), child) + + def add_data_requirements(self, requirements: list[DataRequirementItem]) -> None: + for child in self.children: + child.add_data_requirements(requirements) + self._invalidate_sampling() + + def use_ragged_batches(self, ragged: bool) -> None: + for child in self.children: + configure = getattr(child, "use_ragged_batches", None) + if configure is not None: + configure(ragged) + self._invalidate_sampling() + + def set_frame_transform(self, transform: Any) -> None: + for child in self.children: + configure = getattr(child, "set_frame_transform", None) + if configure is None: + raise ValueError("Frame transforms require LMDB data in every group.") + configure(transform) + + def close(self) -> None: + for child in self.children: + close = getattr(child, "close", None) + if close is not None: + close() + + def get_nsystems(self) -> int: + return sum(child.get_nsystems() for child in self.children) + + def get_ntypes(self) -> int: + return self._ntypes + + def get_type_map(self) -> list[str]: + return self.type_map + + def get_nbatches(self) -> list[int]: + return self.nbatches + + def get_batch_size(self) -> list[int]: + return [size for child in self.children for size in child.get_batch_size()] + + def get_sys(self, index: int) -> Any: + child, local = self._locate( + index, [source.get_nsystems() for source in self.children] + ) + return self.children[child].get_sys(local) + + def get_data_dict(self, ii: int = 0) -> dict[str, Any]: + child, local = self._locate( + ii, [source.get_nsystems() for source in self.children] + ) + return self.children[child].get_data_dict(local) + + @property + def data_systems(self) -> list[Any]: + return [system for child in self.children for system in child.data_systems] + + @property + def system_dirs(self) -> list[str]: + return [path for child in self.children for path in child.system_dirs] + + def add_dict(self, requirements: dict[str, dict[str, Any]]) -> None: + for child in self.children: + child.add_dict(requirements) + + def add(self, *args: Any, **kwargs: Any) -> None: + for child in self.children: + child.add(*args, **kwargs) + + def reduce(self, key_out: str, key_in: str) -> None: + for child in self.children: + child.reduce(key_out, key_in) + + def print_summary(self, name: str, prob: Any = None) -> None: + for index, (child, probability) in enumerate( + zip(self.children, self.component_probabilities, strict=True) + ): + label = f"{name}/group_{index}" + log.info("%s: batch probability %.6g", label, probability) + child.print_summary(label) + + +class GroupedDataset: + """Own native loader datasets and expose their combined statistics view.""" + + def __init__( + self, + children: Sequence[Any], + configs: Sequence[Mapping[str, Any]], + probabilities: Sequence[float | None], + ) -> None: + if not children or not len(children) == len(configs) == len(probabilities): + raise ValueError( + "A dataset group requires matching datasets, configs, and probabilities." + ) + self.children = tuple(children) + self.configs = tuple(configs) + self.probability_specs = tuple(probabilities) + self.filter_enabled: tuple[bool, ...] | None = None + + def __len__(self) -> int: + """Preserve native dataset sizes for default multi-task probabilities.""" + return sum(len(child) for child in self.children) + + def set_filter_enabled(self, enabled: Sequence[bool]) -> None: + """Bind the configured frame-filter policy to each child stream.""" + if len(enabled) != len(self.children): + raise ValueError("Each data group requires one frame-filter policy.") + self.filter_enabled = tuple(enabled) + + @property + def stat_filter_enabled(self) -> list[bool] | None: + if self.filter_enabled is None: + return None + return [ + enabled + for child, enabled in zip(self.children, self.filter_enabled, strict=True) + for _ in child.systems + ] + + @property + def systems(self) -> list[Any]: + return [system for child in self.children for system in child.systems] + + @property + def dataloaders(self) -> list[Any]: + return [loader for child in self.children for loader in child.dataloaders] + + @property + def sampler_list(self) -> list[Any]: + return [sampler for child in self.children for sampler in child.sampler_list] + + @property + def index(self) -> list[int]: + return [count for child in self.children for count in child.index] + + @property + def total_batch(self) -> int: + return sum(self.index) + + def add_data_requirement(self, requirements: list[DataRequirementItem]) -> None: + for child in self.children: + child.add_data_requirement(requirements) + + def close(self) -> None: + for child in self.children: + close = getattr(child, "close", None) + if close is not None: + close() + + +class TrainingDataLoader: + """Expose native loader iteration, sampling metadata, and resource ownership.""" + + def __init__( + self, + dataset: Any, + loader: Any, + probabilities: np.ndarray, + batch_counts: Callable[[], Sequence[int]], + *, + sampling_sizes: Callable[[], Sequence[int]] | None = None, + iterator_factory: Callable[[], Iterator[dict[str, Any]]] | None = None, + ) -> None: + self.dataset = dataset + self.loader = loader + self.probabilities = np.asarray(probabilities, dtype=np.float64) + self._batch_counts = batch_counts + self._sampling_sizes = sampling_sizes or batch_counts + self._iterator_factory = iterator_factory or loader.__iter__ + + @property + def batch_counts(self) -> np.ndarray: + return np.asarray(self._batch_counts(), dtype=np.float64) + + @property + def sampling_sizes(self) -> np.ndarray: + """Return natural source masses, independently of per-rank padding.""" + return np.asarray(self._sampling_sizes(), dtype=np.float64) + + @property + def system_probabilities(self) -> np.ndarray: + return self.probabilities + + @property + def epoch_length(self) -> int: + return compute_total_numb_batch(self.batch_counts, self.system_probabilities) + + def __iter__(self) -> Iterator[dict[str, Any]]: + return self._iterator_factory() + + def __len__(self) -> int: + return len(self.loader) + + def print_summary(self, name: str) -> None: + self.dataset.print_summary(name, self.probabilities) + + def close(self) -> None: + close_loader = getattr(self.loader, "close", None) + if close_loader is not None: + close_loader() + close_dataset = getattr(self.dataset, "close", None) + if close_dataset is not None: + close_dataset() + + +class GroupedTrainingDataLoader: + """Interleave native batch streams without moving decoding into a wrapper worker.""" + + def __init__( + self, + dataset: GroupedDataset, + children: Sequence[TrainingDataLoader | GroupedTrainingDataLoader], + seed: int | None, + ) -> None: + self.dataset = dataset + self.children = tuple(children) + self._rng = np.random.default_rng(seed) + self._iterators: dict[int, Iterator[dict[str, Any]]] = {} + + @property + def probabilities(self) -> np.ndarray: + return group_probabilities( + [float(child.sampling_sizes.sum()) for child in self.children], + self.dataset.probability_specs, + ) + + @property + def batch_counts(self) -> np.ndarray: + return np.concatenate([child.batch_counts for child in self.children]) + + @property + def sampling_sizes(self) -> np.ndarray: + return np.concatenate([child.sampling_sizes for child in self.children]) + + @property + def system_probabilities(self) -> np.ndarray: + return np.concatenate( + [ + probability * child.system_probabilities + for child, probability in zip( + self.children, self.probabilities, strict=True + ) + ] + ) + + @property + def epoch_length(self) -> int: + return compute_total_numb_batch(self.batch_counts, self.system_probabilities) + + def __len__(self) -> int: + return compute_total_numb_batch( + [len(child) for child in self.children], self.probabilities + ) + + def __iter__(self) -> Iterator[dict[str, Any]]: + probabilities = self.probabilities + cdf = np.cumsum(probabilities) + cdf[-1] = 1.0 + count = compute_total_numb_batch( + [len(child) for child in self.children], probabilities + ) + offsets = np.cumsum([0, *(len(child.batch_counts) for child in self.children)]) + for _ in range(count): + index = int(np.searchsorted(cdf, self._rng.random(), side="right")) + if index not in self._iterators: + self._iterators[index] = iter(self.children[index]) + batch = next(self._iterators[index], None) + if batch is None: + self._iterators[index] = iter(self.children[index]) + batch = next(self._iterators[index]) + result = {**batch, "sid": batch["sid"] + int(offsets[index])} + if self.dataset.filter_enabled is not None: + result = DataBatch(result, self.dataset.filter_enabled[index]) + elif isinstance(batch, DataBatch): + result = DataBatch(result, batch.pair_filter_enabled) + yield result + + def print_summary(self, name: str) -> None: + for index, (child, probability) in enumerate( + zip(self.children, self.probabilities, strict=True) + ): + label = f"{name}/group_{index}" + log.info("%s: batch probability %.6g", label, probability) + child.print_summary(label) + + def close(self) -> None: + for iterator in self._iterators.values(): + close = getattr(iterator, "close", None) + if close is not None: + close() + self._iterators.clear() + for child in self.children: + child.close() diff --git a/deepmd/utils/data_path.py b/deepmd/utils/data_path.py new file mode 100644 index 0000000000..986018fd0c --- /dev/null +++ b/deepmd/utils/data_path.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-independent discovery of on-disk data systems.""" + +from pathlib import ( + Path, +) +from typing import ( + Literal, +) + +from deepmd.utils.path import ( + DPOSPath, + DPPath, +) + +DatasetFormat = Literal["deepmd", "lmdb"] + + +def is_lmdb(systems: str) -> bool: + """Return whether a path identifies an LMDB dataset.""" + return systems.endswith(".lmdb") or Path(systems, "data.mdb").is_file() + + +def _inside_lmdb(path: Path, root: Path) -> bool: + for parent in path.parents: + if parent == root: + break + if is_lmdb(str(parent)): + return True + return False + + +def expand_sys_str( + root_dir: str | Path, + *, + formats: tuple[DatasetFormat, ...] = ("deepmd",), + patterns: list[str] | None = None, +) -> list[str]: + """Discover data systems beneath one directory or HDF5 container. + + Parameters + ---------- + root_dir : str or Path + A data system, a parent directory, or a DeePMD HDF5 path. + formats : tuple of str + Accepted leaf formats. ``deepmd`` recognizes ``type.raw`` in NPY + directories and HDF5 groups; ``lmdb`` recognizes LMDB directories. + patterns : list of str, optional + Recursive glob patterns selecting system paths. A directly supplied + LMDB dataset is terminal and does not require a glob match. + + Returns + ------- + list[str] + Discovered system paths. Unpatterned DeePMD traversal retains its + established order, with the root system last. Overlapping patterns + include each matched path once. + """ + root_name = str(root_dir) + if is_lmdb(root_name): + return [root_name] if "lmdb" in formats else [] + + root = DPPath(root_name) + candidates = ( + [*root.rglob("*"), root] + if patterns is None + else [path for pattern in patterns for path in root.rglob(pattern)] + ) + systems: dict[str, None] = {} + for candidate in candidates: + name = str(candidate) + if isinstance(candidate, DPOSPath): + path = candidate.path + if _inside_lmdb(path, root.path): + continue + if candidate.is_dir() and is_lmdb(name): + if "lmdb" in formats: + systems[name] = None + continue + if "deepmd" in formats and (candidate / "type.raw").is_file(): + systems[name] = None + return list(systems) + + +def rglob_sys_str(root_dir: str, patterns: list[str]) -> list[str]: + """Discover DeePMD systems matching recursive glob patterns.""" + return expand_sys_str(root_dir, patterns=patterns) + + +def process_systems( + systems: str | list[str], + patterns: list[str] | None = None, + *, + formats: tuple[DatasetFormat, ...] = ("deepmd", "lmdb"), +) -> list[str]: + """Expand every input path through the same format-aware discovery. + + Explicit input order and repeated input paths are preserved, so an + input list defines the order used by per-system sampling probabilities. + Each entry may name a dataset directly or a parent containing datasets. + """ + paths = [systems] if isinstance(systems, str) else systems + if not isinstance(paths, list) or any(not isinstance(path, str) for path in paths): + raise ValueError( + f"Invalid systems value {systems!r}: expected str or list[str]." + ) + return [ + system + for path in paths + for system in expand_sys_str(path, formats=formats, patterns=patterns) + ] diff --git a/deepmd/utils/data_sampling.py b/deepmd/utils/data_sampling.py new file mode 100644 index 0000000000..b0757417ef --- /dev/null +++ b/deepmd/utils/data_sampling.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Sampling probabilities shared by data formats and training backends.""" + +from collections.abc import ( + Sequence, +) + +import numpy as np + +from deepmd.utils import random as dp_random + + +def process_sys_probs( + sys_probs: Sequence[float], nbatch: Sequence[int] | np.ndarray +) -> np.ndarray: + """Fill unspecified probabilities in proportion to the batch counts. + + Non-negative entries specify absolute probabilities. Negative entries + share the remaining probability in proportion to their batch counts. + A fully assigned probability vector does not depend on batch counts. + """ + probabilities = np.asarray(sys_probs, dtype=np.float64) + if probabilities.ndim != 1: + raise ValueError("Sampling probabilities must be one-dimensional.") + if not np.all(np.isfinite(probabilities)): + raise ValueError("Sampling probabilities must be finite.") + assigned = probabilities >= 0 + assigned_sum = float(probabilities[assigned].sum()) + if assigned_sum > 1.0 + 1e-8: + raise ValueError("The sum of assigned probabilities must not exceed 1.") + result = np.where(assigned, probabilities, 0.0) + remaining = max(0.0, 1.0 - assigned_sum) + if remaining > 1e-8: + sizes = np.asarray(nbatch, dtype=np.float64) + if probabilities.shape != sizes.shape: + raise ValueError( + "Sampling probabilities and batch counts must have the same length." + ) + if not np.all(np.isfinite(sizes)) or np.any(sizes < 0): + raise ValueError("Batch counts must be finite and non-negative.") + remaining_size = float(sizes[~assigned].sum()) + if remaining_size <= 0: + raise ValueError( + "The remaining probability requires at least one unassigned, " + "nonempty dataset." + ) + result[~assigned] = remaining * sizes[~assigned] / remaining_size + if result.size == 0 or result.sum() <= 0: + raise ValueError("Sampling probabilities must contain a positive entry.") + return result / result.sum() + + +def prob_sys_size_ext( + keywords: str, nsystems: int, nbatch: Sequence[int] | np.ndarray +) -> np.ndarray: + """Distribute each configured block's weight by its system sizes.""" + blocks = [part.split(":") for part in keywords.split(";")[1:]] + weights = np.asarray([float(block[2]) for block in blocks]) + if np.any(weights < 0): + raise ValueError("The weight of a block must be non-negative.") + block_probs = weights / weights.sum() + sys_probs = np.zeros(nsystems, dtype=np.float64) + sizes = np.asarray(nbatch, dtype=np.float64) + for block, probability in zip(blocks, block_probs, strict=True): + start, end = int(block[0]), int(block[1]) + block_sizes = sizes[start:end] + sys_probs[start:end] = block_sizes / block_sizes.sum() * probability + return sys_probs + + +def system_probabilities( + nbatches: Sequence[int] | np.ndarray, + *, + sys_probs: Sequence[float] | None = None, + auto_prob: str = "prob_sys_size", +) -> np.ndarray: + """Resolve the established per-system sampling rules.""" + nsystems = len(nbatches) + if nsystems == 0: + raise ValueError("No systems provided.") + if sys_probs is not None: + return process_sys_probs(sys_probs, nbatches) + if auto_prob == "prob_uniform": + return np.full(nsystems, 1.0 / nsystems, dtype=np.float64) + if auto_prob == "prob_sys_size": + auto_prob = f"prob_sys_size;0:{nsystems}:1.0" + if not auto_prob.startswith("prob_sys_size;"): + raise RuntimeError(f"Unknown auto prob style: {auto_prob}") + return prob_sys_size_ext(auto_prob, nsystems, nbatches) + + +def group_probabilities( + sizes: Sequence[int] | np.ndarray, probabilities: Sequence[float | None] +) -> np.ndarray: + """Assign group probabilities, sharing any remainder by group size.""" + if len(sizes) != len(probabilities): + raise ValueError("Group sizes and probabilities must have the same length.") + if any( + probability is not None and probability < 0 for probability in probabilities + ): + raise ValueError("A group's prob must be non-negative or omitted.") + resolved = process_sys_probs( + [-1.0 if probability is None else probability for probability in probabilities], + sizes, + ) + if np.any((np.asarray(sizes) == 0) & (resolved > 0)): + raise ValueError("A data group with positive probability must contain batches.") + return resolved + + +def data_child_seeds(seed: int | None, count: int) -> list[int]: + """Derive independent selection streams for nodes in a dataset hierarchy.""" + if seed is None: + seed = int(dp_random.choice(2**32)) + return [ + int(child.generate_state(1, dtype=np.uint64)[0]) + for child in np.random.SeedSequence(seed).spawn(count) + ] + + +def global_batch_counts( + frame_counts: Sequence[int], batch_sizes: Sequence[int] +) -> list[int]: + """Count source batches before distributed sampling and rank-local padding.""" + return [ + (frames + batch_size - 1) // batch_size + for frames, batch_size in zip(frame_counts, batch_sizes, strict=True) + ] diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 890224e063..ae124859ad 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -2,9 +2,15 @@ import collections import logging import warnings +from collections.abc import ( + Callable, +) from functools import ( cached_property, ) +from types import ( + SimpleNamespace, +) from typing import ( Any, ) @@ -13,9 +19,7 @@ import deepmd.utils.random as dp_random from deepmd.common import ( - expand_sys_str, make_default_mesh, - rglob_sys_str, ) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, @@ -24,6 +28,24 @@ DataRequirementItem, DeepmdData, ) +from deepmd.utils.data_config import ( + DataConfig, + check_type_map_consistency, + resolve_data_groups, + resolve_data_sources, +) +from deepmd.utils.data_path import ( + is_lmdb, +) +from deepmd.utils.data_path import process_systems as process_systems +from deepmd.utils.data_sampling import ( + data_child_seeds, +) +from deepmd.utils.data_sampling import prob_sys_size_ext as prob_sys_size_ext +from deepmd.utils.data_sampling import process_sys_probs as process_sys_probs +from deepmd.utils.data_sampling import ( + system_probabilities, +) from deepmd.utils.out_stat import ( compute_stats_from_redu, ) @@ -230,7 +252,7 @@ def __init__( self.data_systems[ii].get_sys_numb_batch(self.batch_size[ii]) ) type_map_list.append(self.data_systems[ii].get_type_map()) - self.type_map = self._check_type_map_consistency(type_map_list) + self.type_map = check_type_map_consistency(type_map_list) # ! altered by Marián Rynik # test size @@ -434,23 +456,9 @@ def set_sys_probs( sys_probs: list[float] | None = None, auto_prob_style: str = "prob_sys_size", ) -> None: - if sys_probs is None: - if auto_prob_style == "prob_uniform": - prob_v = 1.0 / float(self.nsystems) - probs = [prob_v for ii in range(self.nsystems)] - elif auto_prob_style[:13] == "prob_sys_size": - if auto_prob_style == "prob_sys_size": - prob_style = f"prob_sys_size;0:{self.get_nsystems()}:1.0" - else: - prob_style = auto_prob_style - probs = prob_sys_size_ext( - prob_style, self.get_nsystems(), self.nbatches - ) - else: - raise RuntimeError("Unknown auto prob style: " + auto_prob_style) - else: - probs = process_sys_probs(sys_probs, self.nbatches) - self.sys_probs = probs + self.sys_probs = system_probabilities( + self.nbatches, sys_probs=sys_probs, auto_prob=auto_prob_style + ) def get_batch(self, sys_idx: int | None = None) -> dict: # batch generation style altered by Ziyao Li: @@ -668,6 +676,10 @@ def get_batch_pass_length(self) -> int: return compute_total_numb_batch(self.nbatches, self.sys_probs) + def get_epoch_batch_counts(self, world_size: int) -> np.ndarray: + """Return optimizer workloads for the replicated NumPy data reader.""" + return np.asarray(self.nbatches, dtype=np.float64) / world_size + def print_summary(self, name: str) -> None: print_summary( name, @@ -700,24 +712,6 @@ def _make_auto_ts(self, percent: float) -> list[int]: return ts - def _check_type_map_consistency( - self, type_map_list: list[list[str] | None] - ) -> list[str]: - ret = [] - for ii in type_map_list: - if ii is not None: - min_len = min([len(ii), len(ret)]) - for idx in range(min_len): - if ii[idx] != ret[idx]: - raise RuntimeError( - f"Inconsistent type map: {ret!s} {ii!s} in different data systems. " - "If you didn't set model/type_map, please set it, " - "since the type map of the model cannot be decided by data." - ) - if len(ii) > len(ret): - ret = ii - return ret - def _format_name_length(name: str, width: int) -> str: if len(name) <= width: @@ -820,115 +814,22 @@ def print_summary( ) -def process_sys_probs(sys_probs: list[float], nbatch: int) -> np.ndarray: - sys_probs = np.array(sys_probs) - type_filter = sys_probs >= 0 - assigned_sum_prob = np.sum(type_filter * sys_probs) - # 1e-8 is to handle floating point error; See #1917 - assert assigned_sum_prob <= 1.0 + 1e-8, ( - "the sum of assigned probability should be less than 1" - ) - rest_sum_prob = 1.0 - assigned_sum_prob - if not np.isclose(rest_sum_prob, 0): - rest_nbatch = (1 - type_filter) * nbatch - rest_prob = rest_sum_prob * rest_nbatch / np.sum(rest_nbatch) - ret_prob = rest_prob + type_filter * sys_probs - else: - ret_prob = sys_probs - assert np.isclose(np.sum(ret_prob), 1), "sum of probs should be 1" - return ret_prob - - -def prob_sys_size_ext(keywords: str, nsystems: int, nbatch: int) -> list[float]: - block_str = keywords.split(";")[1:] - block_stt = [] - block_end = [] - block_weights = [] - for ii in block_str: - stt = int(ii.split(":")[0]) - end = int(ii.split(":")[1]) - weight = float(ii.split(":")[2]) - assert weight >= 0, "the weight of a block should be no less than 0" - block_stt.append(stt) - block_end.append(end) - block_weights.append(weight) - nblocks = len(block_str) - block_probs = np.array(block_weights) / np.sum(block_weights) - sys_probs = np.zeros([nsystems], dtype=np.float64) - for ii in range(nblocks): - nbatch_block = nbatch[block_stt[ii] : block_end[ii]] - tmp_prob = [float(i) for i in nbatch_block] / np.sum(nbatch_block) - sys_probs[block_stt[ii] : block_end[ii]] = tmp_prob * block_probs[ii] - return sys_probs - - -def process_systems( - systems: str | list[str], patterns: list[str] | None = None -) -> list[str]: - """Process the user-input systems. - - If it is a single directory, search for all the systems in the directory. - If it is a list, each item in the list is treated as a directory to search. - If it is a single LMDB path, return it directly without expansion. - Check if the systems are valid. - - Parameters - ---------- - systems : str or list of str - The user-input systems - patterns : list of str, optional - The patterns to match the systems, by default None - - Returns - ------- - result_systems: list of str - The valid systems - """ - from deepmd.dpmodel.utils.lmdb_data import ( - is_lmdb, - ) - - # LMDB path: return directly without expansion - if isinstance(systems, str) and is_lmdb(systems): - return [systems] - - # Normalize input to a list of paths to search - if isinstance(systems, str): - search_paths = [systems] - elif isinstance(systems, list): - search_paths = systems - else: - # Handle unsupported input types - raise ValueError( - f"Invalid systems type: {type(systems)}. Must be str or list[str]." - ) - - # Iterate over the search_paths list and apply expansion logic to each path - result_systems = [] - for path in search_paths: - if patterns is None: - expanded_paths = expand_sys_str(path) - else: - expanded_paths = rglob_sys_str(path, patterns) - - result_systems.extend(expanded_paths) - - return result_systems - - def get_data( - jdata: dict[str, Any], + jdata: DataConfig, rcut: float, type_map: list[str] | None, modifier: Any | None, multi_task_mode: bool = False, -) -> DeepmdDataSystem: - """Get the data system. + *, + lmdb_factory: Callable[[str, int | str, str | None], Any] | None = None, + seed: int | None = None, +) -> Any: + """Build native data readers and compose independently configured groups. Parameters ---------- jdata - The json data + A data dictionary or a nonempty list of group dictionaries. rcut The cut-off radius, not used type_map @@ -937,33 +838,100 @@ def get_data( The data modifier multi_task_mode If in multi task mode + lmdb_factory + Optional backend adapter accepting an LMDB path, batch size, and + internal auto-probability rule. Without it, only NPY/HDF5 sources + are discovered. + seed + Seed for group selection. If omitted, draws from the shared training + random generator seed the group-selection streams. Returns ------- - DeepmdDataSystem - The data system + object + A native data system for a single source, or a grouped data system + with the same batch, statistics, and sampling-metadata interface. """ - systems = jdata["systems"] - rglob_patterns = jdata.get("rglob_patterns") - systems = process_systems(systems, patterns=rglob_patterns) - - batch_size = jdata["batch_size"] - sys_probs = jdata.get("sys_probs") - auto_prob = jdata.get("auto_prob", "prob_sys_size") - optional_type_map = not multi_task_mode - - data = DeepmdDataSystem( - systems=systems, - batch_size=batch_size, - test_size=1, # to satisfy the old api - shuffle_test=True, # to satisfy the old api - rcut=rcut, - type_map=type_map, - optional_type_map=optional_type_map, - modifier=modifier, - trn_all_set=True, # sample from all sets - sys_probs=sys_probs, - auto_prob_style=auto_prob, + from deepmd.utils.data_group import ( + GroupedDataSystem, + ) + + groups = resolve_data_groups( + jdata, formats=("deepmd", "lmdb") if lmdb_factory is not None else ("deepmd",) + ) + needs_composition = len(groups) > 1 or any( + len(group.systems) > 1 and any(is_lmdb(path) for path in group.systems) + for group in groups + ) + seeds = ( + iter(data_child_seeds(seed, len(groups) + 1)) if needs_composition else iter(()) + ) + datasets = [] + for group in groups: + sources = resolve_data_sources(group) + has_lmdb = any(source.format == "lmdb" for source in sources) + children = [] + for source in sources: + if source.format == "deepmd": + child = DeepmdDataSystem( + systems=list(source.systems), + batch_size=source.batch_size, + test_size=1, + shuffle_test=True, + rcut=rcut, + type_map=type_map, + optional_type_map=not multi_task_mode, + modifier=modifier, + trn_all_set=True, + sys_probs=group.sys_probs if not has_lmdb else None, + auto_prob_style=group.auto_prob or "prob_sys_size", + ) + else: + assert lmdb_factory is not None + child = lmdb_factory( + source.systems[0], source.batch_size, group.auto_prob + ) + children.append(child) + data = ( + children[0] + if len(children) == 1 + else GroupedDataSystem(children, [None] * len(children), seed=next(seeds)) + ) + datasets.append(data) + return ( + datasets[0] + if len(datasets) == 1 + else GroupedDataSystem( + datasets, [group.prob for group in groups], seed=next(seeds) + ) ) - return data + +def get_neighbor_stat_data(jdata: DataConfig, type_map: list[str] | None) -> Any: + """Expose every discovered source to backend neighbor-statistic kernels.""" + sources = [] + for group in resolve_data_groups(jdata): + for source in resolve_data_sources(group): + if source.format == "deepmd": + sources.append( + DeepmdDataSystem( + list(source.systems), + source.batch_size, + 1, + type_map=type_map, + trn_all_set=True, + ) + ) + else: + from deepmd.dpmodel.utils.lmdb_data import ( + make_neighbor_stat_data, + ) + + sources.append(make_neighbor_stat_data(source.systems[0], type_map)) + return SimpleNamespace( + data_systems=[system for source in sources for system in source.data_systems], + system_dirs=[path for source in sources for path in source.system_dirs], + mixed_type=any(source.mixed_type for source in sources), + get_ntypes=lambda: max(source.get_ntypes() for source in sources), + get_batch=lambda: None, + ) diff --git a/deepmd/utils/json_schema.py b/deepmd/utils/json_schema.py index 1a080ec777..6430bbff8a 100644 --- a/deepmd/utils/json_schema.py +++ b/deepmd/utils/json_schema.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""JSON Schema for model inputs before preset expansion.""" +"""JSON Schema for model presets and grouped data input forms.""" import json import re @@ -14,6 +14,24 @@ _REGIONS = ("type_map", "descriptor", "fitting_net") +def with_data_groups(schema: dict[str, Any], multi_task: bool) -> dict[str, Any]: + """Describe the same training-data fields for dictionaries and list items.""" + training = schema["properties"]["training"] + if multi_task: + training = training["properties"]["data_dict"]["additionalProperties"] + data = training["properties"]["training_data"] + data["items"] = { + "type": "object", + "properties": data["properties"], + "required": data["required"], + } + if "allOf" in data: + data["items"]["allOf"] = data["allOf"] + data["minItems"] = 1 + data["properties"]["prob"].update(minimum=0, maximum=1) + return schema + + def _preset_pattern(names: list[str]) -> str: """Match complete preset names case-insensitively with ECMA-262 syntax.""" alternatives = [ diff --git a/deepmd/utils/model_stat.py b/deepmd/utils/model_stat.py index 109de19fd5..ac927d0e4f 100644 --- a/deepmd/utils/model_stat.py +++ b/deepmd/utils/model_stat.py @@ -99,7 +99,8 @@ def _collect_stat_batches( if "natoms_vec" in stat_data: stat_data["natoms_vec"] = stat_data["natoms_vec"].astype(np.int32) batch = normalize_batch(stat_data) - frame_mask = pair_margin_frame_mask(batch, min_pair_dist > 0.0) + filter_enabled = getattr(stat_data, "pair_filter_enabled", min_pair_dist > 0.0) + frame_mask = pair_margin_frame_mask(batch, filter_enabled) if frame_mask is not None and not frame_mask.all(): if not frame_mask.any(): continue diff --git a/deepmd/utils/path.py b/deepmd/utils/path.py index 44bbe34edf..c511ea52fb 100644 --- a/deepmd/utils/path.py +++ b/deepmd/utils/path.py @@ -22,6 +22,7 @@ Self, ) from wcmatch.glob import ( + GLOBSTAR, globfilter, ) @@ -393,7 +394,7 @@ def glob(self, pattern: str) -> list["DPPath"]: ] return [ type(self)(f"{self.root_path}#{pp}", mode=self.mode) - for pp in globfilter(subpaths, self._connect_path(pattern)) + for pp in globfilter(subpaths, self._connect_path(pattern), flags=GLOBSTAR) ] def rglob(self, pattern: str) -> list["DPPath"]: @@ -410,7 +411,7 @@ def rglob(self, pattern: str) -> list["DPPath"]: list[DPPath] list of paths """ - return self.glob("**" + pattern) + return self.glob("**/" + pattern) @property def _keys(self) -> list[str]: diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 31609c09b8..824d92d887 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -53,11 +53,12 @@ Other training parameters are given in the {ref}`training ` section. The sections {ref}`training_data ` and {ref}`validation_data ` give the training dataset and validation dataset, respectively. Taking the training dataset for example, the keys are explained below: - {ref}`systems ` provide paths of the training data systems. DeePMD-kit allows you to provide multiple systems with different numbers of atoms. This key can be a `list` or a `str`. - - `str`: {ref}`systems ` should be a valid path. It can be a system directory path (containing 'type.raw') or a parent directory path to recursively search for all system subdirectories. + - `str`: {ref}`systems ` should be a valid path. It can be a system directory path (containing 'type.raw'), a DeePMD HDF5 container or group, or a parent directory path to recursively search for all system subdirectories. - `list`: {ref}`systems ` gives a list of paths. Each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories. + - The PyTorch and PyTorch Exportable backends also accept LMDB paths, lists of LMDB paths, and parent directories containing LMDB datasets. An LMDB dataset is terminal: discovery does not inspect its contents as additional datasets. NPY and LMDB sources may appear in the same group. Other backends discover DeePMD NPY/HDF5 systems only. - At each training step, DeePMD-kit randomly picks {ref}`batch_size ` frame(s) from one of the systems. The probability of using a system is by default in proportion to the number of batches in the system. More options are available for automatically determining the probability of using systems. One can set the key {ref}`auto_prob ` to - `"prob_uniform"` all systems are used with the same probability. - - `"prob_sys_size"` the probability of using a system is proportional to its size (number of frames). + - `"prob_sys_size"` the probability of using a system is proportional to its number of batches. - `"prob_sys_size; sidx_0:eidx_0:w_0; sidx_1:eidx_1:w_1;..."` the `list` of systems is divided into blocks. Block `i` has systems ranging from `sidx_i` to `eidx_i`. The probability of using a system from block `i` is proportional to `w_i`. Within one block, the probability of using a system is proportional to its size. - An example of using `"auto_prob"` is given below. The probability of using `systems[2]` is 0.4, and the sum of the probabilities of using `systems[0]` and `systems[1]` is 0.6. If the number of frames in `systems[1]` is twice of `system[0]`, then the probability of using `system[1]` is 0.4 and that of `system[0]` is 0.2. @@ -103,6 +104,53 @@ The sections {ref}`training_data ` and {ref}`validation_ Filling batches also makes the *sampling weight* more faithful. Per-atom loss terms (force, atomic energy) pool over a batch's real labels, so a frame already counts in proportion to its atom count there. Frame-level terms (energy, virial) weigh the frames of a batch equally, so a frame counts for `1 / nframes`, and an atom budget makes `nframes` follow the atom count. `"max:N"` distorts that wherever an atom-count group is too small to fill a batch — its few frames form a short batch and each is over-weighted many times over — while `"mix:N"` packs them with their neighbours and lands much closer to the intended weighting. - The key {ref}`numb_batch ` in {ref}`validate_data ` gives the number of batches of model validation. Note that the batches may not be from the same system +### Data groups + +`training_data` accepts either the dictionary shown above or a nonempty list of dictionaries. Each dictionary defines a group with its own `systems`, `batch_size`, and within-group sampling options. The optional `prob` specifies the probability of drawing a training batch from the entire group, including every dataset found recursively under its paths: + +```json +"training_data": [ + { + "systems": "../data/collection_a", + "batch_size": "auto:32", + "prob": 0.6 + }, + { + "systems": ["../data/collection_b", "../data/extra"], + "batch_size": "auto:64", + "prob": 0.4 + } +] +``` + +The first group's datasets together receive probability 0.6, regardless of how many datasets its path contains. The second group's datasets together receive probability 0.4. Batches do not combine frames from different groups. + +`prob` is an absolute probability between 0 and 1. If every group specifies it, the probabilities must sum to 1. Groups without `prob` share the remaining probability in proportion to their batch counts. Omitting `prob` everywhere therefore preserves size-based sampling; it does not give equally sized shares to differently sized groups. Within each group, `auto_prob` and `sys_probs` retain their existing meanings. + +For a group containing exactly one LMDB, `auto_prob` addresses that LMDB's internal systems, independently of the group's `prob`: + +```json +"training_data": [ + { + "systems": "../data/first.lmdb", + "batch_size": "filter:450", + "auto_prob": "prob_sys_size;0:10:0.5;10:30:0.5", + "prob": 0.6 + }, + { + "systems": "../data/second.lmdb", + "batch_size": "filter:450", + "prob": 0.4 + } +] +``` + +If a group expands to multiple sources and includes LMDB, it uses size-based sampling: `auto_prob` must remain `"prob_sys_size"`, its default, and `sys_probs` must be omitted. Custom block rules and `"prob_uniform"` are rejected because internal system indices belong to individual LMDB datasets. To assign source weights or configure different internal rules, put each LMDB in a separate group. Existing dictionary configurations, NPY `sys_probs`, and single-LMDB `auto_prob` rules remain supported. + +`validation_data` remains a dictionary. Its `systems` accepts the same recursive paths and multiple-LMDB inputs on supported backends; `numb_btch` continues to apply to the whole validation dataset. + +### Mixed precision + The section {ref}`mixed_precision ` specifies the mixed precision settings, which will enable the mixed precision training workflow for DeePMD-kit. The keys are explained below: - {ref}`output_prec ` precision used in the output tensors, only `float32` is supported currently. diff --git a/source/tests/common/test_argcheck_training.py b/source/tests/common/test_argcheck_training.py index cb5d2ffafa..3e783eefda 100644 --- a/source/tests/common/test_argcheck_training.py +++ b/source/tests/common/test_argcheck_training.py @@ -1,7 +1,18 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Tests for normalized training arguments.""" +import copy +from typing import ( + Any, +) + +import pytest +from dargs.dargs import ( + ArgumentError, +) + from deepmd.utils.argcheck import ( + normalize, training_args, ) @@ -15,3 +26,62 @@ def test_ema_checkpoint_retention_is_left_for_runtime_inheritance() -> None: training_argument.check_value(normalized, strict=True) assert normalized["ema_ckpt_keep"] is None + + +def _training_input(data: dict[str, Any] | list[dict[str, Any]]) -> dict[str, Any]: + return { + "model": { + "type_map": ["H"], + "descriptor": {"type": "se_e2_a", "sel": [1]}, + "fitting_net": {}, + }, + "training": {"numb_steps": 1, "training_data": data}, + } + + +def test_training_groups_normalize_aliases_without_mutating_input() -> None: + source = _training_input( + [ + { + "systems": "first.lmdb", + "prob": 0.6, + "auto_prob_style": "prob_sys_size;0:1:1", + }, + {"systems": ["second.lmdb", "third.lmdb"], "prob": 0.4}, + ] + ) + original = copy.deepcopy(source) + + groups = normalize(source)["training"]["training_data"] + + assert source == original + assert groups[0]["auto_prob"] == "prob_sys_size;0:1:1" + assert "auto_prob_style" not in groups[0] + assert groups[1]["auto_prob"] == "prob_sys_size" + assert groups[0]["batch_size"] == groups[1]["batch_size"] == "auto" + assert [group["prob"] for group in groups] == [0.6, 0.4] + + +def test_legacy_training_dictionary_remains_a_dictionary() -> None: + group = normalize(_training_input({"systems": "dataset"}))["training"][ + "training_data" + ] + + assert isinstance(group, dict) + assert group["auto_prob"] == "prob_sys_size" + assert group["batch_size"] == "auto" + + +@pytest.mark.parametrize( + "groups", + [ + [], + [{"systems": "dataset", "unknown_option": True}], + [{"systems": "dataset", "prob": -0.1}], + [{"systems": "dataset", "prob": 0.6}], + [{"systems": "dataset", "prob": float("nan")}], + ], +) +def test_invalid_training_groups_are_rejected(groups: list[dict[str, Any]]) -> None: + with pytest.raises((ArgumentError, ValueError)): + normalize(_training_input(groups)) diff --git a/source/tests/common/test_data_config.py b/source/tests/common/test_data_config.py new file mode 100644 index 0000000000..096caeb0d6 --- /dev/null +++ b/source/tests/common/test_data_config.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Shared dataset configuration and sampling contracts.""" + +from pathlib import ( + Path, +) + +import numpy as np +import pytest + +from deepmd.utils import random as dp_random +from deepmd.utils.data_config import ( + resolve_data_groups, +) +from deepmd.utils.data_sampling import ( + group_probabilities, + system_probabilities, +) +from deepmd.utils.data_system import ( + get_data, +) + + +def test_group_probabilities_do_not_equalize_unweighted_groups() -> None: + np.testing.assert_allclose(group_probabilities([2, 8], [None, None]), [0.2, 0.8]) + np.testing.assert_allclose(group_probabilities([2, 8], [0.6, 0.4]), [0.6, 0.4]) + np.testing.assert_allclose( + group_probabilities([2, 3, 9], [0.6, None, None]), [0.6, 0.1, 0.3] + ) + np.testing.assert_allclose(group_probabilities([2, 8], [1.0, None]), [1.0, 0.0]) + + +def test_group_and_system_probabilities_are_independent() -> None: + outer = group_probabilities([10, 90], [0.6, 0.4]) + inner = system_probabilities([1, 2, 7], auto_prob="prob_sys_size;0:1:0.5;1:3:0.5") + + np.testing.assert_allclose(outer[0] * inner, [0.3, 0.6 / 9, 2.1 / 9]) + assert np.isclose((outer[0] * inner).sum(), 0.6) + + +@pytest.mark.parametrize("systems_as_list", [False, True]) +def test_multi_lmdb_requires_default_auto_prob( + tmp_path: Path, systems_as_list: bool +) -> None: + first = tmp_path / "first.lmdb" + second = tmp_path / "nested" / "second.lmdb" + first.mkdir() + second.mkdir(parents=True) + systems = [str(first), str(second)] if systems_as_list else str(tmp_path) + + groups = resolve_data_groups({"systems": systems}) + assert set(groups[0].systems) == {str(first), str(second)} + assert groups[0].auto_prob is None + explicit_default = resolve_data_groups( + {"systems": systems, "auto_prob": "prob_sys_size"} + ) + assert explicit_default[0].systems == groups[0].systems + for custom in ("prob_uniform", "prob_sys_size;0:1:1"): + with pytest.raises(ValueError, match="auto_prob must be prob_sys_size"): + resolve_data_groups({"systems": systems, "auto_prob": custom}) + with pytest.raises(ValueError, match="sys_probs is not supported"): + resolve_data_groups({"systems": systems, "sys_probs": [-1, -1]}) + + +def test_empty_groups_cannot_receive_positive_probability() -> None: + with pytest.raises(ValueError, match="must contain batches"): + group_probabilities([0, 8], [0.5, 0.5]) + np.testing.assert_allclose(group_probabilities([0, 8], [None, None]), [0, 1]) + + +def test_single_lmdb_groups_retain_internal_auto_prob() -> None: + groups = resolve_data_groups( + [ + { + "systems": "first.lmdb", + "auto_prob": "prob_sys_size;0:10:0.5;10:30:0.5", + "prob": 0.6, + }, + {"systems": "second.lmdb", "prob": 0.4}, + ] + ) + + assert groups[0].auto_prob == "prob_sys_size;0:10:0.5;10:30:0.5" + assert groups[1].auto_prob is None + assert [group.prob for group in groups] == [0.6, 0.4] + + +def test_numpy_groups_follow_the_training_seed() -> None: + system = str(Path(__file__).resolve().parents[1] / "pt" / "property" / "double") + config = [{"systems": system, "batch_size": size} for size in (1, 2)] + sequences = [] + for _ in range(2): + dp_random.seed(123) + data = get_data(config, 0, None, None) + sequences.append([data.get_batch()["coord"].shape[0] for _ in range(24)]) + assert sequences[0] == sequences[1] + assert set(sequences[0]) == {1, 2} + + +def test_numpy_backend_parent_discovery_preserves_format_scope(tmp_path: Path) -> None: + numpy = tmp_path / "numpy" + numpy.mkdir() + (numpy / "type.raw").touch() + lmdb = tmp_path / "frames.lmdb" + lmdb.mkdir() + groups = resolve_data_groups({"systems": str(tmp_path)}, formats=("deepmd",)) + assert groups[0].systems == (str(numpy),) diff --git a/source/tests/common/test_json_schema.py b/source/tests/common/test_json_schema.py index 843af94da8..eecc27866f 100644 --- a/source/tests/common/test_json_schema.py +++ b/source/tests/common/test_json_schema.py @@ -49,6 +49,27 @@ def _assert_valid(validator: Draft202012Validator, data: dict[str, Any]) -> None assert not errors, [(list(error.path), error.message) for error in errors] +@pytest.mark.parametrize("multi_task", [False, True]) +def test_training_group_schema( + validators: dict[bool, Draft202012Validator], multi_task: bool +) -> None: + model = {"preset": "dpa4-nano-v20260911"} + if multi_task: + model = {"model_dict": {"a": model}} + data = _input(model, multi_task) + task = data["training"]["data_dict"]["a"] if multi_task else data["training"] + task["training_data"] = [ + {"systems": "first.lmdb", "prob": 0.6}, + {"systems": ["second.lmdb", "third.lmdb"], "prob": 0.4}, + ] + _assert_valid(validators[multi_task], data) + + task["training_data"][1]["batch_size"] = False + assert not validators[multi_task].is_valid(data) + task["training_data"] = [] + assert not validators[multi_task].is_valid(data) + + def test_preset_user_example(validators: dict[bool, Draft202012Validator]) -> None: data = _input( { diff --git a/source/tests/common/test_path.py b/source/tests/common/test_path.py index ec3f175c0a..7573c69a5b 100644 --- a/source/tests/common/test_path.py +++ b/source/tests/common/test_path.py @@ -8,6 +8,9 @@ import h5py import numpy as np +from deepmd.utils.data_system import ( + process_systems, +) from deepmd.utils.path import ( DPH5Path, DPPath, @@ -74,3 +77,56 @@ def test_write_operations_are_rejected(self) -> None: (self.path / "value").save_numpy(np.ones(1)) with self.assertRaisesRegex(ValueError, "read-only"): (self.path / "group").mkdir() + + +class TestSystemDiscovery(unittest.TestCase): + def test_lmdb_paths_and_parent_directories(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + direct = root / "direct.lmdb" + direct.mkdir() + parent = root / "collection" + nested = parent / "nested" / "frames" + nested.mkdir(parents=True) + (nested / "data.mdb").touch() + npy = parent / "npy" + npy.mkdir() + (npy / "type.raw").touch() + + found = process_systems([str(direct), str(parent)]) + + self.assertEqual(found[0], str(direct)) + self.assertCountEqual(found, [str(direct), str(nested), str(npy)]) + + def test_lmdb_pattern_expansion(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + first = root / "train.lmdb" + first.mkdir() + second = root / "test.lmdb" + second.mkdir() + + self.assertEqual( + process_systems(str(root), patterns=["train*.lmdb"]), + [str(first)], + ) + + def test_lmdb_is_terminal_even_when_patterns_exclude_it(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + internal = root / "frames.lmdb" / "nested" + internal.mkdir(parents=True) + (internal / "type.raw").touch() + self.assertEqual(process_systems(str(root), patterns=["nested"]), []) + + def test_hdf5_paths_expand_inside_containers(self) -> None: + with tempfile.TemporaryDirectory() as directory: + container = Path(directory) / "systems.h5" + with h5py.File(container, "w") as handle: + handle.create_dataset("first/type.raw", data=[0]) + handle.create_dataset("nested/second/type.raw", data=[0, 1]) + expected = [f"{container}#/first", f"{container}#/nested/second"] + self.assertEqual(process_systems(str(container)), expected) + self.assertEqual( + process_systems([expected[1], expected[0]]), expected[::-1] + ) diff --git a/source/tests/pt/test_lmdb_dataloader.py b/source/tests/pt/test_lmdb_dataloader.py index ed27015a92..0ba4f61c7d 100644 --- a/source/tests/pt/test_lmdb_dataloader.py +++ b/source/tests/pt/test_lmdb_dataloader.py @@ -6,6 +6,13 @@ Consistency tests (dpmodel vs pt) live in source/tests/consistent/test_lmdb_data.py. """ +from pathlib import ( + Path, +) +from typing import ( + Any, +) + import lmdb import msgpack import numpy as np @@ -30,6 +37,10 @@ from deepmd.pt.loss.ener import ( EnergyStdLoss, ) +from deepmd.pt.utils.dataloader import ( + create_dataloader, + create_dataset, +) from deepmd.pt.utils.lmdb_dataset import ( LmdbBatchDataLoader, LmdbDataset, @@ -38,9 +49,18 @@ from deepmd.pt.utils.stat import ( make_stat_input, ) +from deepmd.pt_expt.utils.lmdb_dataset import ( + LmdbDataSystem, +) from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_group import ( + iter_data_sources, +) +from deepmd.utils.data_system import ( + get_data, +) def _make_frame(natoms: int = 6, seed: int = 0) -> dict: @@ -1491,7 +1511,13 @@ class TestMultitaskLmdbTraining: All assertions are in a single test to avoid creating multiple trainers. """ - def test_multitask_lmdb_end_to_end(self, multitask_lmdb_setup, monkeypatch): + @pytest.mark.parametrize("grouped", [False, True]) + def test_multitask_lmdb_end_to_end( + self, + multitask_lmdb_setup: tuple[dict[str, Any], Path], + monkeypatch: pytest.MonkeyPatch, + grouped: bool, + ) -> None: from copy import ( deepcopy, ) @@ -1510,6 +1536,20 @@ def test_multitask_lmdb_end_to_end(self, multitask_lmdb_setup, monkeypatch): ) config, tmp_path = multitask_lmdb_setup + if grouped: + branches = config["training"]["data_dict"] + branches["model_1"]["training_data"] = [ + { + **branches["model_1"]["training_data"], + "prob": 0.6, + "min_pair_dist": 0.1, + }, + {**branches["model_2"]["training_data"], "prob": 0.4}, + ] + branches["model_1"]["validation_data"]["systems"] = [ + branches["model_1"]["validation_data"]["systems"], + branches["model_2"]["validation_data"]["systems"], + ] monkeypatch.chdir(tmp_path) monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "2") config = update_deepmd_input(deepcopy(config), warning=True) @@ -1536,21 +1576,112 @@ def test_multitask_lmdb_end_to_end(self, multitask_lmdb_setup, monkeypatch): ) assert "coord" in input_dict assert "sid" in log_dict - assert ( - trainer.training_dataloader["model_1"]._batch_iterator._pool - is trainer.training_dataloader["model_2"]._batch_iterator._pool - ) + native_loaders = [ + source.loader + for loader in trainer.training_dataloader.values() + for source in iter_data_sources(loader) + ] + for loader in native_loaders: + next(iter(loader)) + assert len({id(loader._batch_iterator._pool) for loader in native_loaders}) == 1 # -- training run assertions -- trainer.run() assert len(list(tmp_path.glob("model.ckpt*.pt"))) > 0 - assert trainer.training_dataloader["model_1"]._batch_iterator.closed - assert trainer.training_dataloader["model_2"]._batch_iterator.closed - assert trainer.training_dataloader["model_1"].dataset._reader.closed - assert trainer.training_dataloader["model_2"].dataset._reader.closed + assert all(loader._batch_iterator.closed for loader in native_loaders) + assert all(loader.dataset._reader.closed for loader in native_loaders) # Explicit cleanup to free memory on CI import gc del trainer gc.collect() + + +def test_multiple_lmdb_loader_uses_native_batch_sizes( + lmdb_dir: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0") + parent = tmp_path / "nested" + parent.mkdir() + second = parent / "second.lmdb" + _create_test_lmdb(str(second), nframes=4, natoms=8) + config = {"systems": [lmdb_dir, str(parent)], "batch_size": 2} + data = create_dataset(config, ["O", "H"]) + loader = create_dataloader(data, config, seed=0) + try: + np.testing.assert_allclose(loader.probabilities, [5 / 7, 2 / 7]) + assert loader.epoch_length == 7 + batches = [batch for _ in range(3) for batch in loader] + assert {int(batch["sid"]) for batch in batches} == {0, 1} + assert {batch["coord"].shape[1] for batch in batches} == {6, 8} + finally: + loader.close() + assert all(source._reader.closed for source in iter_data_sources(data)) + + +def test_filter_can_remove_the_numpy_subset_of_a_mixed_group( + lmdb_dir: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0") + numpy = str(Path(__file__).parent / "water" / "data" / "single") + config = {"systems": [numpy, lmdb_dir], "batch_size": "filter:8"} + pt_data = create_dataset(config, ["O", "H"]) + numpy_data = get_data( + config, + 0, + ["O", "H"], + None, + lmdb_factory=lambda path, batch_size, auto_prob: LmdbDataSystem( + path, ["O", "H"], batch_size, auto_prob, num_workers=0 + ), + ) + try: + assert isinstance(pt_data, LmdbDataset) + assert numpy_data.get_batch()["coord"].shape == (1, 6, 3) + finally: + pt_data.close() + numpy_data.close() + + +def test_group_probability_preserves_single_lmdb_auto_prob( + auto_prob_lmdb: str, lmdb_dir: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0") + config = [ + { + "systems": auto_prob_lmdb, + "batch_size": 2, + "auto_prob": "prob_sys_size;0:1:0.5;1:3:0.5", + "prob": 0.6, + }, + {"systems": lmdb_dir, "batch_size": 2, "prob": 0.4}, + ] + data = create_dataset(config, ["O", "H"]) + loader = create_dataloader(data, config, seed=0) + try: + assert data.children[0]._block_targets + np.testing.assert_allclose(loader.probabilities, [0.6, 0.4]) + assert len(data.dataloaders) == 2 + finally: + loader.close() + + +def test_distributed_lmdb_groups_separate_mass_from_local_padding( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0") + paths = [str(tmp_path / f"source_{index}.lmdb") for index in range(2)] + for path, frames in zip(paths, (1, 3), strict=True): + _create_test_lmdb(path, nframes=frames) + config = {"systems": paths, "batch_size": 1} + for rank in (0, 1): + data = create_dataset(config, ["O", "H"]) + loader = create_dataloader(data, config, seed=0, rank=rank, world_size=2) + try: + np.testing.assert_allclose(loader.probabilities, [0.25, 0.75]) + np.testing.assert_array_equal(loader.batch_counts, [1, 2]) + np.testing.assert_array_equal(loader.sampling_sizes, [1, 3]) + assert loader.epoch_length == 4 + finally: + loader.close() diff --git a/source/tests/pt/test_sampler.py b/source/tests/pt/test_sampler.py index 88848c7e02..ee35acd6f8 100644 --- a/source/tests/pt/test_sampler.py +++ b/source/tests/pt/test_sampler.py @@ -17,6 +17,9 @@ from deepmd.dpmodel.utils import ( compute_total_numb_batch, ) +from deepmd.dpmodel.utils.training_utils import ( + resolve_model_prob, +) from deepmd.pt.utils import ( dp_random, ) @@ -168,6 +171,29 @@ def test_sampler_debug_info(self) -> None: frame = self.my_dataset.systems[sid].__getitem__(fid) self.assertTrue(np.allclose(coord, frame["coord"])) + def test_grouping_preserves_default_model_probabilities(self) -> None: + system = self.systems[0] + grouped = pt_dataloader.create_dataset( + [ + {"systems": [system] * 2, "batch_size": 1}, + {"systems": [system] * 3, "batch_size": 1}, + ], + self.type_map, + ) + flat = pt_dataloader.create_dataset( + {"systems": [system] * 5, "batch_size": 1}, self.type_map + ) + single = pt_dataloader.create_dataset( + {"systems": system, "batch_size": 1}, self.type_map + ) + probabilities = [ + resolve_model_prob( + ["large", "small"], None, {"large": data, "small": single} + ) + for data in (flat, grouped) + ] + np.testing.assert_allclose(probabilities, [[5 / 6, 1 / 6]] * 2) + def test_auto_prob_uniform(self) -> None: auto_prob_style = "prob_uniform" sampler = pt_dataloader.get_weighted_sampler( diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index 539b72f2c8..2a4156eaa5 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -35,6 +35,7 @@ collate_lmdb_frames, ) from deepmd.pt_expt.entrypoints.main import ( + _build_data_system, get_trainer, ) from deepmd.pt_expt.loss import ( @@ -43,6 +44,9 @@ from deepmd.pt_expt.train.training import ( Trainer, ) +from deepmd.pt_expt.train.validation import ( + FullValidator, +) from deepmd.pt_expt.utils.lmdb_dataset import ( LmdbDataSystem, ) @@ -58,6 +62,12 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_group import ( + GroupedDataSystem, +) +from deepmd.utils.data_system import ( + get_neighbor_stat_data, +) from .compile_utils import ( REQUIRES_SUPPORTED_COMPILE, @@ -452,7 +462,7 @@ def test_ragged_filter_uses_the_consumed_pass_length(self) -> None: trainer.multi_task = False trainer.world_size = 1 trainer.training_data_by_task = {"Default": data} - trainer.min_pair_dist_by_task = {"Default": 0.5} + trainer.pair_filter_enabled_by_task = {"Default": True} try: # This shuffle packs five batches and leaves frame 3 last; # the following shuffle packs four batches. @@ -700,6 +710,97 @@ def test_stat_batches_restart_after_non_divisible_pass(self) -> None: ) +class TestLmdbDatasetGroups(unittest.TestCase): + def setUp(self) -> None: + directory = tempfile.TemporaryDirectory() + self.addCleanup(directory.cleanup) + self.directory = directory.name + self.first = os.path.join(self.directory, "first.lmdb") + self.parent = os.path.join(self.directory, "nested") + os.mkdir(self.parent) + self.second = os.path.join(self.parent, "second.lmdb") + _create_test_lmdb(self.first, nframes=2, natoms=6) + _create_test_lmdb(self.second, nframes=8, natoms=8) + workers = patch.dict(os.environ, {"DP_LMDB_NUM_WORKERS": "0"}) + workers.start() + self.addCleanup(workers.stop) + + def test_multiple_paths_and_recursive_parent_use_flat_sampling(self) -> None: + for systems in ([self.first, self.parent], self.directory): + with self.subTest(systems=systems): + data = _build_data_system( + {"systems": systems, "batch_size": 1, "auto_prob": "prob_sys_size"}, + ["O", "H"], + seed=0, + ) + self.addCleanup(data.close) + self.assertIsInstance(data, GroupedDataSystem) + np.testing.assert_allclose(data.sys_probs, [0.2, 0.8]) + shapes = {data.get_batch()["coord"].shape[1] for _ in range(20)} + self.assertEqual(shapes, {6, 8}) + + def test_group_probability_and_statistics_keep_sources_separate(self) -> None: + data = _build_data_system( + [ + {"systems": self.first, "batch_size": 1, "prob": 0.6}, + {"systems": self.parent, "batch_size": 2, "prob": 0.4}, + ], + ["O", "H"], + seed=0, + ) + self.addCleanup(data.close) + np.testing.assert_allclose(data.sys_probs, [0.6, 0.4]) + samples = make_stat_input(data, 1) + self.assertEqual(len(samples), 2) + self.assertEqual([sample["coord"].shape[1] for sample in samples], [6, 8]) + self.assertEqual([sample["coord"].shape[0] for sample in samples], [1, 2]) + + def test_full_validation_caches_each_lmdb_separately(self) -> None: + data = _build_data_system( + {"systems": [self.first, self.second], "batch_size": 1}, ["O", "H"] + ) + self.addCleanup(data.close) + validator = FullValidator.__new__(FullValidator) + validator.validation_data = data + validator._lmdb_snapshots = {} + first = list(validator._iter_validation_data_systems()) + snapshots = dict(validator._lmdb_snapshots) + second = list(validator._iter_validation_data_systems()) + self.assertEqual(len(first), 2) + self.assertEqual(len(snapshots), 2) + self.assertEqual(validator._lmdb_snapshots, snapshots) + self.assertEqual( + [system.get_test()["coord"].shape[0] for system in second], [2, 8] + ) + + def test_distributed_epoch_accounts_for_each_leaf_padding(self) -> None: + paths = [ + os.path.join(self.directory, f"tiny_{index}.lmdb") for index in range(2) + ] + for path in paths: + _create_test_lmdb(path, nframes=1, natoms=6) + for rank in (0, 1): + data = _build_data_system( + {"systems": paths, "batch_size": 1}, + ["O", "H"], + seed=0, + rank=rank, + world_size=2, + ) + self.addCleanup(data.close) + trainer = Trainer.__new__(Trainer) + trainer.training_data_by_task = {"Default": data} + trainer.world_size = 2 + self.assertEqual(trainer._epoch_length("Default"), 2) + + def test_neighbor_statistics_visit_both_lmdbs(self) -> None: + data = get_neighbor_stat_data( + {"systems": [self.first, self.parent], "batch_size": 1}, ["O", "H"] + ) + self.assertEqual(data.get_ntypes(), 2) + self.assertEqual([system.get_natoms() for system in data.data_systems], [6, 8]) + + class TestLmdbTrainingLoop(unittest.TestCase): """End-to-end: get_trainer routes an LMDB path and runs training steps.""" @@ -786,6 +887,31 @@ def test_get_trainer_routes_lmdb(self) -> None: finally: os.chdir(cwd) + def test_get_trainer_runs_weighted_lmdb_groups(self) -> None: + config = self._make_lmdb_config(numb_steps=2) + config["training"]["training_data"] = [ + { + "systems": self.lmdb_path, + "batch_size": 1, + "prob": 0.6, + "min_pair_dist": 0.1, + }, + {"systems": self.val_lmdb_path, "batch_size": 2, "prob": 0.4}, + ] + config["training"]["validation_data"]["systems"] = [ + self.lmdb_path, + self.val_lmdb_path, + ] + config = normalize(update_deepmd_input(config, warning=False)) + cwd = os.getcwd() + os.chdir(self.tmpdir) + try: + trainer = get_trainer(config) + np.testing.assert_allclose(trainer.training_data.sys_probs, [0.6, 0.4]) + trainer.run() + finally: + os.chdir(cwd) + def test_numb_epoch_counts_passes_over_the_lmdb(self) -> None: """One epoch is one pass over the frames of the LMDB.""" config = self._make_lmdb_config() diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index fa4199195b..67306e6ee0 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -1112,7 +1112,7 @@ def test_replicated_numpy_filter_reads_a_complete_local_pass( trainer.multi_task = False trainer.world_size = world_size trainer.training_data_by_task = {"Default": data} - trainer.min_pair_dist_by_task = {"Default": 0.5} + trainer.pair_filter_enabled_by_task = {"Default": True} with ( patch("deepmd.utils.data.dp_random.shuffle"), patch.object(data, "get_batch", wraps=data.get_batch) as get_batch, From 2d8285515715e5dcc63549da4143fb45fbc9697a Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 22 Oct 2026 11:23:21 +0800 Subject: [PATCH 2/3] fix(data): balance uniform frame sampling across sources Align PT-expt default task weights with native PT data sizes while preserving explicit probabilities and epoch scheduling. Complete grouped CLI discovery, reader lifetime handling, and focused review regressions. --- deepmd/dpmodel/train/schedule.py | 5 +- deepmd/dpmodel/utils/lmdb_data.py | 80 ++++- deepmd/dpmodel/utils/training_utils.py | 20 +- deepmd/entrypoints/test.py | 50 ++- deepmd/pd/train/training.py | 19 + deepmd/pd/utils/dataloader.py | 90 ++++- deepmd/pt/utils/dataloader.py | 63 +++- deepmd/pt_expt/entrypoints/main.py | 2 + deepmd/pt_expt/train/training.py | 5 +- deepmd/pt_expt/utils/lmdb_dataset.py | 14 + deepmd/utils/argcheck.py | 15 +- deepmd/utils/data.py | 31 +- deepmd/utils/data_config.py | 8 +- deepmd/utils/data_group.py | 77 ++-- deepmd/utils/data_path.py | 14 +- deepmd/utils/data_sampling.py | 25 +- deepmd/utils/data_system.py | 67 +++- doc/train/multi-task-training.md | 2 +- doc/train/training-advanced.md | 10 +- source/tests/common/dpmodel/test_lmdb_data.py | 8 + .../common/dpmodel/test_train_schedule.py | 16 + source/tests/common/test_data_config.py | 9 +- source/tests/common/test_deepmd_data.py | 53 +++ source/tests/common/test_path.py | 4 + source/tests/pd/test_sampler.py | 46 +++ source/tests/pt/test_dp_test.py | 39 ++ .../pt/test_dp_test_spin_virtual_atoms.py | 3 +- source/tests/pt/test_lmdb_dataloader.py | 333 ++++++++++++++++++ source/tests/pt/test_training.py | 10 +- source/tests/pt_expt/test_lmdb_training.py | 19 + 30 files changed, 1020 insertions(+), 117 deletions(-) diff --git a/deepmd/dpmodel/train/schedule.py b/deepmd/dpmodel/train/schedule.py index 8fe2749dfc..1279343fc3 100644 --- a/deepmd/dpmodel/train/schedule.py +++ b/deepmd/dpmodel/train/schedule.py @@ -88,8 +88,9 @@ def resolve_step_schedule( Task keys in the order used by the trainer. Single-task runs pass the one key under which their data is registered. training_data : Mapping[str, Any] - Training data of every task, keyed like ``model_keys``. Only consulted - when multi-task sampling weights fall back to the per-task data size. + Training data of every task, keyed like ``model_keys``. Only consulted + when multi-task sampling weights default to native data sizes: + NumPy/HDF5 system counts or retained LMDB frame counts. epoch_length : Callable[[str], int] Maps a task key to the number of steps *this rank* performs during one epoch of that task. Backends whose data pipeline is sharded across diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index e299b458bb..cb4fc64223 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -1728,6 +1728,7 @@ def __init__( self._sampler = sampler self._epoch = 0 self._queued_pass_length = 0 + self._queued_nframes = 0.0 self._last_pass_length: int | None = None self._iterator = self._iter_epoch() self._num_workers = num_workers @@ -1748,13 +1749,15 @@ def __next__(self) -> dict[str, Any]: batch = self._collect(pending) elif self._deferred_indices is not None: indices, self._deferred_indices = self._deferred_indices, None - batch = self._reader.decode_batch(indices) + batch = self._decode(indices) else: - indices, self._queued_pass_length = self._next_indices() + indices, self._queued_pass_length, self._queued_nframes = ( + self._next_indices() + ) batch = self._decode(indices) self._last_pass_length = self._queued_pass_length - indices, self._queued_pass_length = self._next_indices() + indices, self._queued_pass_length, self._queued_nframes = self._next_indices() self._schedule(indices) return batch @@ -1834,7 +1837,7 @@ def _lose_the_pool(self, futures: "Sequence[Future[dict[str, Any]]]" = ()) -> No "decoders outlive the session that started it." ) - def _next_indices(self) -> tuple[list[int], int]: + def _next_indices(self) -> tuple[list[int], int, float]: """Return the next sampler batch and restart after exhaustion.""" try: return next(self._iterator) @@ -1843,8 +1846,8 @@ def _next_indices(self) -> tuple[list[int], int]: self._iterator = self._iter_epoch() return next(self._iterator) - def _iter_epoch(self) -> Iterator[tuple[list[int], int]]: - """Tag batches with the length of the sampler pass that owns them.""" + def _iter_epoch(self) -> Iterator[tuple[list[int], int, float]]: + """Snapshot each batch's pass length and rank-averaged frame count.""" set_epoch = getattr(self._sampler, "set_epoch", None) if callable(set_epoch): set_epoch(self._epoch) @@ -1852,8 +1855,39 @@ def _iter_epoch(self) -> Iterator[tuple[list[int], int]]: # enter another pass before the current batch reaches its consumer. # Snapshot the length lazily, after data requirements and layout are set. pass_length = len(self._sampler) - for indices in self._sampler: - yield indices, pass_length + frame_counts = getattr(self._sampler, "batch_frame_counts", None) + counts = frame_counts() if frame_counts is not None else None + for index, indices in enumerate(self._sampler): + nframes = len(indices) if counts is None else counts[index] + yield indices, pass_length, float(nframes) + + @property + def uniform_sampling_mass(self) -> float: + """Weight the next batch by its retained systems per sampled frame. + + Only sampler indices are inspected; no extra batch is decoded. The + frame count belongs to the queued batch, even when prefetching has + already advanced the sampler to another pass. Distributed samplers + report the same rank-averaged count on every rank. + """ + nsystems = self._reader.sampling_nsystems + if not nsystems: + return 0.0 + if self._pending is None and self._deferred_indices is None: + ( + self._deferred_indices, + self._queued_pass_length, + self._queued_nframes, + ) = self._next_indices() + return nsystems / self._queued_nframes + + @property + def uniform_nominal_mass(self) -> float: + """Return the systems-per-frame weight averaged over a sampler pass.""" + nsystems = self._reader.sampling_nsystems + if not nsystems: + return 0.0 + return nsystems / float(np.mean(self._sampler.batch_frame_counts())) def _submit(self, indices: list[int]) -> list[Future[dict[str, Any]]]: """Submit one batch as balanced contiguous chunks. @@ -1898,7 +1932,7 @@ def _schedule(self, indices: list[int]) -> None: def batch_pass_length(self) -> int: """Length of the last returned batch's pass, or the first pending pass.""" if self._last_pass_length is None: - return len(self._sampler) + return self._queued_pass_length or len(self._sampler) return self._last_pass_length @property @@ -2267,6 +2301,7 @@ def __init__( ).tolist() else: self._system_nframes = [self.nframes] + self.sampling_nsystems = sum(count > 0 for count in self._system_nframes) # Nominal batch size, reported to callers that want a single number. # The sampler never uses it: same-nloc modes go through @@ -3015,6 +3050,13 @@ def compute_block_targets( prob_sys_size_ext, ) + if auto_prob_style == "prob_uniform": + retained = [index for index, count in enumerate(system_nframes) if count > 0] + target = max(system_nframes, default=0) + if all(system_nframes[index] == target for index in retained): + return [] + return [([index], target) for index in retained] + # Parse block definitions from the auto_prob string # Format: "prob_sys_size;stt:end:weight;stt:end:weight;..." block_str = auto_prob_style.split(";")[1:] @@ -3686,6 +3728,10 @@ def __len__(self) -> int: """Number of batches the pending pass holds.""" return len(self.batches()) + def batch_frame_counts(self) -> np.ndarray: + """Return the number of frames in each pending optimizer batch.""" + return np.fromiter((len(batch) for batch in self.batches()), dtype=np.float64) + @property def total_batches(self) -> int: """Number of batches the pending pass holds over the whole dataset. @@ -3800,6 +3846,22 @@ def __len__(self) -> int: """Number of batches for this rank.""" return len(self._partition_batches(self._global_batches())) + def batch_frame_counts(self) -> np.ndarray: + """Return rank-averaged frame counts for each pending optimizer step. + + Padding matches the strided partition, so every rank obtains identical + source-selection weights even when their local batch sizes differ. + Custom partitioning strategies must provide their corresponding + rank-averaged counts when used with uniform source sampling. + """ + counts = np.fromiter( + (len(batch) for batch in self._global_batches()), dtype=np.float64 + ) + if counts.size == 0: + return counts + padded = np.resize(counts, len(self) * self._world_size) + return padded.reshape(-1, self._world_size).mean(axis=1) + @property def total_batches(self) -> int: """Number of batches one full pass holds, before the per-rank split.""" diff --git a/deepmd/dpmodel/utils/training_utils.py b/deepmd/dpmodel/utils/training_utils.py index 6f46cb8c93..baa5a0dfb0 100644 --- a/deepmd/dpmodel/utils/training_utils.py +++ b/deepmd/dpmodel/utils/training_utils.py @@ -82,7 +82,8 @@ def resolve_model_prob( model_keys : list[str] List of model keys. model_prob_config : dict[str, float] | None - User-specified model probabilities. If None, use data size. + User-specified model probabilities. If omitted or empty, use native + data size: NumPy/HDF5 system counts or retained LMDB frame counts. model_training_data : dict[str, object] Training data for each model. rank : int, optional @@ -112,10 +113,12 @@ def resolve_model_prob( if is_node_main_process(rank): log.info( "training.model_prob is not set or empty; defaulting to the " - "number of systems per task." + "native data size per task (NumPy/HDF5 systems or LMDB frames)." ) for ii, model_key in enumerate(model_keys): - model_prob[ii] = float(_training_data_size(model_training_data[model_key])) + model_prob[ii] = float( + get_training_data_size(model_training_data[model_key]) + ) if not np.all(np.isfinite(model_prob)): raise ValueError("Model prob must be finite.") if np.any(model_prob < 0.0): @@ -126,7 +129,16 @@ def resolve_model_prob( return model_prob / sum_prob -def _training_data_size(training_data: object) -> int: +def get_training_data_size(training_data: object) -> int: + """Return native data size for default multi-task weights. + + A data adapter may expose a size distinct from its batch/statistics system + count. This preserves the PyTorch convention of counting NumPy/HDF5 + systems and retained LMDB frames, including within composed datasets. + """ + get_size = getattr(training_data, "get_model_sampling_size", None) + if get_size is not None: + return int(get_size()) if hasattr(training_data, "get_nsystems"): return int(training_data.get_nsystems()) if not isinstance(training_data, Sized): diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index d9972b967e..a6cdb3bfa9 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -34,7 +34,10 @@ from deepmd.utils.data import ( DeepmdData, ) -from deepmd.utils.data_system import ( +from deepmd.utils.data_config import ( + data_group_configs, +) +from deepmd.utils.data_path import ( process_systems, ) from deepmd.utils.weight_avg import ( @@ -98,34 +101,25 @@ def test( if numb_test == 0: # only float has inf, but should work for min numb_test = float("inf") - if train_json is not None: - jdata = j_loader(train_json) - jdata = update_deepmd_input(jdata) - data_params = jdata.get("training", {}).get("training_data", {}) - systems = data_params.get("systems") - if not systems: - raise RuntimeError("No training data found in input json") - root = Path(train_json).parent - if isinstance(systems, str): - systems = str((root / Path(systems)).resolve()) - else: - systems = [str((root / Path(ss)).resolve()) for ss in systems] - patterns = data_params.get("rglob_patterns", None) - all_sys = process_systems(systems, patterns=patterns) - elif valid_json is not None: - jdata = j_loader(valid_json) + if train_json is not None or valid_json is not None: + input_json = train_json if train_json is not None else valid_json + data_kind = "training" if train_json is not None else "validation" + jdata = j_loader(input_json) jdata = update_deepmd_input(jdata) - data_params = jdata.get("training", {}).get("validation_data", {}) - systems = data_params.get("systems") - if not systems: - raise RuntimeError("No validation data found in input json") - root = Path(valid_json).parent - if isinstance(systems, str): - systems = str((root / Path(systems)).resolve()) - else: - systems = [str((root / Path(ss)).resolve()) for ss in systems] - patterns = data_params.get("rglob_patterns", None) - all_sys = process_systems(systems, patterns=patterns) + data_params = jdata.get("training", {}).get(f"{data_kind}_data") + if not data_params: + raise RuntimeError(f"No {data_kind} data found in input json") + all_sys = [] + for group in data_group_configs(data_params): + if not group.get("systems"): + raise RuntimeError(f"No {data_kind} data found in input json") + all_sys.extend( + process_systems( + group["systems"], + patterns=group.get("rglob_patterns"), + root_dir=Path(input_json).parent, + ) + ) elif datafile is not None: with open(datafile) as datalist: all_sys = datalist.read().splitlines() diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py index 8210ba8022..e90a6583e6 100644 --- a/deepmd/pd/train/training.py +++ b/deepmd/pd/train/training.py @@ -748,6 +748,25 @@ def single_model_finetune( self.profiling_file = training_params.get("profiling_file", "timeline.json") def run(self) -> None: + """Train while releasing native readers and producer threads on exit.""" + try: + self._run_training() + finally: + pipelines = ( + self.training_data, + self.validation_data, + self.training_dataloader, + self.validation_dataloader, + ) + for pipeline in pipelines: + sources = ( + pipeline.values() if isinstance(pipeline, dict) else (pipeline,) + ) + for source in sources: + if source is not None: + source.close() + + def _run_training(self) -> None: fout = ( open( self.disp_file, diff --git a/deepmd/pd/utils/dataloader.py b/deepmd/pd/utils/dataloader.py index 09e92ae7e4..d0e78ce4d8 100644 --- a/deepmd/pd/utils/dataloader.py +++ b/deepmd/pd/utils/dataloader.py @@ -7,10 +7,14 @@ Iterable, Iterator, ) +from functools import ( + cached_property, +) from multiprocessing.dummy import ( Pool, ) from threading import ( + Event, Thread, ) from typing import ( @@ -242,6 +246,22 @@ def __next__(self) -> Any: def __len__(self) -> int: return len(self.dataloaders) + @cached_property + def mean_batch_nframes(self) -> np.ndarray: + """Count frames per native batch, including tails and distributed padding.""" + return np.asarray( + [ + ( + sampler.num_samples + if isinstance(sampler, DistributedBatchSampler) + else len(sampler.sampler) + ) + / count + for sampler, count in zip(self.sampler_list, self.index, strict=True) + ], + dtype=np.float64, + ) + def __getitem__(self, idx: int) -> dict: # log.warning(str(paddle.distributed.get_rank())+" idx: "+str(idx)+" index: "+str(self.index[idx])) try: @@ -291,13 +311,32 @@ def __init__( self._queue = queue self._source = source # Main DL iterator self._max_len = max_len + self._stop_requested = Event() def run(self) -> None: for item in self._source: - self._queue.put(item) # Blocking if the queue is full + if self._stop_requested.is_set(): + return + if not self._put(item): + return # Signal the consumer we are done. - self._queue.put(_sentinel) + self._put(_sentinel) + + def _put(self, item: Any) -> bool: + """Wait for buffer capacity while allowing the owner to stop loading.""" + while not self._stop_requested.is_set(): + try: + self._queue.put(item, timeout=0.1) + return True + except queue.Full: + continue + return False + + def close(self) -> None: + """Stop producing batches and wait for the active read to complete.""" + self._stop_requested.set() + self.join() class BufferedIterator: @@ -322,6 +361,8 @@ def __len__(self) -> int: return self.total def __next__(self) -> Any: + if self._iterable is None: + raise StopIteration # Create consumer if not created yet if self._consumer is None: self._create_consumer() @@ -344,9 +385,19 @@ def __next__(self) -> Any: if isinstance(item, Exception): raise item if item is _sentinel: + self.close() raise StopIteration return item + def close(self) -> None: + """Release buffered batches before the owner closes the native loader.""" + if self._consumer is not None: + self._consumer.close() + self._consumer = None + self._iterable = None + while not self._queue.empty(): + self._queue.get_nowait() + def collate_batch(batch: list[dict]) -> dict: example = batch[0] @@ -366,6 +417,26 @@ def collate_batch(batch: list[dict]) -> dict: return result +class DpDataLoader(DataLoader): + """Own native reader threads and persistent per-system worker cursors.""" + + def __iter__(self) -> Iterator[dict[str, Any]]: + self._active_iterator = super().__iter__() + return self._active_iterator + + def close(self) -> None: + """Release native threads or workers, including a partially consumed pass.""" + iterator = getattr(self, "_active_iterator", None) + if iterator is not None: + iterator._try_shutdown_all() + self._active_iterator = None + self._iterator = None + + def __del__(self) -> None: + """Release workers if construction or initialization ends before training.""" + self.close() + + def get_weighted_sampler( training_data: DpLoaderSet, prob_style: str | list[float], sys_prob: bool = False ) -> WeightedRandomSampler: @@ -373,6 +444,11 @@ def get_weighted_sampler( training_data.index, sys_probs=prob_style if sys_prob else None, auto_prob="prob_sys_size" if sys_prob else prob_style, + batch_nframes=( + training_data.mean_batch_nframes + if not sys_prob and prob_style == "prob_uniform" + else None + ), ) log.debug("Generated weighted sampler with prob array: " + str(probs)) # training_data.total_batch is the size of one epoch, you can increase it to avoid too many rebuilding of iteraters @@ -438,10 +514,16 @@ def create_dataloader( return GroupedTrainingDataLoader(data, children, seeds[0]) params = config[0] if isinstance(config, list) else config sampler = get_sampler_from_params(data, params) - loader = DataLoader( + num_workers = env.NUM_WORKERS if dist.is_available() else 0 + loader = DpDataLoader( data, batch_sampler=BatchSampler(sampler=sampler, drop_last=False), - num_workers=env.NUM_WORKERS if dist.is_available() else 0, + num_workers=num_workers, + persistent_workers=( + num_workers > 0 + and params.get("auto_prob") == "prob_uniform" + and params.get("sys_probs") is None + ), collate_fn=lambda batch: batch[0], ) return TrainingDataLoader( diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py index 1267df0dfc..297895ff24 100644 --- a/deepmd/pt/utils/dataloader.py +++ b/deepmd/pt/utils/dataloader.py @@ -4,6 +4,9 @@ from collections.abc import ( Iterator, ) +from functools import ( + cached_property, +) from multiprocessing.dummy import ( Pool, ) @@ -216,6 +219,23 @@ def construct_dataset(system: str) -> DeepmdDataSetForLoader: def __len__(self) -> int: return len(self.dataloaders) + @cached_property + def mean_batch_nframes(self) -> np.ndarray: + """Count frames per native batch, including tails and distributed padding.""" + return np.asarray( + [len(loader.sampler) / len(loader) for loader in self.dataloaders], + dtype=np.float64, + ) + + @cached_property + def uniform_sampling_mass(self) -> float: + """Return the summed systems-per-frame weights of the native streams.""" + return float(np.reciprocal(self.mean_batch_nframes).sum()) + + @property + def uniform_nominal_mass(self) -> float: + return self.uniform_sampling_mass + def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: # log.warning(str(torch.distributed.get_rank())+" idx: "+str(idx)+" index: "+str(self.index[idx])) with torch.device("cpu"): @@ -278,6 +298,20 @@ def collate_batch(batch: list[dict[str, Any]]) -> dict[str, Any]: return result +class DpDataLoader(DataLoader): + """Own native CPU loading and persistent per-system reader cursors.""" + + def __iter__(self) -> Iterator[dict[str, Any]]: + with torch.device("cpu"): + return super().__iter__() + + def close(self) -> None: + """Release the persistent worker pool owned by this loader.""" + if self._iterator is not None: + self._iterator._shutdown_workers() + self._iterator = None + + def get_weighted_sampler( training_data: Any, prob_style: str | list[float], sys_prob: bool = False ) -> WeightedRandomSampler: @@ -285,6 +319,11 @@ def get_weighted_sampler( training_data.index, sys_probs=prob_style if sys_prob else None, auto_prob="prob_sys_size" if sys_prob else prob_style, + batch_nframes=( + training_data.mean_batch_nframes + if not sys_prob and prob_style == "prob_uniform" + else None + ), ) log.debug("Generated weighted sampler with prob array: " + str(probs)) # training_data.total_batch is the size of one epoch, you can increase it to avoid too many rebuilding of iterators @@ -353,7 +392,12 @@ def create_dataset( datasets.append( children[0] if len(children) == 1 - else GroupedDataset(children, child_configs, [None] * len(children)) + else GroupedDataset( + children, + child_configs, + [None] * len(children), + uniform=group.auto_prob == "prob_uniform", + ) ) return ( datasets[0] @@ -405,23 +449,26 @@ def create_dataloader( np.ones(1, dtype=np.float64), lambda: [len(loader)], sampling_sizes=lambda: [sampler.total_batches], + uniform_sampling=loader._batch_iterator, ) params = config[0] if isinstance(config, list) else config sampler = get_sampler_from_params(data, params) - loader = DataLoader( + num_workers = env.NUM_WORKERS if dist.is_available() else 0 + loader = DpDataLoader( data, sampler=sampler, batch_size=None, - num_workers=env.NUM_WORKERS if dist.is_available() else 0, + num_workers=num_workers, + persistent_workers=( + num_workers > 0 + and params.get("auto_prob") == "prob_uniform" + and params.get("sys_probs") is None + ), drop_last=False, collate_fn=lambda batch: batch, pin_memory=DEVICE.type != "cpu", ) - def iterate() -> Iterator[dict[str, Any]]: - with torch.device("cpu"): - return iter(loader) - return TrainingDataLoader( data, loader, @@ -430,5 +477,5 @@ def iterate() -> Iterator[dict[str, Any]]: sampling_sizes=lambda: global_batch_counts( [len(system) for system in data.systems], data.batch_sizes ), - iterator_factory=iterate, + uniform_sampling=data, ) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 4e6af3a7cb..aeb39688e0 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -104,6 +104,8 @@ def _build_data_system( world_size: int = 1, ) -> Any: """Bind native LMDB batching to the common grouped-data factory.""" + if seed is None and world_size > 1: + seed = 0 def build_lmdb( path: str, batch_size: int | str, auto_prob: str | None diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f93f47c55e..05a98785d7 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -2837,8 +2837,9 @@ def _epoch_length(self, model_key: str) -> int: The epoch length is ``ceil(max_i(local_batches[i] / sys_probs[i]))``. Each LMDB contributes its sampler's rank-local workload, including its own distributed padding. Replicated NumPy readers contribute - their global batch counts divided by world size. Sampling probabilities - use global source sizes and do not change with per-rank padding. + their global batch counts divided by world size. Size-based sampling + uses global source sizes without per-rank padding; uniform sampling + corrects for each native stream's rank-averaged batch frame count. This scheduled workload differs from a replicated reader's pass length, which is exposed separately through ``get_batch_pass_length``. """ diff --git a/deepmd/pt_expt/utils/lmdb_dataset.py b/deepmd/pt_expt/utils/lmdb_dataset.py index 72d407f2b3..20dcf97010 100644 --- a/deepmd/pt_expt/utils/lmdb_dataset.py +++ b/deepmd/pt_expt/utils/lmdb_dataset.py @@ -234,6 +234,20 @@ def get_nsystems(self) -> int: """Return one logical LMDB training dataset.""" return 1 + def get_model_sampling_size(self) -> int: + """Count retained frames, matching native PyTorch LMDB task weights.""" + return len(self._reader) + + @property + def uniform_sampling_mass(self) -> float: + """Return the queued batch's systems-per-frame selection weight.""" + return self._batch_iterator.uniform_sampling_mass + + @property + def uniform_nominal_mass(self) -> float: + """Return the pending pass's mean systems-per-frame selection weight.""" + return self._batch_iterator.uniform_nominal_mass + def get_batch_pass_length(self) -> int: """Return the local pass length belonging to the last returned batch.""" return self._batch_iterator.batch_pass_length diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index caffb5de52..98e71dd45b 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5800,7 +5800,7 @@ def training_data_args() -> Argument: - string "mix:N": LMDB data only. Frames of different atom counts share a batch, filled until the next frame would push the atom axis of the batch past `N`. How that axis is laid out follows from the model and needs no configuration of its own: a descriptor reading a flat node axis (the graph route of the PyTorch Exportable backend) takes the frames of a batch concatenated, so `N` counts real atoms and nothing is padded; every other descriptor takes them padded to the widest frame of the batch, so `N` counts the padded slots `nframes * max_nloc` and the shorter frames carry phantom atoms with `atype = -1` that the neighbor list, the model and the loss all skip. Unlike `"max:N"`, which leaves an under-filled batch whenever an nloc group is small, this keeps every batch close to `N` atoms. A lone frame with nloc > `N` still forms a batch of its own.\n\n\ If MPI is used, the value should be considered as the batch size per task.' doc_auto_prob_style = 'Determine the probability of systems automatically. The method is assigned by this key and can be\n\n\ -- "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems()\n\n\ +- "prob_uniform" : each retained system supplies the same long-run share of sampled frames. This includes internal LMDB systems and corrects for differing batch frame counts.\n\n\ - "prob_sys_size" : the probability of a system is proportional to the number of batches in the system\n\n\ - "prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;..." : the list of systems is divided into blocks. A block is specified by `stt_idx:end_idx:weight`, where `stt_idx` is the starting index of the system, `end_idx` is then ending (not including) index of the system, the probabilities of the systems in this block sums up to `weight`, and the relatively probabilities within this block is proportional to the number of batches in the system.' doc_sys_probs = ( @@ -5852,7 +5852,7 @@ def training_data_args() -> Argument: default="prob_sys_size", doc=doc_auto_prob_style + " When a group expands to multiple sources including LMDB, only " - "prob_sys_size is supported. Use separate groups to configure " + "prob_sys_size and prob_uniform are supported. Use separate groups to configure " "individual LMDB system probabilities.", alias=[ "auto_prob_style", @@ -5927,7 +5927,7 @@ def validation_data_args() -> list[ - string "filter:N": the same as `"max:N"` but additionally drops data whose atom count exceeds `N`. For npy data this removes whole systems with natoms > `N`; for LMDB data this removes individual frames with nloc > `N`.\n\n\ - string "mix:N": LMDB data only. Frames of different atom counts share a batch, filled until the next frame would push the atom axis of the batch past `N`. How that axis is laid out follows from the model: a descriptor reading a flat node axis takes the frames of a batch concatenated, so `N` counts real atoms and nothing is padded; every other descriptor takes them padded to the widest frame of the batch, so `N` counts the padded slots `nframes * max_nloc` and the shorter frames carry phantom atoms that the neighbor list, the model and the loss all skip.' doc_auto_prob_style = 'Determine the probability of systems automatically. The method is assigned by this key and can be\n\n\ -- "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems()\n\n\ +- "prob_uniform" : each retained system supplies the same long-run share of sampled frames. This includes internal LMDB systems and corrects for differing batch frame counts.\n\n\ - "prob_sys_size" : the probability of a system is proportional to the number of batches in the system\n\n\ - "prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;..." : the list of systems is divided into blocks. A block is specified by `stt_idx:end_idx:weight`, where `stt_idx` is the starting index of the system, `end_idx` is then ending (not including) index of the system, the probabilities of the systems in this block sums up to `weight`, and the relatively probabilities within this block is proportional to the number of batches in the system.' doc_sys_probs = ( @@ -5962,7 +5962,7 @@ def validation_data_args() -> list[ optional=True, default="prob_sys_size", doc=doc_auto_prob_style - + " Only prob_sys_size is supported when systems expands to " + + " Only prob_sys_size and prob_uniform are supported when systems expands to " "multiple sources including LMDB.", alias=[ "auto_prob_style", @@ -6161,8 +6161,11 @@ def training_args( doc_model_prob = ( "The visiting probability of each model for each training step in the " "multi-task mode. Only used when num_epoch_dict is not set. If not set " - "or an empty dict, defaults to weights proportional to the number of " - "systems per task." + "or an empty dict, defaults to each task's native data size: the number " + "of NPY/HDF5 systems, or retained LMDB frames on the PyTorch and " + "PyTorch Exportable backends. Grouped inputs sum their native sizes " + "without applying group probabilities. Batch size and nominal epoch " + "length do not affect these default weights." ) doc_data_dict = "The multiple definition of the data, used in the multi-task mode." doc_acc_freq = "Gradient accumulation steps (number of steps to accumulate gradients before performing an update)." diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 4d73e58d87..733be3a626 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -322,6 +322,30 @@ def get_item_paddle( """ return self.get_single_frame(index, num_worker) + @functools.cached_property + def _sampling_set_nframes(self) -> np.ndarray: + """Count frames in each set after its on-disk copy weights are applied.""" + sizes = [] + for set_dir, nframes in zip( + self.dirs, np.diff([0, *self.prefix_sum]), strict=True + ): + _, copies = self._load_data( + set_dir, "numb_copy", int(nframes), 1, must=False, default=1, dtype=int + ) + sizes.append(int(copies.sum())) + return np.asarray(sizes, dtype=np.int64) + + def get_mean_batch_nframes(self, batch_size: int) -> float: + """Return the initial pass's mean batch size, including undersized sets. + + A set smaller than one batch is returned whole. Larger sets drop their + short tail, as in :meth:`get_batch`. Copy weights are applied before + counting batches. + """ + sizes = self._sampling_set_nframes + counts = np.maximum(1, sizes // batch_size) + return float(np.dot(counts, np.minimum(sizes, batch_size)) / counts.sum()) + def get_batch(self, batch_size: int) -> dict: """Get a batch of data with `batch_size` frames. The frames are randomly picked from the data system. @@ -794,7 +818,12 @@ def _load_batch_set(self, set_name: DPPath) -> None: self.batch_set = self._load_set(set_name) if self.modifier is not None: self.modifier.modify_data(self.batch_set, self) - self.batch_set, _ = self._shuffle_data(self.batch_set) + self.batch_set, _ = self._shuffle_data(self.batch_set) + else: + # Resident frames already contain their requested copies. + indices = np.arange(self.batch_set["coord"].shape[0], dtype=np.int64) + dp_random.shuffle(indices) + self.batch_set = self._get_subdata(self.batch_set, indices) self.reset_get_batch() def reset_get_batch(self) -> None: diff --git a/deepmd/utils/data_config.py b/deepmd/utils/data_config.py index 20b6b1fb45..917e51849e 100644 --- a/deepmd/utils/data_config.py +++ b/deepmd/utils/data_config.py @@ -160,10 +160,14 @@ def resolve_data_groups( if not systems: raise ValueError(f"Data group {index} contains no supported data systems.") if len(systems) > 1 and any(is_lmdb(path) for path in systems): - if group.get("auto_prob", "prob_sys_size") != "prob_sys_size": + if group.get("auto_prob", "prob_sys_size") not in ( + "prob_sys_size", + "prob_uniform", + ): raise ValueError( f"Data group {index} expands to multiple sources including LMDB, " - "so auto_prob must be prob_sys_size. Use separate training_data groups " + "so auto_prob must be prob_sys_size or prob_uniform. " + "Use separate training_data groups " "to configure each LMDB's internal system probabilities." ) if group.get("sys_probs") is not None: diff --git a/deepmd/utils/data_group.py b/deepmd/utils/data_group.py index eb0f5c6fc7..fb4933bb01 100644 --- a/deepmd/utils/data_group.py +++ b/deepmd/utils/data_group.py @@ -15,6 +15,7 @@ from deepmd.dpmodel.utils.training_utils import ( compute_total_numb_batch, + get_training_data_size, ) from deepmd.utils.data_config import ( check_type_map_consistency, @@ -77,6 +78,13 @@ class GroupedDataSystem: node owns only the probability of selecting that child. Training and statistical system indices are separate flattened views: an LMDB is one training source but may expose several homogeneous statistical systems. + + Uniform composition selects each native source with weight ``M / n``, + where ``M`` counts its retained systems and ``n`` is its next batch's + frame count, or its mean batch size for a persistent fixed-pass stream. + Each system therefore supplies the same long-run frame share despite + different native batch sizes. Nominal epoch weights use pass means; + runtime weights can change as LMDB packing advances. """ def __init__( @@ -85,11 +93,13 @@ def __init__( probabilities: Sequence[float | None], *, seed: int | None = None, + uniform: bool = False, ) -> None: if not children or len(children) != len(probabilities): raise ValueError("Data groups require one probability per nonempty child.") self.children = tuple(children) self._assigned_probabilities = tuple(probabilities) + self._uniform = uniform self._rng = np.random.default_rng(data_child_seeds(seed, 1)[0]) self._probabilities: np.ndarray | None = None self._cdf: np.ndarray | None = None @@ -99,27 +109,12 @@ def __init__( self._align_metadata() def _align_metadata(self) -> None: - """Preserve dataset-wide NumPy type-map and atom-count contracts.""" - from deepmd.utils.data_system import ( - DeepmdDataSystem, - ) - + """Resolve metadata shared by the native data sources.""" sources = list(iter_data_sources(self)) self.type_map = check_type_map_consistency( [source.get_type_map() for source in sources] ) self._ntypes = max(source.get_ntypes() for source in sources) - numpy_sources = [ - source for source in sources if isinstance(source, DeepmdDataSystem) - ] - if len({source.mixed_type for source in numpy_sources}) > 1: - raise ValueError("All NumPy data systems must agree on mixed_type.") - for source in numpy_sources: - source.sys_ntypes = self._ntypes - source.natoms_vec = [ - system.get_natoms_vec(self._ntypes).astype(int) - for system in source.data_systems - ] self.mixed_type = any(source.mixed_type for source in sources) @property @@ -127,7 +122,10 @@ def component_probabilities(self) -> np.ndarray: """Return this sampling pass's probabilities over child datasets.""" if self._probabilities is None: self._probabilities = group_probabilities( - [sum(child.nbatches) for child in self.children], + [ + child.uniform_nominal_mass if self._uniform else sum(child.nbatches) + for child in self.children + ], self._assigned_probabilities, ) return self._probabilities @@ -195,7 +193,15 @@ def get_batch(self, sys_idx: int | None = None) -> dict[str, Any]: if self._remaining == 0: self._start_pass() assert self._cdf is not None - index = int(np.searchsorted(self._cdf, self._rng.random(), side="right")) + cdf = self._cdf + if self._uniform: + probabilities = group_probabilities( + [child.uniform_sampling_mass for child in self.children], + self._assigned_probabilities, + ) + cdf = np.cumsum(probabilities) + cdf[-1] = 1.0 + index = int(np.searchsorted(cdf, self._rng.random(), side="right")) self._remaining -= 1 return self._batch_context(self.children[index].get_batch(), index) @@ -258,6 +264,10 @@ def close(self) -> None: def get_nsystems(self) -> int: return sum(child.get_nsystems() for child in self.children) + def get_model_sampling_size(self) -> int: + """Sum native data sizes without applying within-task batch probabilities.""" + return sum(get_training_data_size(child) for child in self.children) + def get_ntypes(self) -> int: return self._ntypes @@ -319,6 +329,8 @@ def __init__( children: Sequence[Any], configs: Sequence[Mapping[str, Any]], probabilities: Sequence[float | None], + *, + uniform: bool = False, ) -> None: if not children or not len(children) == len(configs) == len(probabilities): raise ValueError( @@ -327,6 +339,7 @@ def __init__( self.children = tuple(children) self.configs = tuple(configs) self.probability_specs = tuple(probabilities) + self.uniform = uniform self.filter_enabled: tuple[bool, ...] | None = None def __len__(self) -> int: @@ -391,14 +404,22 @@ def __init__( batch_counts: Callable[[], Sequence[int]], *, sampling_sizes: Callable[[], Sequence[int]] | None = None, - iterator_factory: Callable[[], Iterator[dict[str, Any]]] | None = None, + uniform_sampling: Any = None, ) -> None: self.dataset = dataset self.loader = loader self.probabilities = np.asarray(probabilities, dtype=np.float64) self._batch_counts = batch_counts self._sampling_sizes = sampling_sizes or batch_counts - self._iterator_factory = iterator_factory or loader.__iter__ + self._uniform_sampling = uniform_sampling + + @property + def uniform_sampling_mass(self) -> float: + return self._uniform_sampling.uniform_sampling_mass + + @property + def uniform_nominal_mass(self) -> float: + return self._uniform_sampling.uniform_nominal_mass @property def batch_counts(self) -> np.ndarray: @@ -418,7 +439,7 @@ def epoch_length(self) -> int: return compute_total_numb_batch(self.batch_counts, self.system_probabilities) def __iter__(self) -> Iterator[dict[str, Any]]: - return self._iterator_factory() + return iter(self.loader) def __len__(self) -> int: return len(self.loader) @@ -452,7 +473,12 @@ def __init__( @property def probabilities(self) -> np.ndarray: return group_probabilities( - [float(child.sampling_sizes.sum()) for child in self.children], + [ + child.uniform_nominal_mass + if self.dataset.uniform + else float(child.sampling_sizes.sum()) + for child in self.children + ], self.dataset.probability_specs, ) @@ -493,6 +519,13 @@ def __iter__(self) -> Iterator[dict[str, Any]]: ) offsets = np.cumsum([0, *(len(child.batch_counts) for child in self.children)]) for _ in range(count): + if self.dataset.uniform: + probabilities = group_probabilities( + [child.uniform_sampling_mass for child in self.children], + self.dataset.probability_specs, + ) + cdf = np.cumsum(probabilities) + cdf[-1] = 1.0 index = int(np.searchsorted(cdf, self._rng.random(), side="right")) if index not in self._iterators: self._iterators[index] = iter(self.children[index]) diff --git a/deepmd/utils/data_path.py b/deepmd/utils/data_path.py index 986018fd0c..4559121a3f 100644 --- a/deepmd/utils/data_path.py +++ b/deepmd/utils/data_path.py @@ -71,9 +71,11 @@ def expand_sys_str( name = str(candidate) if isinstance(candidate, DPOSPath): path = candidate.path - if _inside_lmdb(path, root.path): + if not candidate.is_dir(): continue - if candidate.is_dir() and is_lmdb(name): + if path != root.path and _inside_lmdb(path, root.path): + continue + if is_lmdb(name): if "lmdb" in formats: systems[name] = None continue @@ -92,18 +94,26 @@ def process_systems( patterns: list[str] | None = None, *, formats: tuple[DatasetFormat, ...] = ("deepmd", "lmdb"), + root_dir: str | Path | None = None, ) -> list[str]: """Expand every input path through the same format-aware discovery. Explicit input order and repeated input paths are preserved, so an input list defines the order used by per-system sampling probabilities. Each entry may name a dataset directly or a parent containing datasets. + When ``root_dir`` is supplied, relative filesystem paths are resolved + against it before discovery; HDF5 group fragments remain unchanged. """ paths = [systems] if isinstance(systems, str) else systems if not isinstance(paths, list) or any(not isinstance(path, str) for path in paths): raise ValueError( f"Invalid systems value {systems!r}: expected str or list[str]." ) + if root_dir is not None: + paths = [ + str((Path(root_dir) / file_name).resolve()) + separator + group + for file_name, separator, group in (path.partition("#") for path in paths) + ] return [ system for path in paths diff --git a/deepmd/utils/data_sampling.py b/deepmd/utils/data_sampling.py index b0757417ef..79fe1ea44d 100644 --- a/deepmd/utils/data_sampling.py +++ b/deepmd/utils/data_sampling.py @@ -68,20 +68,39 @@ def prob_sys_size_ext( return sys_probs +def uniform_sampling_weights(batch_nframes: Sequence[float] | np.ndarray) -> np.ndarray: + """Return inverse batch sizes, excluding streams that supply no frames.""" + sizes = np.asarray(batch_nframes, dtype=np.float64) + return np.divide(1.0, sizes, out=np.zeros_like(sizes), where=sizes > 0) + + def system_probabilities( nbatches: Sequence[int] | np.ndarray, *, sys_probs: Sequence[float] | None = None, auto_prob: str = "prob_sys_size", + batch_nframes: Sequence[float] | np.ndarray | None = None, ) -> np.ndarray: - """Resolve the established per-system sampling rules.""" + """Resolve batch probabilities, correcting uniform sampling for frame counts. + + ``batch_nframes`` gives the mean frames per batch of each native stream. + Its reciprocal makes ``prob_uniform`` allocate equal sampled frame shares + rather than equal batch counts. Without batching metadata, every draw is + treated as one frame. Explicit ``sys_probs`` remain batch probabilities. + """ nsystems = len(nbatches) if nsystems == 0: raise ValueError("No systems provided.") if sys_probs is not None: return process_sys_probs(sys_probs, nbatches) if auto_prob == "prob_uniform": - return np.full(nsystems, 1.0 / nsystems, dtype=np.float64) + sizes = ( + np.ones(nsystems, dtype=np.float64) + if batch_nframes is None + else batch_nframes + ) + weights = uniform_sampling_weights(sizes) + return group_probabilities(weights, [None] * nsystems) if auto_prob == "prob_sys_size": auto_prob = f"prob_sys_size;0:{nsystems}:1.0" if not auto_prob.startswith("prob_sys_size;"): @@ -90,7 +109,7 @@ def system_probabilities( def group_probabilities( - sizes: Sequence[int] | np.ndarray, probabilities: Sequence[float | None] + sizes: Sequence[float] | np.ndarray, probabilities: Sequence[float | None] ) -> np.ndarray: """Assign group probabilities, sharing any remainder by group size.""" if len(sizes) != len(probabilities): diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index ae124859ad..24c62f7811 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -45,6 +45,7 @@ from deepmd.utils.data_sampling import process_sys_probs as process_sys_probs from deepmd.utils.data_sampling import ( system_probabilities, + uniform_sampling_weights, ) from deepmd.utils.out_stat import ( compute_stats_from_redu, @@ -106,7 +107,7 @@ def __init__( automatically by the number of batches in the system. auto_prob_style : str Determine the probability of systems automatically. The method is assigned by this key and can be - - "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems() + - "prob_uniform" : every system supplies the same sampled frame share. - "prob_sys_size" : the probability of a system is proportional to the number of batches in the system - "prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;..." : the list of systems is divided into blocks. A block is specified by `stt_idx:end_idx:weight`, @@ -457,9 +458,41 @@ def set_sys_probs( auto_prob_style: str = "prob_sys_size", ) -> None: self.sys_probs = system_probabilities( - self.nbatches, sys_probs=sys_probs, auto_prob=auto_prob_style + self.nbatches, + sys_probs=sys_probs, + auto_prob=auto_prob_style, + batch_nframes=( + self.mean_batch_nframes + if sys_probs is None and auto_prob_style == "prob_uniform" + else None + ), ) + @cached_property + def mean_batch_nframes(self) -> np.ndarray: + """Return native mean batch sizes for per-system selection.""" + if self.mixed_systems: + return np.ones(self.nsystems, dtype=np.float64) + return np.asarray( + [ + system.get_mean_batch_nframes(size) + for system, size in zip(self.data_systems, self.batch_size, strict=True) + ], + dtype=np.float64, + ) + + @cached_property + def uniform_sampling_mass(self) -> float: + """Return the summed systems-per-frame weights of the native streams.""" + if self.mixed_systems: + return float(self.nsystems / self.batch_size[0]) + return float(uniform_sampling_weights(self.mean_batch_nframes).sum()) + + @property + def uniform_nominal_mass(self) -> float: + """Return the systems-per-frame weight averaged over native passes.""" + return self.uniform_sampling_mass + def get_batch(self, sys_idx: int | None = None) -> dict: # batch generation style altered by Ziyao Li: # one should specify the "sys_prob" and "auto_prob_style" params @@ -505,7 +538,6 @@ def get_batch_standard(self, sys_idx: int | None = None) -> dict: if sys_idx is not None: self.pick_idx = sys_idx else: - # prob = self._get_sys_probs(sys_probs, auto_prob_style) self.pick_idx = dp_random.choice( np.arange(self.nsystems, dtype=np.int32), p=self.sys_probs ) @@ -866,7 +898,7 @@ def get_data( seeds = ( iter(data_child_seeds(seed, len(groups) + 1)) if needs_composition else iter(()) ) - datasets = [] + group_sources = [] for group in groups: sources = resolve_data_sources(group) has_lmdb = any(source.format == "lmdb" for source in sources) @@ -892,12 +924,33 @@ def get_data( source.systems[0], source.batch_size, group.auto_prob ) children.append(child) - data = ( + group_sources.append(children) + + sources = [source for children in group_sources for source in children] + numpy_sources = [ + source for source in sources if isinstance(source, DeepmdDataSystem) + ] + if len({source.mixed_type for source in numpy_sources}) > 1: + raise ValueError("All NumPy data systems must agree on mixed_type.") + ntypes = max(source.get_ntypes() for source in sources) + for source in numpy_sources: + source.sys_ntypes = ntypes + source.natoms_vec = [ + system.get_natoms_vec(ntypes).astype(int) for system in source.data_systems + ] + datasets = [ + ( children[0] if len(children) == 1 - else GroupedDataSystem(children, [None] * len(children), seed=next(seeds)) + else GroupedDataSystem( + children, + [None] * len(children), + seed=next(seeds), + uniform=group.auto_prob == "prob_uniform", + ) ) - datasets.append(data) + for group, children in zip(groups, group_sources, strict=True) + ] return ( datasets[0] if len(datasets) == 1 diff --git a/doc/train/multi-task-training.md b/doc/train/multi-task-training.md index 6ebe5b4fbb..d49f9075fc 100644 --- a/doc/train/multi-task-training.md +++ b/doc/train/multi-task-training.md @@ -76,7 +76,7 @@ Specifically, there are several parts that need to be modified: - (Optional) {ref}`training/model_prob `: The sampling weight settings corresponding to each `model_key`, i.e., the probability weight in the training step. You can specify any positive real number weight for each task. The higher the weight, the higher the probability of being sampled in each training. - This setting is optional, and if not set, tasks will be sampled with weights proportional to the number of systems in each task. It is only used when `num_epoch_dict` is not set. + This setting is optional and is only used when `num_epoch_dict` is not set. The default weights use each task's native data size: the number of NPY/HDF5 systems, or the number of retained LMDB frames on the PyTorch and PyTorch Exportable backends. Grouped inputs add these native sizes, without applying within-task group probabilities. Both PyTorch backends follow the established PyTorch rule; default task weights do not depend on batch size or nominal epoch length. - (Optional) {ref}`training/num_epoch_dict `: The number of training epochs for each model branch, specified as a dictionary mapping `model_key` to epoch values (can be fractional). This allows different tasks to train for different numbers of epochs, which is particularly useful for multi-task fine-tuning scenarios diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 824d92d887..ca3370c217 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -57,7 +57,7 @@ The sections {ref}`training_data ` and {ref}`validation_ - `list`: {ref}`systems ` gives a list of paths. Each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories. - The PyTorch and PyTorch Exportable backends also accept LMDB paths, lists of LMDB paths, and parent directories containing LMDB datasets. An LMDB dataset is terminal: discovery does not inspect its contents as additional datasets. NPY and LMDB sources may appear in the same group. Other backends discover DeePMD NPY/HDF5 systems only. - At each training step, DeePMD-kit randomly picks {ref}`batch_size ` frame(s) from one of the systems. The probability of using a system is by default in proportion to the number of batches in the system. More options are available for automatically determining the probability of using systems. One can set the key {ref}`auto_prob ` to - - `"prob_uniform"` all systems are used with the same probability. + - `"prob_uniform"` each retained system supplies the same long-run share of sampled frames. Batch-selection probabilities account for the number of frames returned per batch. - `"prob_sys_size"` the probability of using a system is proportional to its number of batches. - `"prob_sys_size; sidx_0:eidx_0:w_0; sidx_1:eidx_1:w_1;..."` the `list` of systems is divided into blocks. Block `i` has systems ranging from `sidx_i` to `eidx_i`. The probability of using a system from block `i` is proportional to `w_i`. Within one block, the probability of using a system is proportional to its size. - An example of using `"auto_prob"` is given below. The probability of using `systems[2]` is 0.4, and the sum of the probabilities of using `systems[0]` and `systems[1]` is 0.6. If the number of frames in `systems[1]` is twice of `system[0]`, then the probability of using `system[1]` is 0.4 and that of `system[0]` is 0.2. @@ -125,7 +125,7 @@ The sections {ref}`training_data ` and {ref}`validation_ The first group's datasets together receive probability 0.6, regardless of how many datasets its path contains. The second group's datasets together receive probability 0.4. Batches do not combine frames from different groups. -`prob` is an absolute probability between 0 and 1. If every group specifies it, the probabilities must sum to 1. Groups without `prob` share the remaining probability in proportion to their batch counts. Omitting `prob` everywhere therefore preserves size-based sampling; it does not give equally sized shares to differently sized groups. Within each group, `auto_prob` and `sys_probs` retain their existing meanings. +`prob` is an absolute probability between 0 and 1. If every group specifies it, the probabilities must sum to 1. Groups without `prob` share the remaining probability in proportion to their batch counts. Omitting `prob` everywhere therefore preserves size-based sampling; it does not give equally sized shares to differently sized groups. Within each group, `auto_prob` determines the system sampling rule; explicit `sys_probs` specify batch-selection probabilities. For a group containing exactly one LMDB, `auto_prob` addresses that LMDB's internal systems, independently of the group's `prob`: @@ -145,7 +145,11 @@ For a group containing exactly one LMDB, `auto_prob` addresses that LMDB's inter ] ``` -If a group expands to multiple sources and includes LMDB, it uses size-based sampling: `auto_prob` must remain `"prob_sys_size"`, its default, and `sys_probs` must be omitted. Custom block rules and `"prob_uniform"` are rejected because internal system indices belong to individual LMDB datasets. To assign source weights or configure different internal rules, put each LMDB in a separate group. Existing dictionary configurations, NPY `sys_probs`, and single-LMDB `auto_prob` rules remain supported. +If a group expands to multiple sources and includes LMDB, `auto_prob` can be `"prob_sys_size"` (the default) or `"prob_uniform"`, and `sys_probs` must be omitted. Custom block rules still require a single LMDB because their system indices are local to that dataset. To assign source weights or configure different block rules, put each LMDB in a separate group. Existing dictionary configurations, NPY `sys_probs`, and single-LMDB block rules remain supported. + +For `"prob_uniform"`, each NPY/HDF5 system and each nonempty internal LMDB system counts once after `filter:N` is applied. An LMDB without internal system identifiers counts as one system. The rule is independent of how systems are split among LMDB files: systems containing 100 and 900 frames each supply half the sampled frames. LMDB sampling balances the internal frame counts before applying the native batch packing; `mix:N` continues to combine systems. Source-selection weights also correct for differing batch frame counts. Uniform frame shares do not imply equal contributions to each batch or to the loss, and subsequent geometric frame rejection can change the retained training shares. + +This frame-share definition intentionally changes older NPY/HDF5 `prob_uniform` behavior when batch sizes differ. For example, systems returning batches of 1 and 10 frames receive batch probabilities of 10/11 and 1/11, rather than 1/2 each. Short native batches are included in the correction. Explicit `sys_probs` and outer group `prob` remain batch probabilities. `validation_data` remains a dictionary. Its `systems` accepts the same recursive paths and multiple-LMDB inputs on supported backends; `numb_btch` continues to apply to the whole validation dataset. diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index 304a4ba506..e92b31b7a8 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -2015,6 +2015,14 @@ def test_reader_no_system_ids_backward_compat(self): self.assertIsNone(reader.frame_system_ids) tmpdir.cleanup() + def test_uniform_targets_balance_retained_systems(self) -> None: + self.assertEqual( + compute_block_targets("prob_uniform", 3, [1, 0, 9]), + [([0], 9), ([2], 9)], + ) + self.assertEqual(compute_block_targets("prob_uniform", 2, [9, 9]), []) + self.assertEqual(compute_block_targets("prob_uniform", 2, [0, 0]), []) + def test_compute_block_targets_equal_weight(self): result = compute_block_targets( "prob_sys_size;0:1:0.5;1:2:0.5", nsystems=2, system_nframes=[100, 100] diff --git a/source/tests/common/dpmodel/test_train_schedule.py b/source/tests/common/dpmodel/test_train_schedule.py index dd85991f2c..b4feda6a9e 100644 --- a/source/tests/common/dpmodel/test_train_schedule.py +++ b/source/tests/common/dpmodel/test_train_schedule.py @@ -171,6 +171,22 @@ def test_multi_task_steps_keep_configured_model_prob() -> None: np.testing.assert_allclose(schedule.model_prob, [0.75, 0.25]) +@pytest.mark.parametrize("model_prob", [None, {}]) +def test_default_multi_task_probabilities_use_system_counts( + model_prob: dict[str, float] | None, +) -> None: + schedule = resolve_step_schedule( + {"numb_steps": 12, "model_prob": model_prob}, + multi_task=True, + model_keys=["small", "large"], + training_data={"small": FakeData(1), "large": FakeData(3)}, + epoch_length=_unreachable_epoch_length, + ) + + assert schedule.num_steps == 12 + np.testing.assert_allclose(schedule.model_prob, [0.25, 0.75]) + + def test_multi_task_requires_a_run_length() -> None: with pytest.raises(ValueError, match="num_epoch_dict must be set"): resolve_step_schedule( diff --git a/source/tests/common/test_data_config.py b/source/tests/common/test_data_config.py index 096caeb0d6..01a5b11481 100644 --- a/source/tests/common/test_data_config.py +++ b/source/tests/common/test_data_config.py @@ -39,7 +39,7 @@ def test_group_and_system_probabilities_are_independent() -> None: @pytest.mark.parametrize("systems_as_list", [False, True]) -def test_multi_lmdb_requires_default_auto_prob( +def test_multi_lmdb_accepts_size_and_uniform_auto_prob( tmp_path: Path, systems_as_list: bool ) -> None: first = tmp_path / "first.lmdb" @@ -55,9 +55,10 @@ def test_multi_lmdb_requires_default_auto_prob( {"systems": systems, "auto_prob": "prob_sys_size"} ) assert explicit_default[0].systems == groups[0].systems - for custom in ("prob_uniform", "prob_sys_size;0:1:1"): - with pytest.raises(ValueError, match="auto_prob must be prob_sys_size"): - resolve_data_groups({"systems": systems, "auto_prob": custom}) + uniform = resolve_data_groups({"systems": systems, "auto_prob": "prob_uniform"}) + assert uniform[0].systems == groups[0].systems + with pytest.raises(ValueError, match="auto_prob must be prob_sys_size"): + resolve_data_groups({"systems": systems, "auto_prob": "prob_sys_size;0:1:1"}) with pytest.raises(ValueError, match="sys_probs is not supported"): resolve_data_groups({"systems": systems, "sys_probs": [-1, -1]}) diff --git a/source/tests/common/test_deepmd_data.py b/source/tests/common/test_deepmd_data.py index 9e9e57d9eb..976e2302f9 100644 --- a/source/tests/common/test_deepmd_data.py +++ b/source/tests/common/test_deepmd_data.py @@ -8,7 +8,9 @@ Path, ) +import h5py import numpy as np +import pytest from deepmd.dpmodel.utils.dist_check import ( compute_min_pair_margin_single, @@ -22,6 +24,57 @@ ) +def test_cached_numpy_set_applies_copy_weights_once(tmp_path: Path) -> None: + np.savetxt(tmp_path / "type.raw", [0], fmt="%d") + set_path = tmp_path / "set.000" + set_path.mkdir() + np.save(set_path / "coord.npy", [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]) + np.save(set_path / "box.npy", np.tile(np.eye(3).reshape(1, 9), (2, 1))) + np.save(set_path / "numb_copy.npy", [2, 1]) + data = DeepmdData(str(tmp_path)) + + assert data.get_mean_batch_nframes(8) == 3 + for _ in range(3): + batch = data.get_batch(8) + np.testing.assert_array_equal(np.sort(batch["coord"][:, 0]), [0, 0, 1]) + + +@pytest.mark.parametrize("hdf5", [False, True]) +def test_uniform_numpy_counts_undersized_sets(tmp_path: Path, hdf5: bool) -> None: + paths = [] + container = tmp_path / "systems.h5" + for index, sizes in enumerate(((1, 5), (2, 8))): + root = tmp_path / f"system_{index}" + root.mkdir() + np.savetxt(root / "type.raw", [0], fmt="%d") + for set_index, size in enumerate(sizes): + set_path = root / f"set.{set_index:03d}" + set_path.mkdir() + np.save(set_path / "coord.npy", np.full((size, 3), index, dtype=np.float64)) + np.save(set_path / "box.npy", np.tile(np.eye(3).reshape(1, 9), (size, 1))) + paths.append(str(root)) + if hdf5: + with h5py.File(container, "a") as handle: + group = handle.create_group(root.name) + group.create_dataset("type.raw", data=[0]) + for file in root.rglob("*.npy"): + group.create_dataset( + str(file.relative_to(root)), data=np.load(file) + ) + paths[-1] = f"{container}#/{root.name}" + data = DeepmdDataSystem(paths, [3, 5], 1, auto_prob_style="prob_uniform") + means = [ + np.mean([data.get_batch(sys_idx=index)["coord"].shape[0] for _ in range(4)]) + for index in range(2) + ] + np.testing.assert_allclose(means, [2, 3.5]) + np.testing.assert_allclose(data.sys_probs * means, [14 / 11, 14 / 11]) + data = DeepmdDataSystem(paths, "mixed:4", 1, auto_prob_style="prob_uniform") + np.testing.assert_allclose(data.sys_probs, [0.5, 0.5]) + assert data.uniform_sampling_mass == 0.5 + assert data.get_batch()["coord"].shape[0] == 4 + + class TestDeepmdDataTypeMap(unittest.TestCase): def setUp(self) -> None: self.tmpdir = tempfile.TemporaryDirectory() diff --git a/source/tests/common/test_path.py b/source/tests/common/test_path.py index 7573c69a5b..c5fbac3b3e 100644 --- a/source/tests/common/test_path.py +++ b/source/tests/common/test_path.py @@ -130,3 +130,7 @@ def test_hdf5_paths_expand_inside_containers(self) -> None: self.assertEqual( process_systems([expected[1], expected[0]]), expected[::-1] ) + self.assertEqual( + process_systems("systems.h5#/nested", root_dir=directory), + [expected[1]], + ) diff --git a/source/tests/pd/test_sampler.py b/source/tests/pd/test_sampler.py index e920660704..37a23fed06 100644 --- a/source/tests/pd/test_sampler.py +++ b/source/tests/pd/test_sampler.py @@ -2,9 +2,15 @@ import json import os import unittest +from collections.abc import ( + Iterator, +) from pathlib import ( Path, ) +from threading import ( + Event, +) import numpy as np import paddle @@ -17,6 +23,8 @@ expand_sys_str, ) from deepmd.pd.utils.dataloader import ( + QUEUESIZE, + BufferedIterator, DpLoaderSet, get_weighted_sampler, ) @@ -28,6 +36,44 @@ CUR_DIR = os.path.dirname(__file__) +class TestBufferedIterator(unittest.TestCase): + def test_close_stops_production_with_a_full_queue(self) -> None: + full = Event() + extra_read = Event() + + class Source: + def __len__(self) -> int: + return 10000 + + def __iter__(self) -> Iterator[int]: + for index in range(len(self)): + if index == QUEUESIZE + 1: + full.set() + if index == QUEUESIZE + 2: + extra_read.set() + yield index + + buffered = BufferedIterator(Source()) + self.addCleanup(buffered.close) + self.assertEqual(next(buffered), 0) + self.assertTrue(full.wait(timeout=5)) + self.assertTrue(buffered._queue.full()) + producer = buffered._consumer + buffered.close() + self.assertFalse(producer.is_alive()) + self.assertFalse(extra_read.is_set()) + self.assertTrue(buffered._queue.empty()) + with self.assertRaises(StopIteration): + next(buffered) + buffered.close() + + def test_exhaustion_releases_the_producer(self) -> None: + buffered = BufferedIterator(range(3)) + self.assertEqual(list(buffered), [0, 1, 2]) + self.assertIsNone(buffered._consumer) + self.assertIsNone(buffered._iterable) + + class TestSampler(unittest.TestCase): def setUp(self): with open(str(Path(__file__).parent / "water/se_e2_a.json")) as fin: diff --git a/source/tests/pt/test_dp_test.py b/source/tests/pt/test_dp_test.py index 4826048bce..36e31f7cd5 100644 --- a/source/tests/pt/test_dp_test.py +++ b/source/tests/pt/test_dp_test.py @@ -222,6 +222,45 @@ def test_dp_test_input_json_rglob(self) -> None: os.unlink(tmp_model_path) self.assertTrue(os.path.exists(self.detail_file + ".e.out")) + def test_training_groups_use_input_relative_patterns(self) -> None: + trainer = get_trainer(deepcopy(self.config)) + model = torch.jit.script(trainer.model) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + model_path = root / "model.pth" + torch.jit.save(model, str(model_path)) + data_root = Path(__file__).parent / "water" / "data" + relative_root = os.path.relpath(data_root, root) + config = deepcopy(self.config) + config["training"]["training_data"] = [ + { + "systems": relative_root, + "rglob_patterns": ["single"], + "batch_size": "filter:1", + "prob": 0.0, + }, + { + "systems": [relative_root], + "rglob_patterns": ["data_0"], + "prob": 1.0, + }, + ] + input_path = root / "input.json" + input_path.write_text(json.dumps(config)) + detail = str(root / "details") + dp_test( + model=str(model_path), + system=None, + datafile=None, + train_json=str(input_path), + numb_test=1, + rand_seed=None, + shuffle_test=False, + detail_file=detail, + atomic=False, + ) + self.assertEqual(np.loadtxt(detail + ".e.out", ndmin=2).shape[0], 2) + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith(".pt"): diff --git a/source/tests/pt/test_dp_test_spin_virtual_atoms.py b/source/tests/pt/test_dp_test_spin_virtual_atoms.py index 5318cb39e6..dba6d4c092 100644 --- a/source/tests/pt/test_dp_test_spin_virtual_atoms.py +++ b/source/tests/pt/test_dp_test_spin_virtual_atoms.py @@ -147,9 +147,10 @@ def test_dp_test_spin_virtual_atom_metrics(tmp_path, monkeypatch, padding_label) capture_output=True, text=True, timeout=120, - check=True, + check=False, ) output = result.stdout + result.stderr + assert result.returncode == 0, output for label, key in (("Force atom RMSE", "rmse_fr"), ("Force spin RMSE", "rmse_fm")): assert label in output assert f"{expected[key][0]:e}" in output diff --git a/source/tests/pt/test_lmdb_dataloader.py b/source/tests/pt/test_lmdb_dataloader.py index 0ba4f61c7d..148313c777 100644 --- a/source/tests/pt/test_lmdb_dataloader.py +++ b/source/tests/pt/test_lmdb_dataloader.py @@ -6,6 +6,11 @@ Consistency tests (dpmodel vs pt) live in source/tests/consistent/test_lmdb_data.py. """ +from itertools import ( + chain, + islice, + repeat, +) from pathlib import ( Path, ) @@ -19,6 +24,9 @@ import pytest import torch +from deepmd.dpmodel.train import ( + resolve_step_schedule, +) from deepmd.dpmodel.utils import ( lmdb_data, ) @@ -62,6 +70,10 @@ get_data, ) +from ..common.dpmodel.test_lmdb_data import ( + _create_mixed_sid_nloc_lmdb, +) + def _make_frame(natoms: int = 6, seed: int = 0) -> dict: """Create a synthetic frame dict as stored in LMDB.""" @@ -1685,3 +1697,324 @@ def test_distributed_lmdb_groups_separate_mass_from_local_padding( assert loader.epoch_length == 4 finally: loader.close() + + +@pytest.fixture +def uniform_sources(tmp_path: Path) -> tuple[list[str], list[np.ndarray]]: + paths = [str(tmp_path / f"source_{index}.lmdb") for index in range(3)] + _create_test_lmdb(paths[0], nframes=2, natoms=6) + _create_mixed_sid_nloc_lmdb(paths[1], [(1, 6), (7, 10), (3, 16)]) + _create_test_lmdb(paths[2], nframes=2, natoms=18) + markers = [np.zeros(2, dtype=int), np.repeat([1, 2, 3], [1, 7, 3]), np.full(2, 4)] + for path, values in zip(paths, markers, strict=True): + with lmdb.open(path) as database, database.begin(write=True) as transaction: + for index, value in enumerate(values): + key = format(index, "012d").encode() + frame = msgpack.unpackb(transaction.get(key), raw=False) + frame["energies"] = { + "type": " None: + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0") + paths, _ = uniform_sources + config = { + "systems": str(tmp_path) if recursive else paths, + "batch_size": batch_size, + "auto_prob": "prob_uniform", + } + requirements = [ + DataRequirementItem("energy", 1, atomic=False, must=True, high_prec=True), + DataRequirementItem("virial", 9, atomic=False, must=False, high_prec=True), + ] + if backend == "pt": + data = create_dataset(config, ["O", "H"]) + data.add_data_requirement(requirements) + owner = create_dataloader(data, config, seed=31) + batches = chain.from_iterable(repeat(owner)) + else: + data = get_data( + config, + 0, + ["O", "H"], + None, + seed=31, + lmdb_factory=lambda path, size, auto_prob: LmdbDataSystem( + path, ["O", "H"], size, auto_prob, seed=31, num_workers=0 + ), + ) + data.add_data_requirements(requirements) + data.use_ragged_batches(ragged) + owner = data + batches = (data.get_batch() for _ in repeat(None)) + counts = np.zeros(5, dtype=np.int64) + mixed = False + try: + for batch in islice(batches, 2500): + labels = np.asarray(batch["energy"]).reshape(-1).astype(int) + available = np.asarray(batch["fid"]).reshape(-1).astype(int) % 2 == 0 + assert np.all(available == available[0]) + assert bool(batch["find_virial"]) == bool(available[0]) + counts += np.bincount(labels, minlength=5) + mixed |= np.unique(labels).size > 1 + finally: + owner.close() + eligible = 3 if batch_size == "filter:12" else 5 + np.testing.assert_allclose( + counts[:eligible] / counts.sum(), 1 / eligible, atol=0.025 + ) + assert counts[eligible:].sum() == 0 + if batch_size == "mix:30": + assert mixed + + +@pytest.mark.parametrize("num_workers", [0, 1]) +def test_uniform_numpy_keeps_native_cursors_across_loader_passes( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, num_workers: int +) -> None: + from deepmd.pt.utils import ( + env, + ) + + monkeypatch.setattr(env, "NUM_WORKERS", num_workers) + np.savetxt(tmp_path / "type.raw", [0], fmt="%d") + set_path = tmp_path / "set.000" + set_path.mkdir() + np.save(set_path / "coord.npy", np.arange(24, dtype=float).reshape(8, 3)) + np.save(set_path / "box.npy", np.tile(np.eye(3).reshape(1, 9), (8, 1))) + config = { + "systems": [str(tmp_path), str(tmp_path)], + "batch_size": [7, 8], + "auto_prob": "prob_uniform", + } + data = create_dataset(config, ["O"]) + loader = create_dataloader(data, config, seed=7) + try: + np.testing.assert_allclose(loader.probabilities, [2 / 3, 1 / 3]) + np.testing.assert_allclose( + loader.probabilities * data.mean_batch_nframes, [8 / 3, 8 / 3] + ) + # Deterministic source selection exposes short tails across outer passes. + loader.loader.sampler.weights[:] = torch.tensor( + [1.0, 0.0], dtype=torch.float64, device="cpu" + ) + sizes = [batch["coord"].shape[0] for _ in range(2) for batch in loader] + assert sizes == [7, 1, 7, 1, 7, 1] + workers = tuple(loader.loader._iterator._workers) if num_workers else () + finally: + loader.close() + assert all(not worker.is_alive() for worker in workers) + + +@pytest.mark.parametrize("backend", ["pt", "pt_expt"]) +def test_uniform_group_combines_numpy_and_lmdb_systems( + uniform_sources: tuple[list[str], list[np.ndarray]], + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + backend: str, +) -> None: + from deepmd.pt.utils import ( + env, + ) + + monkeypatch.setattr(env, "NUM_WORKERS", 0) + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0") + paths, _ = uniform_sources + numpy_path = tmp_path / "numpy" + set_path = numpy_path / "set.000" + set_path.mkdir(parents=True) + np.savetxt(numpy_path / "type.raw", [0], fmt="%d") + np.save(set_path / "coord.npy", np.zeros((2, 3), dtype=np.float64)) + np.save(set_path / "box.npy", np.tile(np.eye(3).reshape(1, 9), (2, 1))) + np.save(set_path / "energy.npy", np.full((2, 1), 5.0, dtype=np.float64)) + config = { + "systems": [str(numpy_path), *paths[:2]], + "batch_size": [3, 2, 2], + "auto_prob": "prob_uniform", + } + requirements = [ + DataRequirementItem("energy", 1, atomic=False, must=True, high_prec=True) + ] + if backend == "pt": + data = create_dataset(config, ["O", "H"]) + data.add_data_requirement(requirements) + owner = create_dataloader(data, config, seed=4) + batches = chain.from_iterable(repeat(owner)) + else: + data = get_data( + config, + 0, + ["O", "H"], + None, + seed=4, + lmdb_factory=lambda path, size, auto_prob: LmdbDataSystem( + path, ["O", "H"], size, auto_prob, seed=4, num_workers=0 + ), + ) + data.add_data_requirements(requirements) + owner = data + batches = (data.get_batch() for _ in repeat(None)) + counts = np.zeros(6, dtype=np.int64) + try: + assert owner.children[0].uniform_sampling_mass == 0.5 + for batch in islice(batches, 2000): + labels = np.asarray(batch["energy"]).reshape(-1).astype(int) + counts += np.bincount(labels, minlength=6) + finally: + owner.close() + np.testing.assert_allclose(counts[[0, 1, 2, 3, 5]] / counts.sum(), 0.2, atol=0.025) + assert counts[4] == 0 + + +@pytest.mark.parametrize("grouped", [False, True]) +def test_multitask_defaults_match_native_pt_sizes( + uniform_sources: tuple[list[str], list[np.ndarray]], grouped: bool +) -> None: + paths, _ = uniform_sources + first = {"systems": paths[0], "batch_size": 1} + configs = { + "first": [{**first, "prob": 0.0}, {**first, "prob": 1.0}] if grouped else first, + "second": { + "systems": paths[1:], + "batch_size": "filter:12", + "auto_prob": "prob_uniform", + }, + } + expected = [1 / 3, 2 / 3] if grouped else [0.2, 0.8] + + def unused_epoch_length(key: str) -> int: + raise AssertionError( + f"Default model probabilities must not read epoch length: {key}" + ) + + for backend in ("pt", "pt_expt"): + datasets = { + key: ( + create_dataset(config, ["O", "H"]) + if backend == "pt" + else get_data( + config, + 0, + ["O", "H"], + None, + lmdb_factory=lambda path, size, auto_prob: LmdbDataSystem( + path, ["O", "H"], size, auto_prob, num_workers=0 + ), + ) + ) + for key, config in configs.items() + } + try: + schedule = resolve_step_schedule( + {"numb_steps": 5}, + multi_task=True, + model_keys=list(configs), + training_data=datasets, + epoch_length=unused_epoch_length, + ) + np.testing.assert_allclose(schedule.model_prob, expected) + finally: + for data in datasets.values(): + data.close() + + +@pytest.mark.parametrize("backend", ["pt", "pt_expt"]) +def test_uniform_distributed_sources_use_the_queued_global_step( + uniform_sources: tuple[list[str], list[np.ndarray]], + monkeypatch: pytest.MonkeyPatch, + backend: str, +) -> None: + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0") + paths, _ = uniform_sources + config = {"systems": paths, "batch_size": "mix:30", "auto_prob": "prob_uniform"} + requirements = [ + DataRequirementItem("energy", 1, atomic=False, must=True, high_prec=True) + ] + owners, iterators = [], [] + for rank in range(3): + if backend == "pt": + data = create_dataset(config, ["O", "H"]) + data.add_data_requirement(requirements) + owner = create_dataloader(data, config, seed=12, rank=rank, world_size=3) + iterator = chain.from_iterable(repeat(owner)) + else: + data = get_data( + config, + 0, + ["O", "H"], + None, + seed=12, + lmdb_factory=lambda path, size, auto_prob: LmdbDataSystem( + path, + ["O", "H"], + size, + auto_prob, + seed=12, + rank=rank, + world_size=3, + num_workers=0, + ), + ) + data.add_data_requirements(requirements) + data.use_ragged_batches(True) + owner = data + iterator = (source.get_batch() for source in repeat(data)) + owners.append(owner) + iterators.append(iterator) + counts = np.zeros(5, dtype=np.int64) + try: + for _ in range(1000): + masses = [ + [child.uniform_sampling_mass for child in owner.children] + for owner in owners + ] + np.testing.assert_allclose(masses, np.tile(masses[0], (3, 1))) + batches = [next(iterator) for iterator in iterators] + labels = [ + np.asarray(batch["energy"]).reshape(-1).astype(int) for batch in batches + ] + sources = [ + 0 if values[0] == 0 else 2 if values[0] == 4 else 1 for values in labels + ] + assert sources[0] == sources[1] == sources[2] + mean_frames = np.mean([len(values) for values in labels]) + assert masses[0][sources[0]] * mean_frames == pytest.approx( + [1, 3, 1][sources[0]] + ) + counts += np.bincount(np.concatenate(labels), minlength=5) + finally: + for owner in owners: + owner.close() + np.testing.assert_allclose(counts / counts.sum(), 0.2, atol=0.035) diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index 85a8a1a749..e2c25ea025 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -1267,15 +1267,11 @@ def get_data( class TestEMATraining(unittest.TestCase): def setUp(self) -> None: - import deepmd.pt.train.training as training_module import deepmd.pt.utils.env as env_module - self._num_workers_state = ( - (env_module, env_module.NUM_WORKERS), - (training_module, training_module.NUM_WORKERS), - ) + self._num_workers_state = env_module.NUM_WORKERS + self.addCleanup(setattr, env_module, "NUM_WORKERS", self._num_workers_state) env_module.NUM_WORKERS = 0 - training_module.NUM_WORKERS = 0 self._cwd = os.getcwd() self._tmpdir = tempfile.TemporaryDirectory() os.chdir(self._tmpdir.name) @@ -1306,8 +1302,6 @@ def setUp(self) -> None: } def tearDown(self) -> None: - for module, num_workers in self._num_workers_state: - module.NUM_WORKERS = num_workers os.chdir(self._cwd) self._tmpdir.cleanup() diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index 2a4156eaa5..f710aeb64b 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -22,6 +22,7 @@ import msgpack import numpy as np +import deepmd.utils.random as dp_random from deepmd.dpmodel.utils import lmdb_data as lmdb_data_module from deepmd.dpmodel.utils.batch import ( normalize_batch, @@ -725,6 +726,24 @@ def setUp(self) -> None: workers.start() self.addCleanup(workers.stop) + def test_unseeded_distributed_groups_share_source_choices(self) -> None: + sequences = [] + for rank in range(2): + dp_random.seed(17 + rank) + data = _build_data_system( + { + "systems": [self.first, self.second], + "batch_size": 1, + "auto_prob": "prob_uniform", + }, + ["O", "H"], + rank=rank, + world_size=2, + ) + self.addCleanup(data.close) + sequences.append([data.get_batch()["coord"].shape[1] for _ in range(24)]) + self.assertEqual(sequences[0], sequences[1]) + def test_multiple_paths_and_recursive_parent_use_flat_sampling(self) -> None: for systems in ([self.first, self.parent], self.directory): with self.subTest(systems=systems): From 4b2f6d9708345b010e658bf7a8614cb4538071d8 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 23 Oct 2026 11:38:51 +0800 Subject: [PATCH 3/3] fix(data): bound uniform LMDB sampling passes --- deepmd/dpmodel/utils/lmdb_data.py | 155 ++++++++---------- deepmd/pt/utils/lmdb_dataset.py | 50 ++---- doc/train/training-advanced.md | 2 + source/tests/common/dpmodel/test_lmdb_data.py | 83 ++++++++-- source/tests/common/test_loggers_training.py | 3 +- source/tests/pt/test_lmdb_dataloader.py | 33 +++- 6 files changed, 180 insertions(+), 146 deletions(-) diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index cb4fc64223..edf162c570 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -2264,7 +2264,7 @@ def __init__( # Re-key _frame_nlocs / _frame_system_ids into the dataset-index # space so that every downstream consumer (nloc_groups, system_groups, - # LmdbBatchSampler, _expand_indices_by_blocks) operates in a single, + # LmdbBatchSampler, _resample_indices_by_blocks) operates in a single, # self-consistent indexing scheme. keys_are_identity = retained_keys.size == self.nframes self._keys_are_identity = keys_are_identity @@ -3027,9 +3027,11 @@ def compute_block_targets( ) -> list[tuple[list[int], int]]: """Compute target frame count per block from auto_prob config. - Uses the same ``prob_sys_size_ext`` logic as the npy pipeline to parse - the ``auto_prob`` string, then converts per-system probabilities into - per-block target frame counts using the "max(frames/prob)" strategy. + Uniform sampling assigns each retained system the rounded-up mean frame + count. With ``N`` retained frames and ``M`` nonempty systems, the sampled + pass contains ``M * ceil(N / M)`` frames, fewer than ``N + M``. + Custom block rules use ``prob_sys_size_ext`` and the non-shrinking + ``max(frames / probability)`` strategy. Parameters ---------- @@ -3044,7 +3046,7 @@ def compute_block_targets( ------- list[tuple[list[int], int]] Each element is ``(system_indices_in_block, target_frame_count)``. - Returns empty list if no expansion is needed (all targets == actual). + Returns an empty list if no resampling is needed (all targets == actual). """ from deepmd.utils.data_sampling import ( prob_sys_size_ext, @@ -3052,7 +3054,9 @@ def compute_block_targets( if auto_prob_style == "prob_uniform": retained = [index for index, count in enumerate(system_nframes) if count > 0] - target = max(system_nframes, default=0) + if not retained: + return [] + target = (sum(system_nframes) + len(retained) - 1) // len(retained) if all(system_nframes[index] == target for index in retained): return [] return [([index], target) for index in retained] @@ -3145,118 +3149,76 @@ def compute_block_targets( return result -def _expand_indices_by_blocks( - indices: list[int], +def _resample_indices_by_blocks( + indices: np.ndarray | list[int], frame_system_ids: np.ndarray, - block_targets: list[tuple[list[int], int]], + lookup: np.ndarray, + group_targets: list[int], rng: np.random.Generator, - _block_total_actual: list[int] | None = None, - _sid_to_blk_arr: np.ndarray | None = None, - _group_block_targets: list[int] | None = None, ) -> list[int]: - """Expand frame indices according to block targets. + """Resample a homogeneous group's frames to its exact block targets. - For each block, computes the proportional target for the subset of - indices belonging to that block (within the current nloc group), - then applies full-copy + remainder sampling. + Targets are allocated globally by :func:`_allocate_group_block_targets`. + A smaller target selects frames without replacement; a larger target + uses full-copy + remainder sampling. Parameters ---------- - indices : list[int] - Frame indices in the current nloc group. + indices : numpy.ndarray or list[int] + Frame indices in the current homogeneous group. frame_system_ids : np.ndarray Per-frame system id for the entire dataset. - block_targets : list[tuple[list[int], int]] - Per-block (system_ids, total_target_frames). + lookup : numpy.ndarray + System-id to block-index table from :func:`system_block_lookup`. + group_targets : list[int] + Exact target for each block in this group. rng : np.random.Generator - RNG for remainder sampling. - _block_total_actual : list[int] or None - Pre-computed total actual frame count per block (across all nloc - groups). When provided, avoids an O(N) scan of frame_system_ids. - _sid_to_blk_arr : np.ndarray or None - Pre-computed lookup from :func:`system_block_lookup`. When provided, - avoids rebuilding the mapping for each call. - _group_block_targets : list[int] or None - Exact target for each block in this group. Production samplers - allocate these targets globally across all groups so independent - rounding cannot change a block's total size. + RNG for subset and remainder sampling. Returns ------- list[int] - Expanded indices. + Resampled indices. """ - n_blocks = len(block_targets) - - if _sid_to_blk_arr is None: - _sid_to_blk_arr = system_block_lookup(block_targets) - - sid_arr = np.asarray(frame_system_ids) idx_arr = np.asarray(indices, dtype=np.int64) - idx_blks = resolve_frame_blocks(idx_arr, sid_arr, _sid_to_blk_arr) - - # Pre-compute block_total_actual if not provided - if _block_total_actual is None and _group_block_targets is None: - _block_total_actual = count_group_blocks( - np.arange(sid_arr.size, dtype=np.int64), - sid_arr, - _sid_to_blk_arr, - n_blocks, - ).tolist() - - expanded_parts: list[np.ndarray] = [] + idx_blks = resolve_frame_blocks(idx_arr, frame_system_ids, lookup) + sampled_parts: list[np.ndarray] = [] # Unassigned indices unassigned_mask = idx_blks == -1 if np.any(unassigned_mask): - expanded_parts.append(idx_arr[unassigned_mask]) + sampled_parts.append(idx_arr[unassigned_mask]) - for blk_idx in range(n_blocks): + for blk_idx, target in enumerate(group_targets): blk_mask = idx_blks == blk_idx blk_idxs = idx_arr[blk_mask] n_actual = len(blk_idxs) if n_actual == 0: continue - if _group_block_targets is not None: - target = _group_block_targets[blk_idx] - else: - _, block_total_target = block_targets[blk_idx] - block_total_act = _block_total_actual[blk_idx] - - # Backward-compatible fallback for direct callers. Samplers pass - # exact group targets allocated by _allocate_group_block_targets. - if block_total_act > 0: - target = round(block_total_target * n_actual / block_total_act) - else: - target = n_actual - target = max(target, n_actual) # never shrink - if target < n_actual: - raise ValueError( - "Per-group auto-probability target cannot shrink original frames: " - f"target={target}, actual={n_actual}, block={blk_idx}" - ) + sampled_parts.append(rng.choice(blk_idxs, size=target, replace=False)) + continue # Full copies + remainder deficit = target - n_actual if deficit <= 0: - expanded_parts.append(blk_idxs) + sampled_parts.append(blk_idxs) else: full_copies = deficit // n_actual remainder = deficit % n_actual # Original + full copies if full_copies > 0: - expanded_parts.append(np.tile(blk_idxs, 1 + full_copies)) + sampled_parts.append(np.tile(blk_idxs, 1 + full_copies)) else: - expanded_parts.append(blk_idxs) + sampled_parts.append(blk_idxs) # Remainder: sample without replacement if remainder > 0: sampled = rng.choice(blk_idxs, size=remainder, replace=False) - expanded_parts.append(sampled) + sampled_parts.append(sampled) - if expanded_parts: - return np.concatenate(expanded_parts).tolist() + if sampled_parts: + return np.concatenate(sampled_parts).tolist() return [] @@ -3403,13 +3365,16 @@ def _allocate_group_block_targets( groups: list[np.ndarray], frame_system_ids: np.ndarray, block_targets: list[tuple[list[int], int]], + rng: np.random.Generator, ) -> list[list[int]]: """Allocate every block target exactly across homogeneous groups. - Original frames form the non-shrinking baseline. Each block's expansion + For expansion, original frames form the non-shrinking baseline. Each block's deficit is apportioned by actual group size using the integer largest- remainder method. Stable group order breaks equal-remainder ties, which keeps distributed ranks deterministic without floating-point rounding. + A smaller target instead samples group counts from a uniform subset of + the block's frames, so no small group is permanently excluded by rounding. """ lookup = system_block_lookup(block_targets) group_actual = [ @@ -3422,11 +3387,6 @@ def _allocate_group_block_targets( group_targets = [counts.copy() for counts in group_actual] for block_index, (_system_ids, block_target) in enumerate(block_targets): block_actual = sum(counts[block_index] for counts in group_actual) - if block_target < block_actual: - raise ValueError( - "Auto-probability block target cannot shrink original frames: " - f"target={block_target}, actual={block_actual}, block={block_index}" - ) if block_actual == 0: if block_target != 0: raise ValueError( @@ -3435,6 +3395,20 @@ def _allocate_group_block_targets( ) continue + if block_target < block_actual: + # Uniform positions in the concatenated groups give the group + # counts of a global sample without replacement. Sampling within + # each group then gives every block frame the same inclusion chance. + stops = np.cumsum( + [counts[block_index] for counts in group_actual], dtype=np.int64 + ) + positions = rng.choice(block_actual, size=block_target, replace=False) + sampled_groups = np.searchsorted(stops, positions, side="right") + counts = np.bincount(sampled_groups, minlength=len(groups)) + for group_index, target in enumerate(counts): + group_targets[group_index][block_index] = int(target) + continue + deficit = block_target - block_actual remainders: list[tuple[int, int]] = [] allocated = 0 @@ -3579,7 +3553,7 @@ def _build_all_batches( Random number generator (deterministic for reproducibility). block_targets : list[tuple[list[int], int]] or None Per-block (system_ids, target_frame_count) from compute_block_targets. - When provided, indices are expanded via full-copy + remainder sampling. + When provided, each block is resampled to its target frame count. Returns ------- @@ -3599,22 +3573,21 @@ def _build_all_batches( if block_targets and reader.frame_system_ids is not None: sid_arr = reader.frame_system_ids group_block_targets = _allocate_group_block_targets( - groups, sid_arr, block_targets + groups, sid_arr, block_targets, rng ) sid_to_blk_arr = system_block_lookup(block_targets) for group_index, original_indices in enumerate(groups): indices = original_indices - # Expand each group independently using targets that were allocated + # Resample each group independently using targets that were allocated # globally, so that a block's total size is preserved exactly. if block_targets and sid_arr is not None and group_block_targets is not None: - indices = _expand_indices_by_blocks( + indices = _resample_indices_by_blocks( indices, sid_arr, - block_targets, + sid_to_blk_arr, + group_block_targets[group_index], rng, - _sid_to_blk_arr=sid_to_blk_arr, - _group_block_targets=group_block_targets[group_index], ) if shuffle: # ``permutation`` returns a copy, which matters because a group @@ -3668,7 +3641,7 @@ class LmdbBatchSampler: seed : int or None Random seed for reproducibility. block_targets : list[tuple[list[int], int]] or None - Per-block expansion targets from compute_block_targets. + Per-block resampling targets from compute_block_targets. """ def __init__( @@ -3769,7 +3742,7 @@ class DistributedLmdbBatchSampler: seed : int or None Base seed for deterministic RNG. All ranks must use the same seed. block_targets : list[tuple[list[int], int]] or None - Per-block expansion targets from compute_block_targets. + Per-block resampling targets from compute_block_targets. """ def __init__( diff --git a/deepmd/pt/utils/lmdb_dataset.py b/deepmd/pt/utils/lmdb_dataset.py index 8cce38a853..70dd25b864 100644 --- a/deepmd/pt/utils/lmdb_dataset.py +++ b/deepmd/pt/utils/lmdb_dataset.py @@ -27,10 +27,8 @@ collate_lmdb_frames, collect_lmdb_sampling_groups, compute_block_targets, - count_group_blocks, is_lmdb, resolve_per_atom_keys, - system_block_lookup, ) from deepmd.env import ( get_lmdb_num_workers, @@ -347,14 +345,12 @@ def print_summary(self, name: str, prob: Any) -> None: reader = self._reader # Per-block summary: original vs target frames block_lines = [] - total_original = 0 + total_original = len(reader) total_target = 0 - # Pre-compute block_total_actual for proportional scaling - block_total_actual: list[int] = [] + assigned_systems: set[int] = set() for sys_ids, target in self._block_targets: actual = sum(reader.system_nframes[s] for s in sys_ids) - block_total_actual.append(actual) - total_original += actual + assigned_systems.update(sys_ids) total_target += target # Compact range notation: sys[0-146] instead of sys[0,1,2,...,146] if len(sys_ids) > 3: @@ -367,42 +363,18 @@ def print_summary(self, name: str, prob: Any) -> None: f"{actual}->{target} (x{ratio:.2f})" ) - # A whole nloc group's block membership resolves as one indexing - # operation, which a dataset of 10^8 frames needs it to be. The - # lookup is the one the sampler allocates its targets with, so - # both agree on which systems a block claims. - n_blocks = len(self._block_targets) - lookup = system_block_lookup(self._block_targets) - - # Compute expanded nloc counts analytically (no actual expansion) - expanded_nloc_info = {} - for nloc, indices in sorted(reader.nloc_groups.items()): - if reader.frame_system_ids is None: - expanded_nloc_info[nloc] = len(indices) - continue - counts = count_group_blocks( - indices, reader.frame_system_ids, lookup, n_blocks - ) - expanded = len(indices) - int(counts.sum()) - for blk_idx, (_, blk_target) in enumerate(self._block_targets): - n_actual = int(counts[blk_idx]) - if n_actual == 0: - continue - bta = block_total_actual[blk_idx] - if bta > 0: - t = max(round(blk_target * n_actual / bta), n_actual) - else: - t = n_actual - expanded += t - expanded_nloc_info[nloc] = expanded - - total_expanded = sum(expanded_nloc_info.values()) + # Systems outside explicit blocks retain their original frames. + total_target += sum( + count + for sid, count in enumerate(reader.system_nframes) + if sid not in assigned_systems + ) n_groups = len(reader.nloc_groups) - ratio_all = total_expanded / total_original if total_original > 0 else 0 + ratio_all = total_target / total_original if total_original > 0 else 0 log.info( f"LMDB {name} auto_prob: " - f"{total_original}->{total_expanded} frames (x{ratio_all:.2f}), " + f"{total_original}->{total_target} frames (x{ratio_all:.2f}), " f"{n_groups} nloc groups, {len(self._block_targets)} blocks:" ) for bl in block_lines: diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index ca3370c217..fd12112818 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -149,6 +149,8 @@ If a group expands to multiple sources and includes LMDB, `auto_prob` can be `"p For `"prob_uniform"`, each NPY/HDF5 system and each nonempty internal LMDB system counts once after `filter:N` is applied. An LMDB without internal system identifiers counts as one system. The rule is independent of how systems are split among LMDB files: systems containing 100 and 900 frames each supply half the sampled frames. LMDB sampling balances the internal frame counts before applying the native batch packing; `mix:N` continues to combine systems. Source-selection weights also correct for differing batch frame counts. Uniform frame shares do not imply equal contributions to each batch or to the loss, and subsequent geometric frame rejection can change the retained training shares. +For an LMDB containing `N` retained frames in `M` nonempty systems, a uniform sampling pass draws `ceil(N / M)` frames from each system. Before distributed batch padding, the pass therefore contains fewer than `N + M` frame indices, rather than scaling with the largest system. Larger systems contribute a random subset without replacement; smaller systems retain their frames and repeat them to fill their quota. A pass need not visit every frame of a larger system. Subsets change with the sampler epoch and remain reproducible from its seed. `numb_epoch` and omitted outer-group probabilities use the resulting native batch counts; an explicit outer `prob` is unchanged. Custom block rules retain their existing non-shrinking targets. + This frame-share definition intentionally changes older NPY/HDF5 `prob_uniform` behavior when batch sizes differ. For example, systems returning batches of 1 and 10 frames receive batch probabilities of 10/11 and 1/11, rather than 1/2 each. Short native batches are included in the correction. Explicit `sys_probs` and outer group `prob` remain batch probabilities. `validation_data` remains a dictionary. Its `systems` accepts the same recursive paths and multiple-LMDB inputs on supported backends; `numb_btch` continues to apply to the whole validation dataset. diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index e92b31b7a8..4f3899550e 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -46,9 +46,9 @@ LmdbTestData, LmdbTestDataNlocView, _chop_mixed_nloc, - _expand_indices_by_blocks, _merge_lmdb_chunks, _remap_atom_types, + _resample_indices_by_blocks, collate_lmdb_frames, collect_lmdb_sampling_groups, compute_block_targets, @@ -56,6 +56,7 @@ decode_lmdb_frame, is_lmdb, make_neighbor_stat_data, + system_block_lookup, ) from deepmd.utils import random as dp_random from deepmd.utils.data import ( @@ -1982,7 +1983,7 @@ def test_positive_out_of_range_type_still_raises(self): class TestAutoProb(unittest.TestCase): """Test auto_prob support: frame_system_ids, compute_block_targets, - _expand_indices_by_blocks, and LmdbBatchSampler with block_targets. + _resample_indices_by_blocks, and LmdbBatchSampler with block_targets. """ @classmethod @@ -2018,11 +2019,45 @@ def test_reader_no_system_ids_backward_compat(self): def test_uniform_targets_balance_retained_systems(self) -> None: self.assertEqual( compute_block_targets("prob_uniform", 3, [1, 0, 9]), - [([0], 9), ([2], 9)], + [([0], 5), ([2], 5)], ) self.assertEqual(compute_block_targets("prob_uniform", 2, [9, 9]), []) self.assertEqual(compute_block_targets("prob_uniform", 2, [0, 0]), []) + def test_uniform_target_count_is_bounded(self) -> None: + """Uniform quotas stay within one rounding unit per retained system.""" + counts = [100000] + [1] * 999 + targets = compute_block_targets("prob_uniform", len(counts), counts) + self.assertEqual({target for _, target in targets}, {101}) + total = sum(target for _, target in targets) + self.assertEqual(total, 101000) + self.assertGreaterEqual(total, sum(counts)) + self.assertLess(total, sum(counts) + len(counts)) + + def test_uniform_sampler_balances_a_bounded_pass(self) -> None: + """A uniform pass can undersample large systems without losing them.""" + reader = LmdbDataReader(self._lmdb_path, ["O", "H"], batch_size=7) + self.addCleanup(reader.close) + targets = compute_block_targets( + "prob_uniform", reader.nsystems, reader.system_nframes + ) + sampler = LmdbBatchSampler(reader, seed=42, block_targets=targets) + seen: set[int] = set() + for epoch in range(16): + sampler.set_epoch(epoch) + expected_batches = len(sampler) + batches = list(sampler) + self.assertEqual(len(batches), expected_batches) + indices = np.concatenate(batches) + self.assertEqual(len(indices), 600) + np.testing.assert_array_equal( + np.bincount(reader.frame_system_ids[indices]), [200, 200, 200] + ) + seen.update(indices.tolist()) + sampler.set_epoch(epoch) + self.assertEqual(list(sampler), batches) + self.assertEqual(seen, set(range(len(reader)))) + def test_compute_block_targets_equal_weight(self): result = compute_block_targets( "prob_sys_size;0:1:0.5;1:2:0.5", nsystems=2, system_nframes=[100, 100] @@ -2075,8 +2110,12 @@ def test_expand_indices_basic(self): frame_system_ids = [0] * 5 + [1] * 5 block_targets = [([0], 25), ([1], 25)] rng = np.random.default_rng(42) - expanded = _expand_indices_by_blocks( - list(range(10)), frame_system_ids, block_targets, rng + expanded = _resample_indices_by_blocks( + list(range(10)), + frame_system_ids, + system_block_lookup(block_targets), + [target for _, target in block_targets], + rng, ) sys0 = [i for i in expanded if frame_system_ids[i] == 0] sys1 = [i for i in expanded if frame_system_ids[i] == 1] @@ -2087,8 +2126,12 @@ def test_expand_indices_no_expansion(self): frame_system_ids = [0] * 5 + [1] * 5 block_targets = [([0], 5), ([1], 5)] rng = np.random.default_rng(42) - expanded = _expand_indices_by_blocks( - list(range(10)), frame_system_ids, block_targets, rng + expanded = _resample_indices_by_blocks( + list(range(10)), + frame_system_ids, + system_block_lookup(block_targets), + [target for _, target in block_targets], + rng, ) self.assertEqual(sorted(expanded), list(range(10))) @@ -2100,8 +2143,12 @@ def test_expand_indices_remainder_sampling(self): frame_system_ids = [0] * 10 block_targets = [([0], 23)] rng = np.random.default_rng(42) - expanded = _expand_indices_by_blocks( - list(range(10)), frame_system_ids, block_targets, rng + expanded = _resample_indices_by_blocks( + list(range(10)), + frame_system_ids, + system_block_lookup(block_targets), + [target for _, target in block_targets], + rng, ) self.assertEqual(len(expanded), 23) counts = Counter(expanded) @@ -2114,8 +2161,12 @@ def test_expand_epoch_diversity(self): results = [] for seed in range(5): rng = np.random.default_rng(seed) - expanded = _expand_indices_by_blocks( - list(range(10)), frame_system_ids, block_targets, rng + expanded = _resample_indices_by_blocks( + list(range(10)), + frame_system_ids, + system_block_lookup(block_targets), + [target for _, target in block_targets], + rng, ) results.append(sorted(expanded[10:])) unique = {tuple(r) for r in results} @@ -2169,6 +2220,16 @@ def get_batch_size_for_nloc(nloc): self.assertEqual(sum(map(len, batches)), 3) self.assertEqual(sorted(counts), [1, 2]) + sampler = LmdbBatchSampler( + TwoSignatureReader(), shuffle=False, seed=42, block_targets=[([0], 1)] + ) + seen: set[int] = set() + for _ in range(16): + batches = list(sampler) + self.assertEqual(sum(map(len, batches)), 1) + seen.update(index for batch in batches for index in batch) + self.assertEqual(seen, {0, 1}) + def test_sampler_tolerates_blocks_that_omit_a_system(self): """A system no auto_prob block names is passed through, not rejected. diff --git a/source/tests/common/test_loggers_training.py b/source/tests/common/test_loggers_training.py index 97a1874cfe..4620774e3a 100644 --- a/source/tests/common/test_loggers_training.py +++ b/source/tests/common/test_loggers_training.py @@ -257,7 +257,8 @@ def test_node_summary_does_not_take_checkpoint_ownership( output = capsys.readouterr().err assert "Computed num_steps=20" in output - assert "defaulting to the number of systems per task" in output + assert "defaulting to the native data size per task" in output + assert "NumPy/HDF5 systems or LMDB frames" in output assert task_schedule.model_prob.tolist() == [1 / 3, 2 / 3] assert "Resolved checkpoint retention to 2" in output assert store.max_keep == 2 diff --git a/source/tests/pt/test_lmdb_dataloader.py b/source/tests/pt/test_lmdb_dataloader.py index 148313c777..eb85481f7d 100644 --- a/source/tests/pt/test_lmdb_dataloader.py +++ b/source/tests/pt/test_lmdb_dataloader.py @@ -1075,6 +1075,31 @@ def auto_prob_lmdb(tmp_path): class TestAutoProbDataset: """Test LmdbDataset with auto_prob_style.""" + @pytest.mark.parametrize( + ("auto_prob", "expected_frames"), + [("prob_uniform", 300), ("prob_sys_size;0:1:0.5;1:2:0.5", 350)], + ) + def test_summary_reports_exact_pass_frames( + self, + auto_prob_lmdb: str, + caplog: pytest.LogCaptureFixture, + auto_prob: str, + expected_frames: int, + ) -> None: + ds = LmdbDataset( + auto_prob_lmdb, + type_map=["O", "H"], + batch_size=4, + auto_prob_style=auto_prob, + ) + try: + with caplog.at_level("INFO", logger="deepmd.pt.utils.lmdb_dataset"): + ds.print_summary("training", None) + assert f"300->{expected_frames} frames" in caplog.text + assert sum(len(batch) for batch in ds._batch_sampler) == expected_frames + finally: + ds.close() + def test_dataset_auto_prob_passthrough(self, auto_prob_lmdb): ds = LmdbDataset( auto_prob_lmdb, @@ -1732,11 +1757,11 @@ def uniform_sources(tmp_path: Path) -> tuple[list[str], list[np.ndarray]]: ("pt", 1, False, False), ("pt", "max:24", True, False), ("pt", "filter:12", False, False), - ("pt", "mix:30", True, False), + ("pt", "mix:50", True, False), ("pt_expt", 1, True, False), ("pt_expt", "filter:12", True, False), - ("pt_expt", "mix:30", False, False), - ("pt_expt", "mix:30", False, True), + ("pt_expt", "mix:50", False, False), + ("pt_expt", "mix:50", False, True), ], ) def test_uniform_groups_balance_internal_system_frames( @@ -1796,7 +1821,7 @@ def test_uniform_groups_balance_internal_system_frames( counts[:eligible] / counts.sum(), 1 / eligible, atol=0.025 ) assert counts[eligible:].sum() == 0 - if batch_size == "mix:30": + if batch_size == "mix:50": assert mixed