From 0f91fe0478e5feff1610e9e640a11f75a7749e14 Mon Sep 17 00:00:00 2001 From: Chengqian Zhang <2000011006@stu.pku.edu.cn> Date: Wed, 27 Aug 2025 07:46:08 +0000 Subject: [PATCH 01/13] Init branch --- .../model/atomic_model/base_atomic_model.py | 20 +++++++++++++++++++ .../pt/model/atomic_model/dp_atomic_model.py | 4 +--- deepmd/pt/model/model/make_model.py | 2 ++ 3 files changed, 23 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index a2cbef3eee..6bd6af7f9c 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -487,6 +487,26 @@ def change_out_bias( else: raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) + def compute_fitting_stat( + self, + sample_merged + ) -> None: + """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.. + + Parameters + ---------- + sample_merged : Union[Callable[[], list[dict]], list[dict]] + - list[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + """ + self.fitting_net.compute_input_stats( + sample_merged, protection=self.data_stat_protect + ) + def _get_forward_wrapper_func(self) -> Callable[..., torch.Tensor]: """Get a forward wrapper of the atomic model for output bias calculation.""" diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 62c7d78d75..c988501b8a 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -324,9 +324,7 @@ def wrapped_sampler(): return sampled self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path) - self.fitting_net.compute_input_stats( - wrapped_sampler, protection=self.data_stat_protect - ) + self.compute_fitting_stat(wrapped_sampler) if compute_or_load_out_stat: self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index b9335df747..3773bc12bb 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -230,6 +230,8 @@ def change_out_bias( merged, bias_adjust_mode=bias_adjust_mode, ) + if bias_adjust_mode == "set-by-statistic": + self.atomic_model.compute_fitting_stat(merged) def forward_common_lower( self, From 316a6afba90b92705e8812db70d86613c59dc04e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 27 Aug 2025 08:21:03 +0000 Subject: [PATCH 02/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/atomic_model/base_atomic_model.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index 6bd6af7f9c..b697957a5b 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -487,10 +487,7 @@ def change_out_bias( else: raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) - def compute_fitting_stat( - self, - sample_merged - ) -> None: + def compute_fitting_stat(self, sample_merged) -> None: """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.. Parameters From 2fbe51ce6997609a34eccbcebd5f1bca98dc9fa2 Mon Sep 17 00:00:00 2001 From: Chengqian Zhang <2000011006@stu.pku.edu.cn> Date: Fri, 7 Nov 2025 08:16:15 +0000 Subject: [PATCH 03/13] Move compute_fitting_stat to dp_atomic_model --- .../pt/model/atomic_model/base_atomic_model.py | 6 ++---- .../pt/model/atomic_model/dp_atomic_model.py | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index 4cbfbc6ff4..de37043e05 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -493,7 +493,7 @@ def change_out_bias( else: raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) - def compute_fitting_stat(self, sample_merged) -> None: + def compute_fitting_stat(self, sample_merged: Union[Callable[[], list[dict]], list[dict]],) -> None: """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.. Parameters @@ -506,9 +506,7 @@ def compute_fitting_stat(self, sample_merged) -> None: only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once. """ - self.fitting_net.compute_input_stats( - sample_merged, protection=self.data_stat_protect - ) + pass def _get_forward_wrapper_func(self) -> Callable[..., torch.Tensor]: """Get a forward wrapper of the atomic model for output bias calculation.""" diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index ca611a7a42..cf9f599332 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -3,6 +3,7 @@ import logging from typing import ( Any, + Union, Callable, Optional, ) @@ -332,6 +333,23 @@ def wrapped_sampler() -> list[dict]: if compute_or_load_out_stat: self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) + def compute_fitting_stat(self, sample_merged: Union[Callable[[], list[dict]], list[dict]],) -> None: + """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.. + + Parameters + ---------- + sample_merged : Union[Callable[[], list[dict]], list[dict]] + - list[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + """ + self.fitting_net.compute_input_stats( + sample_merged, protection=self.data_stat_protect + ) + def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.fitting_net.get_dim_fparam() From 939558f21a6c44632982f8c12ea324c1b7605d08 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 7 Nov 2025 08:18:01 +0000 Subject: [PATCH 04/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/atomic_model/base_atomic_model.py | 5 ++++- deepmd/pt/model/atomic_model/dp_atomic_model.py | 7 +++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index de37043e05..9e64f071c0 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -493,7 +493,10 @@ def change_out_bias( else: raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) - def compute_fitting_stat(self, sample_merged: Union[Callable[[], list[dict]], list[dict]],) -> None: + def compute_fitting_stat( + self, + sample_merged: Union[Callable[[], list[dict]], list[dict]], + ) -> None: """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.. Parameters diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index cf9f599332..425fd905a0 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -3,9 +3,9 @@ import logging from typing import ( Any, - Union, Callable, Optional, + Union, ) import torch @@ -333,7 +333,10 @@ def wrapped_sampler() -> list[dict]: if compute_or_load_out_stat: self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) - def compute_fitting_stat(self, sample_merged: Union[Callable[[], list[dict]], list[dict]],) -> None: + def compute_fitting_stat( + self, + sample_merged: Union[Callable[[], list[dict]], list[dict]], + ) -> None: """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.. Parameters From 003233fee133f451f3f13302e1b15d71b2475890 Mon Sep 17 00:00:00 2001 From: Chengqian Zhang <2000011006@stu.pku.edu.cn> Date: Fri, 7 Nov 2025 08:42:46 +0000 Subject: [PATCH 05/13] Change UT to detech this error --- .../model/water/data/data_0/set.000/fparam.npy | Bin 1408 -> 0 bytes source/tests/pt/test_training.py | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) delete mode 100644 source/tests/pt/model/water/data/data_0/set.000/fparam.npy diff --git a/source/tests/pt/model/water/data/data_0/set.000/fparam.npy b/source/tests/pt/model/water/data/data_0/set.000/fparam.npy deleted file mode 100644 index 770b4a5d664c5d3dbc0a803f11c407f1d4c8c26c..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1408 zcmbW1={pn%0EbC-Ct;^H+O=V9YQ$ub%9YXOuDvX?jYR1n#OL^ks= zB2!3YpMap?AP;}{fFLjL|GbsQIjZ-vr}}vWdM~qb8*Sa9dlhLbsB zJaYRWQ@2YD{bvapK3@HZev_vbkzR-S%W?hHG956Db8L!N8iiBZUaLYoE+~q`DT3h} z%v{VW-pOM@Y&z8NwsHtha;$ARg%jv){$XoU6%Wkm!D;&A7#I-0o@fOJHZCWS8)I2=&p!VSx`v5y%;8(xP?Y zWXO7Ia%3UCct)Svt~r8Y;~0%*QWv;6864ZjI=F6?x#4%X7A`i?0i|pns@)HD?i%C5 zzg+9AzA6`;FI7qp7>?pTy=&c3nHGfF+M5S^ir^4u>|3=s4%2{hi9VnB=qFuXxu&lR z^KQC&!Ur8tQ+L_s7%j%zOrt)(nkM+2pSnx2<-Z^-9rtb_262PqbNsfFlRDbT1*%~T z(O1Qq8wc^w99w&Vhn#rDzas5~*kJ$JS)SGlc1Vn7FQFBRX}@0gjA+5Q&Al6$W;`U$ z%atW2H9@~}>o0ppM^Knzw%=}G9Q!tSDp4vWn6Wc_G07f6toa{`(uquD4LZriw6{WP z1>ehYR|rgGAG)5%6Jg3aOkv_t4~k0ocRMU=kkzdlobibbUX0a7+2ziy)=k8MZz8A_ z9Mj3ZB7xMsTJO?S4+wu|Ng}*Da6_D}9F{+bMA?GLr?pJb?fDi8zSVf8Fj2v~H2}lp z!imV=1qd1z7;v|6a5YOkEzMPoI_+U+x&568ZN3zn{9?IF!m0Na+eL63+Z}uRycm;tsE4NWqS$4i0d)7X;sdcCV2~|e zop!Ai@=J-2XFS`$3|qXQv`LJ>7X1QdLN#3Eb7d47g?Pt^BD~5Mg19oUxQdqtn)=E$ zleb#IwJ-5yNwwooMrx*5jtAmndUBqM2sHPP=g77^*qltaDqd>9XZgJPW@QPkjg^eA zD$9T&RprNTgdU{y5&hq3*CDs`>124^D6~db@0hA+gcj-a9x7u1%AJngw&{E*pK+Kz z!5qZ9wK66{_can`Gak?!+j+1mtgU!6w}&>qG>)kj4# zR|NV$q4QSdZP-06pHp#q5KgMR)P$#Og!&~km)6jsn9IL8`<9E@3zxX8&>lz`>sZM8 jv|)~+=@`#00Y`i!mPQ(ei}ja;na#aeS1$0Wwrs?|p%br& diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index da239212b0..37fc030ed0 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -92,7 +92,7 @@ def test_dp_train(self) -> None: state_dict_trained[state_key], state_dict_finetuned_empty[state_key], ) - if "fitting_net" not in state_key: + if ("fitting_net" not in state_key) or ("fparam" in state_key) or ("aparam" in state_key): torch.testing.assert_close( state_dict_trained[state_key], state_dict_finetuned_random[state_key], From 4241d51162dabf6247ad5a345514bd34d77269bf Mon Sep 17 00:00:00 2001 From: Chengqian Zhang <2000011006@stu.pku.edu.cn> Date: Fri, 7 Nov 2025 09:27:35 +0000 Subject: [PATCH 06/13] Add pd support --- .../model/atomic_model/base_atomic_model.py | 18 +++++++++++++++++ .../pd/model/atomic_model/dp_atomic_model.py | 20 +++++++++++++++++++ deepmd/pd/model/model/make_model.py | 2 ++ source/tests/pd/test_training.py | 3 ++- source/tests/pt/test_training.py | 1 + 5 files changed, 43 insertions(+), 1 deletion(-) diff --git a/deepmd/pd/model/atomic_model/base_atomic_model.py b/deepmd/pd/model/atomic_model/base_atomic_model.py index 4f40117fb7..fa76021484 100644 --- a/deepmd/pd/model/atomic_model/base_atomic_model.py +++ b/deepmd/pd/model/atomic_model/base_atomic_model.py @@ -515,6 +515,24 @@ def change_out_bias( else: raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) + def compute_fitting_stat( + self, + sample_merged: Union[Callable[[], list[dict]], list[dict]], + ) -> None: + """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.. + + Parameters + ---------- + sample_merged : Union[Callable[[], list[dict]], list[dict]] + - list[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + """ + pass + def _get_forward_wrapper_func(self) -> Callable[..., paddle.Tensor]: """Get a forward wrapper of the atomic model for output bias calculation.""" diff --git a/deepmd/pd/model/atomic_model/dp_atomic_model.py b/deepmd/pd/model/atomic_model/dp_atomic_model.py index 816245c28a..062c5b967e 100644 --- a/deepmd/pd/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pd/model/atomic_model/dp_atomic_model.py @@ -403,6 +403,26 @@ def wrapped_sampler(): if compute_or_load_out_stat: self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) + def compute_fitting_stat( + self, + sample_merged: Union[Callable[[], list[dict]], list[dict]], + ) -> None: + """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.. + + Parameters + ---------- + sample_merged : Union[Callable[[], list[dict]], list[dict]] + - list[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + """ + self.fitting_net.compute_input_stats( + sample_merged, protection=self.data_stat_protect + ) + def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.fitting_net.get_dim_fparam() diff --git a/deepmd/pd/model/model/make_model.py b/deepmd/pd/model/model/make_model.py index 42c406f8d7..bc57113f4d 100644 --- a/deepmd/pd/model/model/make_model.py +++ b/deepmd/pd/model/model/make_model.py @@ -228,6 +228,8 @@ def change_out_bias( merged, bias_adjust_mode=bias_adjust_mode, ) + if bias_adjust_mode == "set-by-statistic": + self.atomic_model.compute_fitting_stat(merged) def forward_common_lower( self, diff --git a/source/tests/pd/test_training.py b/source/tests/pd/test_training.py index 0dc36fa314..b71829db61 100644 --- a/source/tests/pd/test_training.py +++ b/source/tests/pd/test_training.py @@ -89,7 +89,7 @@ def test_dp_train(self) -> None: state_dict_trained[state_key].numpy(), state_dict_finetuned_empty[state_key].numpy(), ) - if "fitting_net" not in state_key: + if ("fitting_net" not in state_key) or ("fparam" in state_key) or ("aparam" in state_key): np.testing.assert_allclose( state_dict_trained[state_key].numpy(), state_dict_finetuned_random[state_key].numpy(), @@ -190,6 +190,7 @@ def setUp(self) -> None: self.config["training"]["save_freq"] = 1 self.set_path = Path(__file__).parent / "water/data/data_0" / "set.000" shutil.copyfile(self.set_path / "energy.npy", self.set_path / "fparam.npy") + self.config["model"]["data_stat_nbatch"] = 100 def tearDown(self) -> None: (self.set_path / "fparam.npy").unlink(missing_ok=True) diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index 37fc030ed0..df96508464 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -256,6 +256,7 @@ def setUp(self) -> None: self.config["training"]["save_freq"] = 1 self.set_path = Path(__file__).parent / "water/data/data_0" / "set.000" shutil.copyfile(self.set_path / "energy.npy", self.set_path / "fparam.npy") + self.config["model"]["data_stat_nbatch"] = 100 def tearDown(self) -> None: (self.set_path / "fparam.npy").unlink(missing_ok=True) From 12b0136999eee38605ceb7015b91347738d1fb98 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 7 Nov 2025 09:30:15 +0000 Subject: [PATCH 07/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/pd/test_training.py | 6 +++++- source/tests/pt/test_training.py | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/source/tests/pd/test_training.py b/source/tests/pd/test_training.py index b71829db61..f3d7860881 100644 --- a/source/tests/pd/test_training.py +++ b/source/tests/pd/test_training.py @@ -89,7 +89,11 @@ def test_dp_train(self) -> None: state_dict_trained[state_key].numpy(), state_dict_finetuned_empty[state_key].numpy(), ) - if ("fitting_net" not in state_key) or ("fparam" in state_key) or ("aparam" in state_key): + if ( + ("fitting_net" not in state_key) + or ("fparam" in state_key) + or ("aparam" in state_key) + ): np.testing.assert_allclose( state_dict_trained[state_key].numpy(), state_dict_finetuned_random[state_key].numpy(), diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index df96508464..ff4f00f912 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -92,7 +92,11 @@ def test_dp_train(self) -> None: state_dict_trained[state_key], state_dict_finetuned_empty[state_key], ) - if ("fitting_net" not in state_key) or ("fparam" in state_key) or ("aparam" in state_key): + if ( + ("fitting_net" not in state_key) + or ("fparam" in state_key) + or ("aparam" in state_key) + ): torch.testing.assert_close( state_dict_trained[state_key], state_dict_finetuned_random[state_key], From 59c50c8d498612298ff88d5b919150619e521926 Mon Sep 17 00:00:00 2001 From: Chengqian Zhang <2000011006@stu.pku.edu.cn> Date: Fri, 7 Nov 2025 09:31:38 +0000 Subject: [PATCH 08/13] Recover fparam.npy --- .../model/water/data/data_0/set.000/fparam.npy | Bin 0 -> 1408 bytes 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 source/tests/pt/model/water/data/data_0/set.000/fparam.npy diff --git a/source/tests/pt/model/water/data/data_0/set.000/fparam.npy b/source/tests/pt/model/water/data/data_0/set.000/fparam.npy new file mode 100644 index 0000000000000000000000000000000000000000..770b4a5d664c5d3dbc0a803f11c407f1d4c8c26c GIT binary patch literal 1408 zcmbW1={pn%0EbC-Ct;^H+O=V9YQ$ub%9YXOuDvX?jYR1n#OL^ks= zB2!3YpMap?AP;}{fFLjL|GbsQIjZ-vr}}vWdM~qb8*Sa9dlhLbsB zJaYRWQ@2YD{bvapK3@HZev_vbkzR-S%W?hHG956Db8L!N8iiBZUaLYoE+~q`DT3h} z%v{VW-pOM@Y&z8NwsHtha;$ARg%jv){$XoU6%Wkm!D;&A7#I-0o@fOJHZCWS8)I2=&p!VSx`v5y%;8(xP?Y zWXO7Ia%3UCct)Svt~r8Y;~0%*QWv;6864ZjI=F6?x#4%X7A`i?0i|pns@)HD?i%C5 zzg+9AzA6`;FI7qp7>?pTy=&c3nHGfF+M5S^ir^4u>|3=s4%2{hi9VnB=qFuXxu&lR z^KQC&!Ur8tQ+L_s7%j%zOrt)(nkM+2pSnx2<-Z^-9rtb_262PqbNsfFlRDbT1*%~T z(O1Qq8wc^w99w&Vhn#rDzas5~*kJ$JS)SGlc1Vn7FQFBRX}@0gjA+5Q&Al6$W;`U$ z%atW2H9@~}>o0ppM^Knzw%=}G9Q!tSDp4vWn6Wc_G07f6toa{`(uquD4LZriw6{WP z1>ehYR|rgGAG)5%6Jg3aOkv_t4~k0ocRMU=kkzdlobibbUX0a7+2ziy)=k8MZz8A_ z9Mj3ZB7xMsTJO?S4+wu|Ng}*Da6_D}9F{+bMA?GLr?pJb?fDi8zSVf8Fj2v~H2}lp z!imV=1qd1z7;v|6a5YOkEzMPoI_+U+x&568ZN3zn{9?IF!m0Na+eL63+Z}uRycm;tsE4NWqS$4i0d)7X;sdcCV2~|e zop!Ai@=J-2XFS`$3|qXQv`LJ>7X1QdLN#3Eb7d47g?Pt^BD~5Mg19oUxQdqtn)=E$ zleb#IwJ-5yNwwooMrx*5jtAmndUBqM2sHPP=g77^*qltaDqd>9XZgJPW@QPkjg^eA zD$9T&RprNTgdU{y5&hq3*CDs`>124^D6~db@0hA+gcj-a9x7u1%AJngw&{E*pK+Kz z!5qZ9wK66{_can`Gak?!+j+1mtgU!6w}&>qG>)kj4# zR|NV$q4QSdZP-06pHp#q5KgMR)P$#Og!&~km)6jsn9IL8`<9E@3zxX8&>lz`>sZM8 jv|)~+=@`#00Y`i!mPQ(ei}ja;na#aeS1$0Wwrs?|p%br& literal 0 HcmV?d00001 From 646977c5f8a28d3f7c2d165c6771e309e32a9b86 Mon Sep 17 00:00:00 2001 From: Chengqian Zhang <2000011006@stu.pku.edu.cn> Date: Fri, 7 Nov 2025 09:42:04 +0000 Subject: [PATCH 09/13] Update doc --- deepmd/pd/model/atomic_model/base_atomic_model.py | 2 +- deepmd/pt/model/atomic_model/dp_atomic_model.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pd/model/atomic_model/base_atomic_model.py b/deepmd/pd/model/atomic_model/base_atomic_model.py index fa76021484..0a49f6d532 100644 --- a/deepmd/pd/model/atomic_model/base_atomic_model.py +++ b/deepmd/pd/model/atomic_model/base_atomic_model.py @@ -525,7 +525,7 @@ def compute_fitting_stat( ---------- sample_merged : Union[Callable[[], list[dict]], list[dict]] - list[dict]: A list of data samples from various data systems. - Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + Each element, `merged[i]`, is a data dictionary containing `keys`: `paddle.Tensor` originating from the `i`-th data system. - Callable[[], list[dict]]: A lazy function that returns data samples in the above format only when needed. Since the sampling process can be slow and memory-intensive, diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 425fd905a0..a9364252d5 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -343,7 +343,7 @@ def compute_fitting_stat( ---------- sample_merged : Union[Callable[[], list[dict]], list[dict]] - list[dict]: A list of data samples from various data systems. - Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + Each element, `merged[i]`, is a data dictionary containing `keys`: `paddle.Tensor` originating from the `i`-th data system. - Callable[[], list[dict]]: A lazy function that returns data samples in the above format only when needed. Since the sampling process can be slow and memory-intensive, From 52d29ef82ed9d8844ac1c52f14b1271aff51bad8 Mon Sep 17 00:00:00 2001 From: Chengqian Zhang <2000011006@stu.pku.edu.cn> Date: Fri, 7 Nov 2025 09:57:11 +0000 Subject: [PATCH 10/13] Fix typo --- deepmd/pd/model/atomic_model/base_atomic_model.py | 2 +- deepmd/pd/model/atomic_model/dp_atomic_model.py | 4 ++-- deepmd/pt/model/atomic_model/base_atomic_model.py | 2 +- deepmd/pt/model/atomic_model/dp_atomic_model.py | 4 ++-- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/deepmd/pd/model/atomic_model/base_atomic_model.py b/deepmd/pd/model/atomic_model/base_atomic_model.py index 0a49f6d532..a66b075498 100644 --- a/deepmd/pd/model/atomic_model/base_atomic_model.py +++ b/deepmd/pd/model/atomic_model/base_atomic_model.py @@ -519,7 +519,7 @@ def compute_fitting_stat( self, sample_merged: Union[Callable[[], list[dict]], list[dict]], ) -> None: - """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.. + """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. Parameters ---------- diff --git a/deepmd/pd/model/atomic_model/dp_atomic_model.py b/deepmd/pd/model/atomic_model/dp_atomic_model.py index 062c5b967e..c09abd2221 100644 --- a/deepmd/pd/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pd/model/atomic_model/dp_atomic_model.py @@ -407,13 +407,13 @@ def compute_fitting_stat( self, sample_merged: Union[Callable[[], list[dict]], list[dict]], ) -> None: - """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.. + """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. Parameters ---------- sample_merged : Union[Callable[[], list[dict]], list[dict]] - list[dict]: A list of data samples from various data systems. - Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + Each element, `merged[i]`, is a data dictionary containing `keys`: `paddle.Tensor` originating from the `i`-th data system. - Callable[[], list[dict]]: A lazy function that returns data samples in the above format only when needed. Since the sampling process can be slow and memory-intensive, diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index 9e64f071c0..6377a1c3db 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -497,7 +497,7 @@ def compute_fitting_stat( self, sample_merged: Union[Callable[[], list[dict]], list[dict]], ) -> None: - """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.. + """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. Parameters ---------- diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index a9364252d5..fea7779d91 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -337,13 +337,13 @@ def compute_fitting_stat( self, sample_merged: Union[Callable[[], list[dict]], list[dict]], ) -> None: - """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.. + """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. Parameters ---------- sample_merged : Union[Callable[[], list[dict]], list[dict]] - list[dict]: A list of data samples from various data systems. - Each element, `merged[i]`, is a data dictionary containing `keys`: `paddle.Tensor` + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` originating from the `i`-th data system. - Callable[[], list[dict]]: A lazy function that returns data samples in the above format only when needed. Since the sampling process can be slow and memory-intensive, From 3af617c7652a6fe5c9dbaa1b3e847ac46da48c97 Mon Sep 17 00:00:00 2001 From: Chengqian Zhang <2000011006@stu.pku.edu.cn> Date: Tue, 18 Nov 2025 15:26:51 +0000 Subject: [PATCH 11/13] optimize typo --- deepmd/pd/model/atomic_model/base_atomic_model.py | 2 +- deepmd/pd/model/atomic_model/dp_atomic_model.py | 6 ++---- deepmd/pd/model/model/make_model.py | 2 +- deepmd/pt/model/atomic_model/base_atomic_model.py | 4 ++-- deepmd/pt/model/atomic_model/dp_atomic_model.py | 4 ++-- deepmd/pt/model/model/make_model.py | 2 +- 6 files changed, 9 insertions(+), 11 deletions(-) diff --git a/deepmd/pd/model/atomic_model/base_atomic_model.py b/deepmd/pd/model/atomic_model/base_atomic_model.py index a66b075498..65e8097ee2 100644 --- a/deepmd/pd/model/atomic_model/base_atomic_model.py +++ b/deepmd/pd/model/atomic_model/base_atomic_model.py @@ -515,7 +515,7 @@ def change_out_bias( else: raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) - def compute_fitting_stat( + def compute_fitting_input_stat( self, sample_merged: Union[Callable[[], list[dict]], list[dict]], ) -> None: diff --git a/deepmd/pd/model/atomic_model/dp_atomic_model.py b/deepmd/pd/model/atomic_model/dp_atomic_model.py index c09abd2221..a34bd2d08b 100644 --- a/deepmd/pd/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pd/model/atomic_model/dp_atomic_model.py @@ -397,13 +397,11 @@ def wrapped_sampler(): return sampled self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path) - self.fitting_net.compute_input_stats( - wrapped_sampler, protection=self.data_stat_protect - ) + self.compute_fitting_input_stat(wrapped_sampler) if compute_or_load_out_stat: self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) - def compute_fitting_stat( + def compute_fitting_input_stat( self, sample_merged: Union[Callable[[], list[dict]], list[dict]], ) -> None: diff --git a/deepmd/pd/model/model/make_model.py b/deepmd/pd/model/model/make_model.py index bc57113f4d..dafb15e5cd 100644 --- a/deepmd/pd/model/model/make_model.py +++ b/deepmd/pd/model/model/make_model.py @@ -229,7 +229,7 @@ def change_out_bias( bias_adjust_mode=bias_adjust_mode, ) if bias_adjust_mode == "set-by-statistic": - self.atomic_model.compute_fitting_stat(merged) + self.atomic_model.compute_fitting_input_stat(merged) def forward_common_lower( self, diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index 6377a1c3db..6b786a796f 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -493,11 +493,11 @@ def change_out_bias( else: raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) - def compute_fitting_stat( + def compute_fitting_input_stat( self, sample_merged: Union[Callable[[], list[dict]], list[dict]], ) -> None: - """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. + """Compute the input statistics (e.g. mean and stddev) for the atomic model from packed data. Parameters ---------- diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index fea7779d91..d9ff561335 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -329,11 +329,11 @@ def wrapped_sampler() -> list[dict]: return sampled self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path) - self.compute_fitting_stat(wrapped_sampler) + self.compute_fitting_input_stat(wrapped_sampler) if compute_or_load_out_stat: self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) - def compute_fitting_stat( + def compute_fitting_input_stat( self, sample_merged: Union[Callable[[], list[dict]], list[dict]], ) -> None: diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index e18f5e90bf..9df8bd393d 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -233,7 +233,7 @@ def change_out_bias( bias_adjust_mode=bias_adjust_mode, ) if bias_adjust_mode == "set-by-statistic": - self.atomic_model.compute_fitting_stat(merged) + self.atomic_model.compute_fitting_input_stat(merged) def forward_common_lower( self, From c6af0e00fe9b633c306e5ddce7b1de22afd08c0b Mon Sep 17 00:00:00 2001 From: Chengqian Zhang <2000011006@stu.pku.edu.cn> Date: Tue, 18 Nov 2025 16:16:12 +0000 Subject: [PATCH 12/13] Add fitting stat in python backend --- deepmd/dpmodel/fitting/general_fitting.py | 66 +++++++++++++ .../model/atomic_model/base_atomic_model.py | 2 +- .../tests/common/dpmodel/test_fitting_stat.py | 99 +++++++++++++++++++ 3 files changed, 166 insertions(+), 1 deletion(-) create mode 100644 source/tests/common/dpmodel/test_fitting_stat.py diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index a380717927..10a746cbcb 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -4,6 +4,7 @@ ) from typing import ( Any, + Callable, Optional, Union, ) @@ -221,6 +222,71 @@ def __init__( ], ) + def compute_input_stats( + self, + merged: Union[Callable[[], list[dict]], list[dict]], + protection: float = 1e-2, + ) -> None: + """ + Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. + + Parameters + ---------- + merged : Union[Callable[[], list[dict]], list[dict]] + - list[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `numpy.ndarray` + originating from the `i`-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + protection : float + Divided-by-zero protection + """ + if self.numb_fparam == 0 and self.numb_aparam == 0: + # skip data statistics + return + if callable(merged): + sampled = merged() + else: + sampled = merged + # stat fparam + if self.numb_fparam > 0: + cat_data = np.concatenate([frame["fparam"] for frame in sampled], axis=0) + cat_data = np.reshape(cat_data, [-1, self.numb_fparam]) + fparam_avg = np.mean(cat_data, axis=0) + fparam_std = np.std(cat_data, axis=0, ddof=0) # ddof=0 for population std + fparam_std = np.where( + fparam_std < protection, + np.array(protection, dtype=fparam_std.dtype), + fparam_std, + ) + fparam_inv_std = 1.0 / fparam_std + self.fparam_avg = fparam_avg.astype(self.fparam_avg.dtype) + self.fparam_inv_std = fparam_inv_std.astype(self.fparam_inv_std.dtype) + # stat aparam + if self.numb_aparam > 0: + sys_sumv = [] + sys_sumv2 = [] + sys_sumn = [] + for ss_ in [frame["aparam"] for frame in sampled]: + ss = np.reshape(ss_, [-1, self.numb_aparam]) + sys_sumv.append(np.sum(ss, axis=0)) + sys_sumv2.append(np.sum(ss * ss, axis=0)) + sys_sumn.append(ss.shape[0]) + sumv = np.sum(np.stack(sys_sumv), axis=0) + sumv2 = np.sum(np.stack(sys_sumv2), axis=0) + sumn = sum(sys_sumn) + aparam_avg = sumv / sumn + aparam_std = np.sqrt(sumv2 / sumn - (sumv / sumn) ** 2) + aparam_std = np.where( + aparam_std < protection, + np.array(protection, dtype=aparam_std.dtype), + aparam_std, + ) + aparam_inv_std = 1.0 / aparam_std + self.aparam_avg = aparam_avg.astype(self.aparam_avg.dtype) + self.aparam_inv_std = aparam_inv_std.astype(self.aparam_inv_std.dtype) + @abstractmethod def _net_out_dim(self) -> int: """Set the FittingNet output dim.""" diff --git a/deepmd/pd/model/atomic_model/base_atomic_model.py b/deepmd/pd/model/atomic_model/base_atomic_model.py index 65e8097ee2..e4bf23bf53 100644 --- a/deepmd/pd/model/atomic_model/base_atomic_model.py +++ b/deepmd/pd/model/atomic_model/base_atomic_model.py @@ -519,7 +519,7 @@ def compute_fitting_input_stat( self, sample_merged: Union[Callable[[], list[dict]], list[dict]], ) -> None: - """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. + """Compute the input statistics (e.g. mean and stddev) for the atomic model from packed data. Parameters ---------- diff --git a/source/tests/common/dpmodel/test_fitting_stat.py b/source/tests/common/dpmodel/test_fitting_stat.py new file mode 100644 index 0000000000..af7e189571 --- /dev/null +++ b/source/tests/common/dpmodel/test_fitting_stat.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + +from deepmd.dpmodel.descriptor import ( + DescrptSeA, +) +from deepmd.dpmodel.fitting import ( + EnergyFittingNet, +) + + +def _make_fake_data_pt(sys_natoms, sys_nframes, avgs, stds): + merged_output_stat = [] + nsys = len(sys_natoms) + ndof = len(avgs) + for ii in range(nsys): + sys_dict = {} + tmp_data_f = [] + tmp_data_a = [] + for jj in range(ndof): + rng = np.random.default_rng(2025 * ii + 220 * jj) + tmp_data_f.append( + rng.normal(loc=avgs[jj], scale=stds[jj], size=(sys_nframes[ii], 1)) + ) + rng = np.random.default_rng(220 * ii + 1636 * jj) + tmp_data_a.append( + rng.normal( + loc=avgs[jj], scale=stds[jj], size=(sys_nframes[ii], sys_natoms[ii]) + ) + ) + tmp_data_f = np.transpose(tmp_data_f, (1, 2, 0)) + tmp_data_a = np.transpose(tmp_data_a, (1, 2, 0)) + sys_dict["fparam"] = tmp_data_f + sys_dict["aparam"] = tmp_data_a + merged_output_stat.append(sys_dict) + return merged_output_stat + + +def _brute_fparam_pt(data, ndim): + adata = [ii["fparam"] for ii in data] + all_data = [] + for ii in adata: + tmp = np.reshape(ii, [-1, ndim]) + if len(all_data) == 0: + all_data = np.array(tmp) + else: + all_data = np.concatenate((all_data, tmp), axis=0) + avg = np.average(all_data, axis=0) + std = np.std(all_data, axis=0) + return avg, std + + +def _brute_aparam_pt(data, ndim): + adata = [ii["aparam"] for ii in data] + all_data = [] + for ii in adata: + tmp = np.reshape(ii, [-1, ndim]) + if len(all_data) == 0: + all_data = np.array(tmp) + else: + all_data = np.concatenate((all_data, tmp), axis=0) + avg = np.average(all_data, axis=0) + std = np.std(all_data, axis=0) + return avg, std + + +class TestEnerFittingStat(unittest.TestCase): + def test(self) -> None: + descrpt = DescrptSeA(6.0, 5.8, [46, 92], neuron=[25, 50, 100], axis_neuron=16) + fitting = EnergyFittingNet( + descrpt.get_ntypes(), + descrpt.get_dim_out(), + neuron=[240, 240, 240], + resnet_dt=True, + numb_fparam=3, + numb_aparam=3, + ) + avgs = [0, 10, 100] + stds = [2, 0.4, 0.00001] + sys_natoms = [10, 100] + sys_nframes = [5, 2] + all_data = _make_fake_data_pt(sys_natoms, sys_nframes, avgs, stds) + frefa, frefs = _brute_fparam_pt(all_data, len(avgs)) + arefa, arefs = _brute_aparam_pt(all_data, len(avgs)) + fitting.compute_input_stats(all_data, protection=1e-2) + frefs_inv = 1.0 / frefs + arefs_inv = 1.0 / arefs + frefs_inv[frefs_inv > 100] = 100 + arefs_inv[arefs_inv > 100] = 100 + np.testing.assert_almost_equal(frefa, fitting.fparam_avg) + np.testing.assert_almost_equal( + frefs_inv, fitting.fparam_inv_std + ) + np.testing.assert_almost_equal(arefa, fitting.aparam_avg) + np.testing.assert_almost_equal( + arefs_inv, fitting.aparam_inv_std + ) From 14e64e24ba6ce49c02d2e89118946c3c3e0a240c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 18 Nov 2025 16:18:07 +0000 Subject: [PATCH 13/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/common/dpmodel/test_fitting_stat.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/source/tests/common/dpmodel/test_fitting_stat.py b/source/tests/common/dpmodel/test_fitting_stat.py index af7e189571..101d2a9ad7 100644 --- a/source/tests/common/dpmodel/test_fitting_stat.py +++ b/source/tests/common/dpmodel/test_fitting_stat.py @@ -90,10 +90,6 @@ def test(self) -> None: frefs_inv[frefs_inv > 100] = 100 arefs_inv[arefs_inv > 100] = 100 np.testing.assert_almost_equal(frefa, fitting.fparam_avg) - np.testing.assert_almost_equal( - frefs_inv, fitting.fparam_inv_std - ) + np.testing.assert_almost_equal(frefs_inv, fitting.fparam_inv_std) np.testing.assert_almost_equal(arefa, fitting.aparam_avg) - np.testing.assert_almost_equal( - arefs_inv, fitting.aparam_inv_std - ) + np.testing.assert_almost_equal(arefs_inv, fitting.aparam_inv_std)