diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index dcb381d53a..80b0e0c9d8 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -77,6 +77,7 @@ def __init__( repformer_update_style: str = "res_avg", repformer_set_davg_zero: bool = True, # TODO repformer_add_type_ebd_to_seq: bool = False, + trainable: bool = True, type: Optional[ str ] = None, # work around the bad design in get_trainer and DpLoaderSet! diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 47535580db..8385f06697 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy import logging +import os +import tempfile from abc import ( abstractmethod, ) @@ -12,6 +14,9 @@ import numpy as np import torch +from deepmd.infer.deep_eval import ( + DeepEval, +) from deepmd.pt.model.network.mlp import ( FittingNet, NetworkCollection, @@ -25,9 +30,6 @@ from deepmd.pt.utils import ( env, ) -from deepmd.pt.utils.dataloader import ( - DpLoaderSet, -) from deepmd.pt.utils.env import ( DEFAULT_PRECISION, DEVICE, @@ -36,13 +38,16 @@ from deepmd.pt.utils.exclude_mask import ( AtomExcludeMask, ) -from deepmd.pt.utils.stat import ( - make_stat_input, -) from deepmd.pt.utils.utils import ( to_numpy_array, to_torch_tensor, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) +from deepmd.utils.finetune import ( + change_energy_bias_lower, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -86,7 +91,13 @@ def share_params(self, base_class, shared_level, resume=False): raise NotImplementedError def change_energy_bias( - self, config, model, old_type_map, new_type_map, bias_shift="delta", ntest=10 + self, + config, + model, + old_type_map: List[str], + new_type_map: List[str], + bias_shift="delta", + ntest=10, ): """Change the energy bias according to the input data and the pretrained model. @@ -96,9 +107,9 @@ def change_energy_bias( The configuration. model : EnergyModel Energy model loaded pre-trained model. - new_type_map : list + new_type_map : List[str] The original type_map in dataset, they are targets to change the energy bias. - old_type_map : str + old_type_map : List[str] The full type_map in pretrained model bias_shift : str The mode for changing energy bias : ['delta', 'statistic'] @@ -114,93 +125,32 @@ def change_energy_bias( ) # data systems = config["training"]["training_data"]["systems"] - finetune_data = DpLoaderSet(systems, ntest, config["model"]) - sampled = make_stat_input(finetune_data.systems, finetune_data.dataloaders, 1) - # map - sorter = np.argsort(old_type_map) - idx_type_map = sorter[ - np.searchsorted(old_type_map, new_type_map, sorter=sorter) - ] - data_mixed_types = np.all([i.mixed_type for i in finetune_data.systems]) - numb_type = len(old_type_map) - type_numbs, energy_ground_truth, energy_predict = [], [], [] - for test_data in sampled: - nframes = test_data["energy"].shape[0] - if data_mixed_types: - atype = test_data["atype"].detach().cpu().numpy() - else: - atype = test_data["atype"][0].detach().cpu().numpy() - assert np.array( - [i.item() in idx_type_map for i in list(set(atype.reshape(-1)))] - ).all(), "Some types are not in 'type_map'!" - energy_ground_truth.append(test_data["energy"].cpu().numpy()) - if data_mixed_types: - type_numbs.append( - np.array( - [(atype == i).sum(axis=-1) for i in idx_type_map], - dtype=np.int32, - ).T - ) - else: - type_numbs.append( - np.tile( - np.bincount(atype, minlength=numb_type)[idx_type_map], - (nframes, 1), - ) - ) - if bias_shift == "delta": - coord = test_data["coord"].to(DEVICE) - atype = test_data["atype"].to(DEVICE) - box = ( - test_data["box"].to(DEVICE) - if test_data["box"] is not None - else None - ) - ret = model(coord, atype, box) - energy_predict.append( - ret["energy"].reshape([nframes, 1]).detach().cpu().numpy() - ) - type_numbs = np.concatenate(type_numbs) - energy_ground_truth = np.concatenate(energy_ground_truth) - old_bias = self.bias_atom_e[idx_type_map] - if bias_shift == "delta": - energy_predict = np.concatenate(energy_predict) - bias_diff = energy_ground_truth - energy_predict - delta_bias = np.linalg.lstsq(type_numbs, bias_diff, rcond=None)[0] - unbias_e = energy_predict + type_numbs @ delta_bias - atom_numbs = type_numbs.sum(-1) - rmse_ae = np.sqrt( - np.mean( - np.square( - (unbias_e.ravel() - energy_ground_truth.ravel()) / atom_numbs - ) - ) - ) - self.bias_atom_e[idx_type_map] += torch.from_numpy( - delta_bias.reshape(-1) - ).to(DEVICE) - log.info( - f"RMSE of atomic energy after linear regression is: {rmse_ae:10.5e} eV/atom." - ) - elif bias_shift == "statistic": - statistic_bias = np.linalg.lstsq( - type_numbs, energy_ground_truth, rcond=None - )[0] - self.bias_atom_e[idx_type_map] = ( - torch.from_numpy(statistic_bias.reshape(-1)) - .type_as(self.bias_atom_e[idx_type_map]) - .to(DEVICE) - ) - else: - raise RuntimeError("Unknown bias_shift mode: " + bias_shift) - log.info( - "Change energy bias of {} from {} to {}.".format( - str(new_type_map), - str(old_bias.detach().cpu().numpy()), - str(self.bias_atom_e[idx_type_map].detach().cpu().numpy()), - ) + finetune_data = DeepmdDataSystem( + systems=systems, + batch_size=config["training"]["training_data"].get("batch_size", "auto"), + test_size=1, + ) + finetune_data.add("energy", ndof=1, atomic=False, must=True, high_prec=True) + model = torch.jit.script(model) + tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth") + torch.jit.save(model, tmp_model.name) + dp = DeepEval(tmp_model.name) + os.unlink(tmp_model.name) + bias = change_energy_bias_lower( + finetune_data, + dp, + new_type_map, + old_type_map, + self.bias_atom_e.detach().cpu().numpy().reshape(-1), + bias_shift=bias_shift, + ntest=ntest, + ) + self.bias_atom_e = ( + torch.from_numpy(bias) + .type_as(self.bias_atom_e) + .reshape(self.bias_atom_e.shape) + .to(DEVICE) ) - return None class GeneralFitting(Fitting): diff --git a/deepmd/pt/utils/finetune.py b/deepmd/pt/utils/finetune.py index c8fa1e5185..d555478af4 100644 --- a/deepmd/pt/utils/finetune.py +++ b/deepmd/pt/utils/finetune.py @@ -27,7 +27,6 @@ def change_finetune_model_params( last_model_params = state_dict["_extra_state"]["model_params"] finetune_multi_task = "model_dict" in last_model_params trainable_param = { - "type_embedding": True, "descriptor": True, "fitting_net": True, } @@ -74,7 +73,7 @@ def change_finetune_model_params( assert set(new_type_map).issubset( old_type_map ), "Only support for smaller type map when finetuning or resuming." - for key_item in ["type_map", "type_embedding", "descriptor"]: + for key_item in ["type_map", "descriptor"]: if key_item in model_dict_params[model_branch_chosen]: model_config[key_item] = model_dict_params[model_branch_chosen][ key_item diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index a842df50bd..82375694a7 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -53,6 +53,9 @@ from deepmd.tf.utils.spin import ( Spin, ) +from deepmd.utils.finetune import ( + change_energy_bias_lower, +) from deepmd.utils.out_stat import ( compute_stats_from_redu, ) @@ -793,121 +796,18 @@ def change_energy_bias( bias_shift="delta", ntest=10, ) -> None: - """Change the energy bias according to the input data and the pretrained model. - - Parameters - ---------- - data : DeepmdDataSystem - The training data. - frozen_model : str - The path file of frozen model. - origin_type_map : list - The original type_map in dataset, they are targets to change the energy bias. - full_type_map : str - The full type_map in pretrained model - bias_shift : str - The mode for changing energy bias : ['delta', 'statistic'] - 'delta' : perform predictions on energies of target dataset, - and do least sqaure on the errors to obtain the target shift as bias. - 'statistic' : directly use the statistic energy bias in the target dataset. - ntest : int - The number of test samples in a system to change the energy bias. - """ - type_numbs = [] - energy_ground_truth = [] - energy_predict = [] - sorter = np.argsort(full_type_map) - idx_type_map = sorter[ - np.searchsorted(full_type_map, origin_type_map, sorter=sorter) - ] - mixed_type = data.mixed_type - numb_type = len(full_type_map) dp = None if bias_shift == "delta": # init model dp = DeepPotential(frozen_model) - for sys in data.data_systems: - test_data = sys.get_test() - nframes = test_data["box"].shape[0] - numb_test = min(nframes, ntest) - if mixed_type: - atype = test_data["type"][:numb_test].reshape([numb_test, -1]) - else: - atype = test_data["type"][0] - assert np.array( - [i in idx_type_map for i in list(set(atype.reshape(-1)))] - ).all(), "Some types are not in 'type_map'!" - energy_ground_truth.append( - test_data["energy"][:numb_test].reshape([numb_test, 1]) - ) - if mixed_type: - type_numbs.append( - np.array( - [(atype == i).sum(axis=-1) for i in idx_type_map], - dtype=np.int32, - ).T - ) - else: - type_numbs.append( - np.tile( - np.bincount(atype, minlength=numb_type)[idx_type_map], - (numb_test, 1), - ) - ) - if bias_shift == "delta": - coord = test_data["coord"][:numb_test].reshape([numb_test, -1]) - if sys.pbc: - box = test_data["box"][:numb_test] - else: - box = None - if dp.get_dim_fparam() > 0: - fparam = test_data["fparam"][:numb_test] - else: - fparam = None - if dp.get_dim_aparam() > 0: - aparam = test_data["aparam"][:numb_test] - else: - aparam = None - ret = dp.eval( - coord, - box, - atype, - mixed_type=mixed_type, - fparam=fparam, - aparam=aparam, - ) - energy_predict.append(ret[0].reshape([numb_test, 1])) - type_numbs = np.concatenate(type_numbs) - energy_ground_truth = np.concatenate(energy_ground_truth) - old_bias = self.bias_atom_e[idx_type_map] - if bias_shift == "delta": - energy_predict = np.concatenate(energy_predict) - bias_diff = energy_ground_truth - energy_predict - delta_bias = np.linalg.lstsq(type_numbs, bias_diff, rcond=None)[0] - unbias_e = energy_predict + type_numbs @ delta_bias - atom_numbs = type_numbs.sum(-1) - rmse_ae = np.sqrt( - np.mean( - np.square( - (unbias_e.ravel() - energy_ground_truth.ravel()) / atom_numbs - ) - ) - ) - self.bias_atom_e[idx_type_map] += delta_bias.reshape(-1) - log.info( - f"RMSE of atomic energy after linear regression is: {rmse_ae} eV/atom." - ) - elif bias_shift == "statistic": - statistic_bias = np.linalg.lstsq( - type_numbs, energy_ground_truth, rcond=None - )[0] - self.bias_atom_e[idx_type_map] = statistic_bias.reshape(-1) - else: - raise RuntimeError("Unknown bias_shift mode: " + bias_shift) - log.info( - "Change energy bias of {} from {} to {}.".format( - str(origin_type_map), str(old_bias), str(self.bias_atom_e[idx_type_map]) - ) + self.bias_atom_e = change_energy_bias_lower( + data, + dp, + origin_type_map, + full_type_map, + self.bias_atom_e, + bias_shift=bias_shift, + ntest=ntest, ) def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: diff --git a/deepmd/utils/finetune.py b/deepmd/utils/finetune.py new file mode 100644 index 0000000000..b6d04b9bc5 --- /dev/null +++ b/deepmd/utils/finetune.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from typing import ( + TYPE_CHECKING, + List, +) + +import numpy as np + +from deepmd.infer.deep_eval import ( + DeepEval, +) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) + +if TYPE_CHECKING: + pass + +log = logging.getLogger(__name__) + + +def change_energy_bias_lower( + data: DeepmdDataSystem, + dp: DeepEval, + origin_type_map: List[str], + full_type_map: List[str], + bias_atom_e: np.ndarray, + bias_shift="delta", + ntest=10, +): + """Change the energy bias according to the input data and the pretrained model. + + Parameters + ---------- + data : DeepmdDataSystem + The training data. + dp : str + The DeepEval object. + origin_type_map : list + The original type_map in dataset, they are targets to change the energy bias. + full_type_map : str + The full type_map in pretrained model + bias_atom_e : np.ndarray + The old energy bias in the pretrained model. + bias_shift : str + The mode for changing energy bias : ['delta', 'statistic'] + 'delta' : perform predictions on energies of target dataset, + and do least sqaure on the errors to obtain the target shift as bias. + 'statistic' : directly use the statistic energy bias in the target dataset. + ntest : int + The number of test samples in a system to change the energy bias. + """ + type_numbs = [] + energy_ground_truth = [] + energy_predict = [] + sorter = np.argsort(full_type_map) + idx_type_map = sorter[ + np.searchsorted(full_type_map, origin_type_map, sorter=sorter) + ] + mixed_type = data.mixed_type + numb_type = len(full_type_map) + for sys in data.data_systems: + test_data = sys.get_test() + nframes = test_data["box"].shape[0] + numb_test = min(nframes, ntest) + if mixed_type: + atype = test_data["type"][:numb_test].reshape([numb_test, -1]) + else: + atype = test_data["type"][0] + assert np.array( + [i in idx_type_map for i in list(set(atype.reshape(-1)))] + ).all(), "Some types are not in 'type_map'!" + energy_ground_truth.append( + test_data["energy"][:numb_test].reshape([numb_test, 1]) + ) + if mixed_type: + type_numbs.append( + np.array( + [(atype == i).sum(axis=-1) for i in idx_type_map], + dtype=np.int32, + ).T + ) + else: + type_numbs.append( + np.tile( + np.bincount(atype, minlength=numb_type)[idx_type_map], + (numb_test, 1), + ) + ) + if bias_shift == "delta": + coord = test_data["coord"][:numb_test].reshape([numb_test, -1]) + if sys.pbc: + box = test_data["box"][:numb_test] + else: + box = None + if dp.get_dim_fparam() > 0: + fparam = test_data["fparam"][:numb_test] + else: + fparam = None + if dp.get_dim_aparam() > 0: + aparam = test_data["aparam"][:numb_test] + else: + aparam = None + ret = dp.eval( + coord, + box, + atype, + mixed_type=mixed_type, + fparam=fparam, + aparam=aparam, + ) + energy_predict.append(ret[0].reshape([numb_test, 1])) + type_numbs = np.concatenate(type_numbs) + energy_ground_truth = np.concatenate(energy_ground_truth) + old_bias = bias_atom_e[idx_type_map] + if bias_shift == "delta": + energy_predict = np.concatenate(energy_predict) + bias_diff = energy_ground_truth - energy_predict + delta_bias = np.linalg.lstsq(type_numbs, bias_diff, rcond=None)[0] + unbias_e = energy_predict + type_numbs @ delta_bias + atom_numbs = type_numbs.sum(-1) + rmse_ae = np.sqrt( + np.mean( + np.square((unbias_e.ravel() - energy_ground_truth.ravel()) / atom_numbs) + ) + ) + bias_atom_e[idx_type_map] += delta_bias.reshape(-1) + log.info( + f"RMSE of atomic energy after linear regression is: {rmse_ae} eV/atom." + ) + elif bias_shift == "statistic": + statistic_bias = np.linalg.lstsq(type_numbs, energy_ground_truth, rcond=None)[0] + bias_atom_e[idx_type_map] = statistic_bias.reshape(-1) + else: + raise RuntimeError("Unknown bias_shift mode: " + bias_shift) + log.info( + "Change energy bias of {} from {} to {}.".format( + str(origin_type_map), str(old_bias), str(bias_atom_e[idx_type_map]) + ) + ) + return bias_atom_e diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index 4e73fc4f8a..dcc70ce63a 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -24,8 +24,13 @@ class DPTrainTest: def test_dp_train(self): + # test training from scratch trainer = get_trainer(deepcopy(self.config)) trainer.run() + + # test fine-tuning + trainer_finetune = get_trainer(deepcopy(self.config), finetune_model="model.pt") + trainer_finetune.run() self.tearDown() def tearDown(self):