Skip to content

Commit 565f4be

Browse files
authored
feat(pt/dpmodel): add lmdb dataloader (#5283)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Full LMDB data support across training, validation and testing with PyTorch-friendly dataset/dataloader integration, per-atom-count (nloc) grouping, and deterministic distributed batching. * **Tests** * Extensive unit and integration tests covering LMDB reader, dataset/dataloader, samplers, collation, and mixed-/uniform-nloc scenarios. * **Documentation** * Added an example LMDB training configuration and updated test runner to handle LMDB-grouped test runs. * **Chores** * Added runtime dependencies for LMDB and msgpack support. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent 52792f0 commit 565f4be

18 files changed

Lines changed: 4212 additions & 111 deletions

File tree

deepmd/dpmodel/utils/__init__.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,14 @@
66
AtomExcludeMask,
77
PairExcludeMask,
88
)
9+
from .lmdb_data import (
10+
DistributedSameNlocBatchSampler,
11+
LmdbDataReader,
12+
LmdbTestData,
13+
SameNlocBatchSampler,
14+
is_lmdb,
15+
make_neighbor_stat_data,
16+
)
917
from .network import (
1018
EmbeddingNet,
1119
FittingNet,
@@ -44,13 +52,17 @@
4452

4553
__all__ = [
4654
"AtomExcludeMask",
55+
"DistributedSameNlocBatchSampler",
4756
"EmbeddingNet",
4857
"EnvMat",
4958
"FittingNet",
59+
"LmdbDataReader",
60+
"LmdbTestData",
5061
"NativeLayer",
5162
"NativeNet",
5263
"NetworkCollection",
5364
"PairExcludeMask",
65+
"SameNlocBatchSampler",
5466
"aggregate",
5567
"build_multiple_neighbor_list",
5668
"build_neighbor_list",
@@ -59,10 +71,12 @@
5971
"get_graph_index",
6072
"get_multiple_nlist_key",
6173
"inter2phys",
74+
"is_lmdb",
6275
"load_dp_model",
6376
"make_embedding_network",
6477
"make_fitting_network",
6578
"make_multilayer_network",
79+
"make_neighbor_stat_data",
6680
"nlist_distinguish_types",
6781
"normalize_coord",
6882
"phys2inter",

0 commit comments

Comments
 (0)