From d5a7c04384917be8955a0a2603019010ed90b70d Mon Sep 17 00:00:00 2001 From: theNotfish Date: Sat, 21 Sep 2024 13:59:24 +0800 Subject: [PATCH] Revert "fix: bugs in uts for property fit (#4120)" This reverts commit 96ed5df709f3cd2013dcf57583f3535f8dfb408d. --- backend/read_env.py | 2 - deepmd/pt/train/training.py | 21 ++ doc/install/install-from-source.md | 3 +- .../tests/pt/model/test_property_fitting.py | 256 ++++++------------ 4 files changed, 107 insertions(+), 175 deletions(-) diff --git a/backend/read_env.py b/backend/read_env.py index ae82778f4e..c3fe2d5127 100644 --- a/backend/read_env.py +++ b/backend/read_env.py @@ -60,8 +60,6 @@ def get_argument_from_env() -> Tuple[str, list, list, dict, str, str]: cmake_minimum_required_version = "3.21" cmake_args.append("-DUSE_ROCM_TOOLKIT:BOOL=TRUE") rocm_root = os.environ.get("ROCM_ROOT") - if not rocm_root: - rocm_root = os.environ.get("ROCM_PATH") if rocm_root: cmake_args.append(f"-DCMAKE_HIP_COMPILER_ROCM_ROOT:STRING={rocm_root}") hipcc_flags = os.environ.get("HIP_HIPCC_FLAGS") diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index a7b9e25b4e..40c5006ef6 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -392,6 +392,21 @@ def get_lr(lr_params): # JIT if JIT: self.model = torch.jit.script(self.model) + + # Initialize the fparam + if model_params["fitting_net"]["numb_fparam"] > 0: + nbatches = 10 + datasets = training_data.systems + dataloaders = training_data.dataloaders + fparams = [] + for i in range(len(datasets)): + iterator = iter(dataloaders[i]) + numb_batches = min(nbatches, len(dataloaders[i])) + for _ in range(numb_batches): + stat_data = next(iterator) + fparams.append(stat_data['fparam']) + fparams = torch.tensor(fparams) + init_fparam(self.model, fparams) # Model Wrapper self.wrapper = ModelWrapper(self.model, self.loss, model_params=model_params) @@ -1212,6 +1227,12 @@ def get_additional_data_requirement(_model): return additional_data_requirement +def init_fparam(_model, fparams): + fitting = _model.get_fitting_net() + fitting['fparam_avg'] = torch.unsqueeze(torch.mean(fparams, dim=0), dim=-1).to(DEVICE) + fitting['fparam_inv_std'] = torch.unsqueeze(1. / torch.std(fparams, dim=0), dim=-1).to(DEVICE) + + def get_loss(loss_params, start_lr, _ntypes, _model): loss_type = loss_params.get("type", "ener") if loss_type == "ener": diff --git a/doc/install/install-from-source.md b/doc/install/install-from-source.md index a725be0133..6f17a272c6 100644 --- a/doc/install/install-from-source.md +++ b/doc/install/install-from-source.md @@ -155,8 +155,7 @@ The path to the CUDA toolkit directory. CUDA 9.0 or later is supported. NVCC is **Type**: Path; **Default**: Detected automatically -The path to the ROCM toolkit directory. If `ROCM_ROOT` is not set, it will look for `ROCM_PATH`; if `ROCM_PATH` is also not set, it will be detected using `hipconfig --rocmpath`. - +The path to the ROCM toolkit directory. ::: :::{envvar} DP_ENABLE_TENSORFLOW diff --git a/source/tests/pt/model/test_property_fitting.py b/source/tests/pt/model/test_property_fitting.py index 59a5b1b172..df11f5b2f3 100644 --- a/source/tests/pt/model/test_property_fitting.py +++ b/source/tests/pt/model/test_property_fitting.py @@ -45,6 +45,7 @@ def setUp(self): self.rng = np.random.default_rng() self.nf, self.nloc, _ = self.nlist.shape self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) + self.scale = self.rng.uniform(0, 1, self.nt).tolist() def test_consistency( self, @@ -58,15 +59,16 @@ def test_consistency( self.atype_ext[:, : self.nloc], dtype=int, device=env.DEVICE ) - for nfp, nap, bias_atom_p, intensive, bias_method in itertools.product( + for mixed_types, nfp, nap, fit_diag, scale, bias_atom_p in itertools.product( + [True, False], [0, 3], [0, 4], + [True, False], + [None, self.scale], [ np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]), np.array([[11, 12, 13, 4, 15], [16, 17, 18, 9, 20]]), ], - [True, False], - ["normal", "no_bias"], ): ft0 = PropertyFittingNet( self.nt, @@ -74,10 +76,10 @@ def test_consistency( task_dim=5, numb_fparam=nfp, numb_aparam=nap, - mixed_types=self.dd0.mixed_types(), + mixed_types=mixed_types, + fit_diag=fit_diag, + scale=scale, bias_atom_p=bias_atom_p, - intensive=intensive, - bias_method=bias_method, ).to(env.DEVICE) ft1 = DPProperFittingNet.deserialize(ft0.serialize()) @@ -131,11 +133,11 @@ def test_consistency( def test_jit( self, ): - for nfp, nap, intensive, bias_method in itertools.product( + for mixed_types, nfp, nap, fit_diag in itertools.product( + [True, False], [0, 3], [0, 4], [True, False], - ["normal", "no_bias"], ): ft0 = PropertyFittingNet( self.nt, @@ -143,92 +145,13 @@ def test_jit( task_dim=5, numb_fparam=nfp, numb_aparam=nap, - mixed_types=self.dd0.mixed_types(), - intensive=intensive, - bias_method=bias_method, + mixed_types=mixed_types, + fit_diag=fit_diag, ).to(env.DEVICE) torch.jit.script(ft0) -class TestInvarianceOutCell(unittest.TestCase): - def setUp(self) -> None: - self.natoms = 5 - self.rcut = 4 - self.rcut_smth = 0.5 - self.sel = [46, 92, 4] - self.nf = 1 - self.nt = 3 - self.rng = np.random.default_rng() - self.coord = torch.tensor( - [ - [1.1042, 0.6852, 1.3582], - [1.8812, 1.6277, 0.3153], - [1.5655, 1.0383, 0.4152], - [0.9594, 1.2298, 0.8124], - [0.7905, 0.5014, 0.6654], - ], - dtype=dtype, - device=env.DEVICE, - ) - self.shift = torch.tensor([1000, 1000, 1000], dtype=dtype, device=env.DEVICE) - self.atype = torch.tensor([0, 0, 0, 1, 1], dtype=torch.int32, device=env.DEVICE) - self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) - self.cell = torch.tensor( - [ - [0.7333, 0.9166, 0.6533], - [0.1151, 0.9078, 0.2058], - [0.6907, 0.0370, 0.4863], - ], - dtype=dtype, - device=env.DEVICE, - ) - self.cell = (self.cell + self.cell.T) + 5.0 * torch.eye(3, device=env.DEVICE) - - def test_trans(self): - atype = self.atype.reshape(1, 5) - coord_s = torch.matmul( - torch.remainder( - torch.matmul(self.coord + self.shift, torch.linalg.inv(self.cell)), 1.0 - ), - self.cell, - ) - ft0 = PropertyFittingNet( - self.nt, - self.dd0.dim_out, - task_dim=11, - numb_fparam=0, - numb_aparam=0, - mixed_types=self.dd0.mixed_types(), - ).to(env.DEVICE) - res = [] - for xyz in [self.coord, coord_s]: - ( - extended_coord, - extended_atype, - _, - nlist, - ) = extend_input_and_build_neighbor_list( - xyz, - atype, - self.rcut, - self.sel, - self.dd0.mixed_types(), - box=self.cell, - ) - - rd0, gr0, _, _, _ = self.dd0( - extended_coord, - extended_atype, - nlist, - ) - - ret0 = ft0(rd0, atype, gr0, fparam=0, aparam=0) - res.append(ret0["property"]) - - np.testing.assert_allclose(to_numpy_array(res[0]), to_numpy_array(res[1])) - - -class TestInvarianceRandomShift(unittest.TestCase): +class TestInvariance(unittest.TestCase): def setUp(self) -> None: self.natoms = 5 self.rcut = 4 @@ -243,19 +166,19 @@ def setUp(self) -> None: self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) self.cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE) self.cell = (self.cell + self.cell.T) + 5.0 * torch.eye(3, device=env.DEVICE) + self.scale = self.rng.uniform(0, 1, self.nt).tolist() def test_rot(self): atype = self.atype.reshape(1, 5) rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype, device=env.DEVICE) coord_rot = torch.matmul(self.coord, rmat) - # use larger cell to rotate only coord and shift to the center of cell - cell_rot = 10.0 * torch.eye(3, dtype=dtype, device=env.DEVICE) - for nfp, nap, intensive, bias_method in itertools.product( + for mixed_types, nfp, nap, fit_diag, scale in itertools.product( + [True, False], [0, 3], [0, 4], [True, False], - ["normal", "no_bias"], + [None, self.scale], ): ft0 = PropertyFittingNet( self.nt, @@ -263,9 +186,9 @@ def test_rot(self): task_dim=9, numb_fparam=nfp, numb_aparam=nap, - mixed_types=self.dd0.mixed_types(), - intensive=intensive, - bias_method=bias_method, + mixed_types=True, + fit_diag=fit_diag, + scale=scale, ).to(env.DEVICE) if nfp > 0: ifp = torch.tensor( @@ -290,12 +213,7 @@ def test_rot(self): _, nlist, ) = extend_input_and_build_neighbor_list( - xyz + self.shift, - atype, - self.rcut, - self.sel, - self.dd0.mixed_types(), - box=cell_rot, + xyz + self.shift, atype, self.rcut, self.sel, mixed_types ) rd0, gr0, _, _, _ = self.dd0( @@ -304,7 +222,7 @@ def test_rot(self): nlist, ) - ret0 = ft0(rd0, atype, gr0, fparam=ifp, aparam=iap) + ret0 = ft0(rd0, extended_atype, gr0, fparam=ifp, aparam=iap) res.append(ret0["property"]) np.testing.assert_allclose( to_numpy_array(res[1]), @@ -313,44 +231,42 @@ def test_rot(self): def test_permu(self): coord = torch.matmul(self.coord, self.cell) - ft0 = PropertyFittingNet( - self.nt, - self.dd0.dim_out, - task_dim=8, - numb_fparam=0, - numb_aparam=0, - mixed_types=self.dd0.mixed_types(), - ).to(env.DEVICE) - res = [] - for idx_perm in [[0, 1, 2, 3, 4], [1, 0, 4, 3, 2]]: - atype = self.atype[idx_perm].reshape(1, 5) - ( - extended_coord, - extended_atype, - _, - nlist, - ) = extend_input_and_build_neighbor_list( - coord[idx_perm], - atype, - self.rcut, - self.sel, - self.dd0.mixed_types(), - box=self.cell, - ) + for fit_diag, scale in itertools.product([True, False], [None, self.scale]): + ft0 = PropertyFittingNet( + self.nt, + self.dd0.dim_out, + task_dim=8, + numb_fparam=0, + numb_aparam=0, + mixed_types=True, + fit_diag=fit_diag, + scale=scale, + ).to(env.DEVICE) + res = [] + for idx_perm in [[0, 1, 2, 3, 4], [1, 0, 4, 3, 2]]: + atype = self.atype[idx_perm].reshape(1, 5) + ( + extended_coord, + extended_atype, + _, + nlist, + ) = extend_input_and_build_neighbor_list( + coord[idx_perm], atype, self.rcut, self.sel, False + ) - rd0, gr0, _, _, _ = self.dd0( - extended_coord, - extended_atype, - nlist, - ) + rd0, gr0, _, _, _ = self.dd0( + extended_coord, + extended_atype, + nlist, + ) - ret0 = ft0(rd0, atype, gr0, fparam=None, aparam=None) - res.append(ret0["property"]) + ret0 = ft0(rd0, extended_atype, gr0, fparam=None, aparam=None) + res.append(ret0["property"]) - np.testing.assert_allclose( - to_numpy_array(res[0][:, idx_perm]), - to_numpy_array(res[1]), - ) + np.testing.assert_allclose( + to_numpy_array(res[0][:, idx_perm]), + to_numpy_array(res[1]), + ) def test_trans(self): atype = self.atype.reshape(1, 5) @@ -360,40 +276,38 @@ def test_trans(self): ), self.cell, ) - ft0 = PropertyFittingNet( - self.nt, - self.dd0.dim_out, - task_dim=11, - numb_fparam=0, - numb_aparam=0, - mixed_types=self.dd0.mixed_types(), - ).to(env.DEVICE) - res = [] - for xyz in [self.coord, coord_s]: - ( - extended_coord, - extended_atype, - _, - nlist, - ) = extend_input_and_build_neighbor_list( - xyz, - atype, - self.rcut, - self.sel, - self.dd0.mixed_types(), - box=self.cell, - ) + for fit_diag, scale in itertools.product([True, False], [None, self.scale]): + ft0 = PropertyFittingNet( + self.nt, + self.dd0.dim_out, + task_dim=11, + numb_fparam=0, + numb_aparam=0, + mixed_types=True, + fit_diag=fit_diag, + scale=scale, + ).to(env.DEVICE) + res = [] + for xyz in [self.coord, coord_s]: + ( + extended_coord, + extended_atype, + _, + nlist, + ) = extend_input_and_build_neighbor_list( + xyz, atype, self.rcut, self.sel, False + ) - rd0, gr0, _, _, _ = self.dd0( - extended_coord, - extended_atype, - nlist, - ) + rd0, gr0, _, _, _ = self.dd0( + extended_coord, + extended_atype, + nlist, + ) - ret0 = ft0(rd0, atype, gr0, fparam=0, aparam=0) - res.append(ret0["property"]) + ret0 = ft0(rd0, extended_atype, gr0, fparam=0, aparam=0) + res.append(ret0["property"]) - np.testing.assert_allclose(to_numpy_array(res[0]), to_numpy_array(res[1])) + np.testing.assert_allclose(to_numpy_array(res[0]), to_numpy_array(res[1])) class TestPropertyModel(unittest.TestCase): @@ -415,7 +329,7 @@ def setUp(self): task_dim=3, numb_fparam=0, numb_aparam=0, - mixed_types=self.dd0.mixed_types(), + mixed_types=True, intensive=True, ).to(env.DEVICE) self.type_mapping = ["O", "H", "B"]