diff --git a/deepmd/dpmodel/common.py b/deepmd/dpmodel/common.py index 920364edc0..1f9d4817a2 100644 --- a/deepmd/dpmodel/common.py +++ b/deepmd/dpmodel/common.py @@ -43,7 +43,7 @@ } assert VALID_PRECISION.issubset(PRECISION_DICT.keys()) -RESERVED_PRECISON_DICT = { +RESERVED_PRECISION_DICT = { np.float16: "float16", np.float32: "float32", np.float64: "float64", @@ -52,7 +52,7 @@ ml_dtypes.bfloat16: "bfloat16", np.bool_: "bool", } -assert set(RESERVED_PRECISON_DICT.keys()) == set(PRECISION_DICT.values()) +assert set(RESERVED_PRECISION_DICT.keys()) == set(PRECISION_DICT.values()) DEFAULT_PRECISION = "float64" @@ -74,9 +74,9 @@ def get_xp_precision( elif precision == "bool": return bool elif precision == "default": - return get_xp_precision(xp, RESERVED_PRECISON_DICT[PRECISION_DICT[precision]]) + return get_xp_precision(xp, RESERVED_PRECISION_DICT[PRECISION_DICT[precision]]) elif precision == "global": - return get_xp_precision(xp, RESERVED_PRECISON_DICT[GLOBAL_NP_FLOAT_PRECISION]) + return get_xp_precision(xp, RESERVED_PRECISION_DICT[GLOBAL_NP_FLOAT_PRECISION]) elif precision == "bfloat16": return ml_dtypes.bfloat16 else: @@ -225,6 +225,6 @@ def safe_cast_array( "GLOBAL_ENER_FLOAT_PRECISION", "GLOBAL_NP_FLOAT_PRECISION", "PRECISION_DICT", - "RESERVED_PRECISON_DICT", + "RESERVED_PRECISION_DICT", "NativeOP", ] diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index ccad72c6a5..ec0b986394 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -14,7 +14,7 @@ GLOBAL_ENER_FLOAT_PRECISION, GLOBAL_NP_FLOAT_PRECISION, PRECISION_DICT, - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, NativeOP, ) from deepmd.dpmodel.model.base_model import ( @@ -169,7 +169,7 @@ def __init__( self.atomic_model: T_AtomicModel = T_AtomicModel(*args, **kwargs) self.precision_dict = PRECISION_DICT # not supported by flax - # self.reverse_precision_dict = RESERVED_PRECISON_DICT + # self.reverse_precision_dict = RESERVED_PRECISION_DICT self.global_np_float_precision = GLOBAL_NP_FLOAT_PRECISION self.global_ener_float_precision = GLOBAL_ENER_FLOAT_PRECISION @@ -373,7 +373,7 @@ def input_type_cast( str, ]: """Cast the input data to global float type.""" - input_prec = RESERVED_PRECISON_DICT[self.precision_dict[coord.dtype.name]] + input_prec = RESERVED_PRECISION_DICT[self.precision_dict[coord.dtype.name]] ### ### type checking would not pass jit, convert to coord prec anyway ### @@ -382,7 +382,7 @@ def input_type_cast( for vv in [box, fparam, aparam] ] box, fparam, aparam = _lst - if input_prec == RESERVED_PRECISON_DICT[self.global_np_float_precision]: + if input_prec == RESERVED_PRECISION_DICT[self.global_np_float_precision]: return coord, box, fparam, aparam, input_prec else: pp = self.global_np_float_precision @@ -401,7 +401,7 @@ def output_type_cast( ) -> dict[str, np.ndarray]: """Convert the model output to the input prec.""" do_cast = ( - input_prec != RESERVED_PRECISON_DICT[self.global_np_float_precision] + input_prec != RESERVED_PRECISION_DICT[self.global_np_float_precision] ) pp = self.precision_dict[input_prec] odef = self.model_output_def() diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index c31170ad71..d13e837161 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -35,7 +35,7 @@ from deepmd.pd.utils.env import ( DEVICE, GLOBAL_PD_FLOAT_PRECISION, - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, enable_prim, ) from deepmd.pd.utils.utils import ( @@ -355,7 +355,7 @@ def _eval_model( request_defs: list[OutputVariableDef], ): model = self.dp.to(DEVICE) - prec = NP_PRECISION_DICT[RESERVED_PRECISON_DICT[GLOBAL_PD_FLOAT_PRECISION]] + prec = NP_PRECISION_DICT[RESERVED_PRECISION_DICT[GLOBAL_PD_FLOAT_PRECISION]] nframes = coords.shape[0] if len(atom_types.shape) == 1: @@ -370,7 +370,7 @@ def _eval_model( place=DEVICE, ) type_input = paddle.to_tensor( - atom_types.astype(NP_PRECISION_DICT[RESERVED_PRECISON_DICT[paddle.int64]]), + atom_types.astype(NP_PRECISION_DICT[RESERVED_PRECISION_DICT[paddle.int64]]), dtype=paddle.int64, place=DEVICE, ) diff --git a/deepmd/pd/model/descriptor/dpa1.py b/deepmd/pd/model/descriptor/dpa1.py index f3f1ea26d6..ebaaf96bc3 100644 --- a/deepmd/pd/model/descriptor/dpa1.py +++ b/deepmd/pd/model/descriptor/dpa1.py @@ -23,7 +23,7 @@ ) from deepmd.pd.utils.env import ( PRECISION_DICT, - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, ) from deepmd.pd.utils.update_sel import ( UpdateSel, @@ -503,7 +503,7 @@ def serialize(self) -> dict: "use_tebd_bias": self.use_tebd_bias, "type_map": self.type_map, # make deterministic - "precision": RESERVED_PRECISON_DICT[obj.prec], + "precision": RESERVED_PRECISION_DICT[obj.prec], "embeddings": obj.filter_layers.serialize(), "attention_layers": obj.dpa1_attention.serialize(), "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(), diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py index 0af6d082b8..6e5832670d 100644 --- a/deepmd/pd/model/descriptor/se_a.py +++ b/deepmd/pd/model/descriptor/se_a.py @@ -23,7 +23,7 @@ ) from deepmd.pd.utils.env import ( PRECISION_DICT, - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, ) from deepmd.pd.utils.env_mat_stat import ( EnvMatStatSe, @@ -334,7 +334,7 @@ def serialize(self) -> dict: "set_davg_zero": obj.set_davg_zero, "activation_function": obj.activation_function, # make deterministic - "precision": RESERVED_PRECISON_DICT[obj.prec], + "precision": RESERVED_PRECISION_DICT[obj.prec], "embeddings": obj.filter_layers.serialize(), "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(), "exclude_types": obj.exclude_types, diff --git a/deepmd/pd/model/descriptor/se_t_tebd.py b/deepmd/pd/model/descriptor/se_t_tebd.py index a8b9a6a417..b28f79e436 100644 --- a/deepmd/pd/model/descriptor/se_t_tebd.py +++ b/deepmd/pd/model/descriptor/se_t_tebd.py @@ -30,7 +30,7 @@ ) from deepmd.pd.utils.env import ( PRECISION_DICT, - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, ) from deepmd.pd.utils.env_mat_stat import ( EnvMatStatSe, @@ -358,7 +358,7 @@ def serialize(self) -> dict: "use_econf_tebd": self.use_econf_tebd, "type_map": self.type_map, # make deterministic - "precision": RESERVED_PRECISON_DICT[obj.prec], + "precision": RESERVED_PRECISION_DICT[obj.prec], "embeddings": obj.filter_layers.serialize(), "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(), "type_embedding": self.type_embedding.embedding.serialize(), diff --git a/deepmd/pd/model/model/make_model.py b/deepmd/pd/model/model/make_model.py index 2b9a4b5bec..acb237b5ac 100644 --- a/deepmd/pd/model/model/make_model.py +++ b/deepmd/pd/model/model/make_model.py @@ -28,7 +28,7 @@ GLOBAL_PD_ENER_FLOAT_PRECISION, GLOBAL_PD_FLOAT_PRECISION, PRECISION_DICT, - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, ) from deepmd.pd.utils.nlist import ( extend_input_and_build_neighbor_list, @@ -76,7 +76,7 @@ def __init__( else: self.atomic_model: T_AtomicModel = T_AtomicModel(*args, **kwargs) self.precision_dict = PRECISION_DICT - self.reverse_precision_dict = RESERVED_PRECISON_DICT + self.reverse_precision_dict = RESERVED_PRECISION_DICT self.global_pd_float_precision = GLOBAL_PD_FLOAT_PRECISION self.global_pd_ener_float_precision = GLOBAL_PD_ENER_FLOAT_PRECISION diff --git a/deepmd/pd/utils/dataloader.py b/deepmd/pd/utils/dataloader.py index 80b3e7cb8b..221a5e776d 100644 --- a/deepmd/pd/utils/dataloader.py +++ b/deepmd/pd/utils/dataloader.py @@ -36,7 +36,7 @@ from deepmd.pd.utils.dataset import ( DeepmdDataSetForLoader, ) -from deepmd.pt.utils.utils import ( +from deepmd.pd.utils.utils import ( mix_entropy, ) from deepmd.utils import random as dp_random diff --git a/deepmd/pd/utils/env.py b/deepmd/pd/utils/env.py index 041c231282..e2abe9a6e5 100644 --- a/deepmd/pd/utils/env.py +++ b/deepmd/pd/utils/env.py @@ -55,7 +55,7 @@ PRECISION_DICT["default"] = GLOBAL_PD_FLOAT_PRECISION assert VALID_PRECISION.issubset(PRECISION_DICT.keys()) # cannot automatically generated -RESERVED_PRECISON_DICT = { +RESERVED_PRECISION_DICT = { paddle.float16: "float16", paddle.float32: "float32", paddle.float64: "float64", @@ -64,7 +64,7 @@ paddle.bfloat16: "bfloat16", paddle.bool: "bool", } -assert set(PRECISION_DICT.values()) == set(RESERVED_PRECISON_DICT.keys()) +assert set(PRECISION_DICT.values()) == set(RESERVED_PRECISION_DICT.keys()) DEFAULT_PRECISION = "float64" # throw warnings if threads not set @@ -163,7 +163,7 @@ def enable_prim(enable: bool = True): "LOCAL_RANK", "NUM_WORKERS", "PRECISION_DICT", - "RESERVED_PRECISON_DICT", + "RESERVED_PRECISION_DICT", "SAMPLER_RECORD", "enable_prim", ] diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index d22bddd97a..06126b5dab 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -58,7 +58,7 @@ from deepmd.pt.utils.env import ( DEVICE, GLOBAL_PT_FLOAT_PRECISION, - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, ) from deepmd.pt.utils.utils import ( to_numpy_array, @@ -413,7 +413,7 @@ def _eval_model( request_defs: list[OutputVariableDef], ): model = self.dp.to(DEVICE) - prec = NP_PRECISION_DICT[RESERVED_PRECISON_DICT[GLOBAL_PT_FLOAT_PRECISION]] + prec = NP_PRECISION_DICT[RESERVED_PRECISION_DICT[GLOBAL_PT_FLOAT_PRECISION]] nframes = coords.shape[0] if len(atom_types.shape) == 1: @@ -428,7 +428,7 @@ def _eval_model( device=DEVICE, ) type_input = torch.tensor( - atom_types.astype(NP_PRECISION_DICT[RESERVED_PRECISON_DICT[torch.long]]), + atom_types.astype(NP_PRECISION_DICT[RESERVED_PRECISION_DICT[torch.long]]), dtype=torch.long, device=DEVICE, ) @@ -560,7 +560,7 @@ def _eval_model_spin( np.abs(shape), np.nan, dtype=NP_PRECISION_DICT[ - RESERVED_PRECISON_DICT[GLOBAL_PT_FLOAT_PRECISION] + RESERVED_PRECISION_DICT[GLOBAL_PT_FLOAT_PRECISION] ], ) ) # this is kinda hacky diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index ba2fd1b6c6..47b2a4d15c 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -23,7 +23,7 @@ ) from deepmd.pt.utils.env import ( PRECISION_DICT, - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, ) from deepmd.pt.utils.tabulate import ( DPTabulate, @@ -505,7 +505,7 @@ def serialize(self) -> dict: "use_tebd_bias": self.use_tebd_bias, "type_map": self.type_map, # make deterministic - "precision": RESERVED_PRECISON_DICT[obj.prec], + "precision": RESERVED_PRECISION_DICT[obj.prec], "embeddings": obj.filter_layers.serialize(), "attention_layers": obj.dpa1_attention.serialize(), "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(), diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 408cd51d3d..550154300b 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -23,7 +23,7 @@ ) from deepmd.pt.utils.env import ( PRECISION_DICT, - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, ) from deepmd.pt.utils.env_mat_stat import ( EnvMatStatSe, @@ -379,7 +379,7 @@ def serialize(self) -> dict: "set_davg_zero": obj.set_davg_zero, "activation_function": obj.activation_function, # make deterministic - "precision": RESERVED_PRECISON_DICT[obj.prec], + "precision": RESERVED_PRECISION_DICT[obj.prec], "embeddings": obj.filter_layers.serialize(), "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(), "exclude_types": obj.exclude_types, diff --git a/deepmd/pt/model/descriptor/se_atten_v2.py b/deepmd/pt/model/descriptor/se_atten_v2.py index 11d783261e..533d7887e0 100644 --- a/deepmd/pt/model/descriptor/se_atten_v2.py +++ b/deepmd/pt/model/descriptor/se_atten_v2.py @@ -20,7 +20,7 @@ env, ) from deepmd.pt.utils.env import ( - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, ) from deepmd.utils.version import ( check_version_compatibility, @@ -223,7 +223,7 @@ def serialize(self) -> dict: "use_tebd_bias": self.use_tebd_bias, "type_map": self.type_map, # make deterministic - "precision": RESERVED_PRECISON_DICT[obj.prec], + "precision": RESERVED_PRECISION_DICT[obj.prec], "embeddings": obj.filter_layers.serialize(), "embeddings_strip": obj.filter_layers_strip.serialize(), "attention_layers": obj.dpa1_attention.serialize(), diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index 2cf597d015..f25fb93fa7 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -25,7 +25,7 @@ ) from deepmd.pt.utils.env import ( PRECISION_DICT, - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, ) from deepmd.pt.utils.env_mat_stat import ( EnvMatStatSe, @@ -549,7 +549,7 @@ def serialize(self) -> dict: "set_davg_zero": self.set_davg_zero, "activation_function": self.activation_function, # make deterministic - "precision": RESERVED_PRECISON_DICT[self.prec], + "precision": RESERVED_PRECISION_DICT[self.prec], "embeddings": self.filter_layers.serialize(), "env_mat": DPEnvMat(self.rcut, self.rcut_smth).serialize(), "exclude_types": self.exclude_types, diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index 90565300dc..d41a28d3db 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -23,7 +23,7 @@ ) from deepmd.pt.utils.env import ( PRECISION_DICT, - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, ) from deepmd.pt.utils.env_mat_stat import ( EnvMatStatSe, @@ -413,7 +413,7 @@ def serialize(self) -> dict: "resnet_dt": obj.resnet_dt, "set_davg_zero": obj.set_davg_zero, "activation_function": obj.activation_function, - "precision": RESERVED_PRECISON_DICT[obj.prec], + "precision": RESERVED_PRECISION_DICT[obj.prec], "embeddings": obj.filter_layers.serialize(), "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(), "exclude_types": obj.exclude_types, diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 01380a7fdf..0c7b37bb58 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -30,7 +30,7 @@ ) from deepmd.pt.utils.env import ( PRECISION_DICT, - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, ) from deepmd.pt.utils.env_mat_stat import ( EnvMatStatSe, @@ -354,7 +354,7 @@ def serialize(self) -> dict: "use_econf_tebd": self.use_econf_tebd, "type_map": self.type_map, # make deterministic - "precision": RESERVED_PRECISON_DICT[obj.prec], + "precision": RESERVED_PRECISION_DICT[obj.prec], "embeddings": obj.filter_layers.serialize(), "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(), "type_embedding": self.type_embedding.embedding.serialize(), diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 472eae5329..c32abaa095 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -28,7 +28,7 @@ GLOBAL_PT_ENER_FLOAT_PRECISION, GLOBAL_PT_FLOAT_PRECISION, PRECISION_DICT, - RESERVED_PRECISON_DICT, + RESERVED_PRECISION_DICT, ) from deepmd.pt.utils.nlist import ( extend_input_and_build_neighbor_list, @@ -76,7 +76,7 @@ def __init__( else: self.atomic_model: T_AtomicModel = T_AtomicModel(*args, **kwargs) self.precision_dict = PRECISION_DICT - self.reverse_precision_dict = RESERVED_PRECISON_DICT + self.reverse_precision_dict = RESERVED_PRECISION_DICT self.global_pt_float_precision = GLOBAL_PT_FLOAT_PRECISION self.global_pt_ener_float_precision = GLOBAL_PT_ENER_FLOAT_PRECISION diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py index 12681a304d..1c7a1884d3 100644 --- a/deepmd/pt/utils/dataloader.py +++ b/deepmd/pt/utils/dataloader.py @@ -94,26 +94,23 @@ def __init__( with h5py.File(systems) as file: systems = [os.path.join(systems, item) for item in file.keys()] - self.systems: list[DeepmdDataSetForLoader] = [] - if len(systems) >= 100: - log.info(f"Constructing DataLoaders from {len(systems)} systems") - def construct_dataset(system): return DeepmdDataSetForLoader( system=system, type_map=type_map, ) - with Pool( - os.cpu_count() - // ( - int(os.environ["LOCAL_WORLD_SIZE"]) - if dist.is_available() and dist.is_initialized() - else 1 - ) - ) as pool: - self.systems = pool.map(construct_dataset, systems) - + self.systems: list[DeepmdDataSetForLoader] = [] + global_rank = dist.get_rank() if dist.is_initialized() else 0 + if global_rank == 0: + log.info(f"Constructing DataLoaders from {len(systems)} systems") + with Pool(max(1, env.NUM_WORKERS)) as pool: + self.systems = pool.map(construct_dataset, systems) + else: + self.systems = [None] * len(systems) # type: ignore + if dist.is_initialized(): + dist.broadcast_object_list(self.systems) + assert self.systems[-1] is not None self.sampler_list: list[DistributedSampler] = [] self.index = [] self.total_batch = 0 diff --git a/deepmd/pt/utils/env.py b/deepmd/pt/utils/env.py index 6471fd80a8..9803f8d04d 100644 --- a/deepmd/pt/utils/env.py +++ b/deepmd/pt/utils/env.py @@ -54,7 +54,7 @@ PRECISION_DICT["default"] = GLOBAL_PT_FLOAT_PRECISION assert VALID_PRECISION.issubset(PRECISION_DICT.keys()) # cannot automatically generated -RESERVED_PRECISON_DICT = { +RESERVED_PRECISION_DICT = { torch.float16: "float16", torch.float32: "float32", torch.float64: "float64", @@ -63,7 +63,7 @@ torch.bfloat16: "bfloat16", torch.bool: "bool", } -assert set(PRECISION_DICT.values()) == set(RESERVED_PRECISON_DICT.keys()) +assert set(PRECISION_DICT.values()) == set(RESERVED_PRECISION_DICT.keys()) DEFAULT_PRECISION = "float64" # throw warnings if threads not set @@ -87,6 +87,6 @@ "LOCAL_RANK", "NUM_WORKERS", "PRECISION_DICT", - "RESERVED_PRECISON_DICT", + "RESERVED_PRECISION_DICT", "SAMPLER_RECORD", ] diff --git a/doc/model/dpa2.md b/doc/model/dpa2.md index 300876bf05..a041547a14 100644 --- a/doc/model/dpa2.md +++ b/doc/model/dpa2.md @@ -4,7 +4,7 @@ **Supported backends**: PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} ::: -The DPA-2 model implementation. See https://doi.org/10.1038/s41524-024-01493-2 for more details. +The DPA-2 model implementation. See [DPA-2 paper](https://doi.org/10.1038/s41524-024-01493-2) for more details. Training example: `examples/water/dpa2/input_torch_medium.json`, see [README](../../examples/water/dpa2/README.md) for inputs in different levels. diff --git a/doc/model/train-se-atten.md b/doc/model/train-se-atten.md index 92a56395f6..504b214737 100644 --- a/doc/model/train-se-atten.md +++ b/doc/model/train-se-atten.md @@ -8,7 +8,7 @@ Here we propose DPA-1, a Deep Potential model with a novel attention mechanism, which is highly effective for representing the conformation and chemical spaces of atomic systems and learning the PES. -See [this paper](https://arxiv.org/abs/2208.08236) for more information. DPA-1 is implemented as a new descriptor `"se_atten"` for model training, which can be used after simply editing the input.json. +See [this paper](https://www.nature.com/articles/s41524-024-01278-7) for more information. DPA-1 is implemented as a new descriptor `"se_atten"` for model training, which can be used after simply editing the input.json. ## Theory @@ -71,7 +71,7 @@ Then layer normalization is added in a residual way to finally obtain the self-a Next, we will list the detailed settings in input.json and the data format, especially for large systems with dozens of elements. An example of DPA-1 input can be found in `examples/water/se_atten/input.json`. The notation of `se_atten` is short for the smooth edition of Deep Potential with an attention mechanism. -This descriptor was described in detail in [the DPA-1 paper](https://arxiv.org/abs/2208.08236) and the images above. +This descriptor was described in detail in [the DPA-1 paper](https://www.nature.com/articles/s41524-024-01278-7) and the images above. In this example, we will train a DPA-1 model for a water system. A complete training input script of this example can be found in the directory: diff --git a/doc/train/finetuning.md b/doc/train/finetuning.md index 04d86cfc98..beb6012003 100644 --- a/doc/train/finetuning.md +++ b/doc/train/finetuning.md @@ -9,7 +9,7 @@ to vastly reduce the training cost, while it's not trivial in potential models. Compositions and configurations of data samples or even computational parameters in upstream software (such as VASP) may be different between the pre-trained and target datasets, leading to energy shifts or other diversities of training data. -Recently the emerging of methods such as [DPA-1](https://arxiv.org/abs/2208.08236) has brought us to a new stage where we can +Recently the emerging of methods such as [DPA-1](https://www.nature.com/articles/s41524-024-01278-7) has brought us to a new stage where we can perform similar pretraining-finetuning approaches. They can hopefully learn the common knowledge in the pre-trained dataset (especially the `force` information) and thus reduce the computational cost in downstream training tasks. @@ -19,7 +19,7 @@ and thus reduce the computational cost in downstream training tasks. If you have a pre-trained model `pretrained.pb` (here we support models using [`se_atten`](../model/train-se-atten.md) descriptor and [`ener`](../model/train-energy.md) fitting net) on a large dataset (for example, [OC2M](https://github.com/Open-Catalyst-Project/ocp/blob/main/DATASET.md) in -DPA-1 [paper](https://arxiv.org/abs/2208.08236)), a finetuning strategy can be performed by simply running: +DPA-1 [paper](https://www.nature.com/articles/s41524-024-01278-7)), a finetuning strategy can be performed by simply running: ```bash $ dp train input.json --finetune pretrained.pb diff --git a/doc/train/parallel-training.md b/doc/train/parallel-training.md index 9ea92b4751..00df0a63f0 100644 --- a/doc/train/parallel-training.md +++ b/doc/train/parallel-training.md @@ -27,13 +27,14 @@ In some cases, it won't work well when scaling the learning rate by worker count ### Scaling test Testing `examples/water/se_e2_a` on an 8-GPU host, linear acceleration can be observed with the increasing number of cards. - -| Num of GPU cards | Seconds every 100 samples | Samples per second | Speed up | -| ---------------- | ------------------------- | ------------------ | -------- | -| 1 | 1.4515 | 68.89 | 1.00 | -| 2 | 1.5962 | 62.65\*2 | 1.82 | -| 4 | 1.7635 | 56.71\*4 | 3.29 | -| 8 | 1.7267 | 57.91\*8 | 6.72 | +In this example, the number of samples per batch on a single GPU card ({ref}`batch_size `) is set to `1`. + +| Num of GPU cards | Samples per batch | Seconds every 100 batches | Samples per second | Speed up | +| ---------------- | ----------------- | ------------------------- | ------------------ | -------- | +| 1 | 1 | 1.4515 | 68.89 | 1.00 | +| 2 | 2 | 1.5962 | 62.65\*2 | 1.82 | +| 4 | 4 | 1.7635 | 56.71\*4 | 3.29 | +| 8 | 8 | 1.7267 | 57.91\*8 | 6.72 | ### How to use