From 0e8c046540280068c46e76f6e386857c82b2306f Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 6 Jul 2021 23:30:28 +0800 Subject: [PATCH 1/2] automatically determine the `sel` from the training data. --- deepmd/entrypoints/__init__.py | 4 +- deepmd/entrypoints/main.py | 4 +- deepmd/entrypoints/train.py | 85 ++++++++++++++++++++++++++++++++++ deepmd/utils/argcheck.py | 6 +-- 4 files changed, 93 insertions(+), 6 deletions(-) diff --git a/deepmd/entrypoints/__init__.py b/deepmd/entrypoints/__init__.py index 4a02b995f3..d92cbf3076 100644 --- a/deepmd/entrypoints/__init__.py +++ b/deepmd/entrypoints/__init__.py @@ -5,7 +5,9 @@ from .doc import doc_train_input from .freeze import freeze from .test import test -from .train import train +# import `train` as `train_dp` to avoid the conflict of the +# module name `train` and the function name `train` +from .train import train as train_dp from .transfer import transfer from ..infer.model_devi import make_model_devi from .convert import convert diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index df2db34923..cda9856a24 100644 --- a/deepmd/entrypoints/main.py +++ b/deepmd/entrypoints/main.py @@ -11,7 +11,7 @@ doc_train_input, freeze, test, - train, + train_dp, transfer, make_model_devi, convert, @@ -416,7 +416,7 @@ def main(): dict_args = vars(args) if args.command == "train": - train(**dict_args) + train_dp(**dict_args) elif args.command == "freeze": freeze(**dict_args) elif args.command == "config": diff --git a/deepmd/entrypoints/train.py b/deepmd/entrypoints/train.py index 9fa8344a8b..6c6e616115 100755 --- a/deepmd/entrypoints/train.py +++ b/deepmd/entrypoints/train.py @@ -19,6 +19,7 @@ from deepmd.utils.compat import updata_deepmd_input from deepmd.utils.data_system import DeepmdDataSystem from deepmd.utils.sess import run_sess +from deepmd.utils.neighbor_stat import NeighborStat if TYPE_CHECKING: from deepmd.run_options import TFServerV1 @@ -173,6 +174,9 @@ def train( jdata = updata_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") jdata = normalize(jdata) + + jdata = update_sel(jdata) + with open(output, "w") as fp: json.dump(jdata, fp, indent=4) @@ -327,3 +331,84 @@ def get_modifier(modi_data=None): else: modifier = None return modifier + + +def get_rcut(jdata): + descrpt_data = jdata['model']['descriptor'] + rcut_list = [] + if descrpt_data['type'] == 'hybrid': + for ii in descrpt_data['list']: + rcut_list.append(ii['rcut']) + else: + rcut_list.append(descrpt_data['rcut']) + return max(rcut_list) + + +def get_type_map(jdata): + return jdata['model'].get('type_map', None) + + +def get_sel(jdata, rcut): + max_rcut = get_rcut(jdata) + type_map = get_type_map(jdata) + + if len(type_map) == 0: + type_map = None + train_data = get_data(jdata["training"]["training_data"], max_rcut, type_map, None) + train_data.get_batch() + data_ntypes = train_data.get_ntypes() + if type_map is not None: + map_ntypes = len(type_map) + else: + map_ntypes = data_ntypes + ntypes = max([map_ntypes, data_ntypes]) + + neistat = NeighborStat(ntypes, rcut) + + min_nbor_dist, max_nbor_size = neistat.get_stat(train_data) + + return max_nbor_size + + +def parse_auto_sel(sel): + if type(sel) is not str: + return False + words = sel.split(':') + if words[0] == 'auto': + return True + else: + return False + + +def parse_auto_sel_ratio(sel): + if not parse_auto_sel(sel): + raise RuntimeError(f'invalid auto sel format {sel}') + else: + words = sel.split(':') + if len(words) == 1: + ratio = 1.1 + elif len(words) == 2: + ratio = float(words[1]) + else: + raise RuntimeError(f'invalid auto sel format {sel}') + return ratio + + +def update_one_sel(jdata, descriptor): + if parse_auto_sel(descriptor['sel']) : + ratio = parse_auto_sel_ratio(descriptor['sel']) + rcut = descriptor['rcut'] + tmp_sel = get_sel(jdata, rcut) + descriptor['sel'] = [int(ii * ratio) for ii in tmp_sel] + return descriptor + + +def update_sel(jdata): + descrpt_data = jdata['model']['descriptor'] + if descrpt_data['type'] == 'hybrid': + for ii in range(len(descrpt_data['list'])): + descrpt_data['list'][ii] = update_one_sel(jdata, descrpt_data['list'][ii]) + else: + descrpt_data = update_one_sel(jdata, descrpt_data) + jdata['model']['descriptor'] = descrpt_data + return jdata diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 684b0eacbc..5cf7c2e9c6 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -73,7 +73,7 @@ def descrpt_se_a_args(): doc_set_davg_zero = 'Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used' return [ - Argument("sel", list, optional = False, doc = doc_sel), + Argument("sel", [list,str], optional = True, default = "auto", doc = doc_sel), Argument("rcut", float, optional = True, default = 6.0, doc = doc_rcut), Argument("rcut_smth", float, optional = True, default = 0.5, doc = doc_rcut_smth), Argument("neuron", list, optional = True, default = [10,20,40], doc = doc_neuron), @@ -102,7 +102,7 @@ def descrpt_se_t_args(): doc_set_davg_zero = 'Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used' return [ - Argument("sel", list, optional = False, doc = doc_sel), + Argument("sel", [list,str], optional = True, default = "auto", doc = doc_sel), Argument("rcut", float, optional = True, default = 6.0, doc = doc_rcut), Argument("rcut_smth", float, optional = True, default = 0.5, doc = doc_rcut_smth), Argument("neuron", list, optional = True, default = [10,20,40], doc = doc_neuron), @@ -143,7 +143,7 @@ def descrpt_se_r_args(): doc_set_davg_zero = 'Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used' return [ - Argument("sel", list, optional = False, doc = doc_sel), + Argument("sel", [list,str], optional = True, default = "auto", doc = doc_sel), Argument("rcut", float, optional = True, default = 6.0, doc = doc_rcut), Argument("rcut_smth", float, optional = True, default = 0.5, doc = doc_rcut_smth), Argument("neuron", list, optional = True, default = [10,20,40], doc = doc_neuron), From 0bac52fca246212ee26551139f2dabd08fc10faf Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 9 Jul 2021 16:13:16 +0800 Subject: [PATCH 2/2] add UT for auto sel. update doc. wrap auto determined sel to be 4 divisible --- deepmd/entrypoints/train.py | 6 +- deepmd/utils/argcheck.py | 12 +++- doc/train-input-auto.rst | 72 +++++++++++++--------- source/tests/test_train.py | 120 ++++++++++++++++++++++++++++++++++++ 4 files changed, 178 insertions(+), 32 deletions(-) create mode 100644 source/tests/test_train.py diff --git a/deepmd/entrypoints/train.py b/deepmd/entrypoints/train.py index 6c6e616115..cffac25947 100755 --- a/deepmd/entrypoints/train.py +++ b/deepmd/entrypoints/train.py @@ -394,12 +394,16 @@ def parse_auto_sel_ratio(sel): return ratio +def wrap_up_4(xx): + return 4 * ((int(xx) + 3) // 4) + + def update_one_sel(jdata, descriptor): if parse_auto_sel(descriptor['sel']) : ratio = parse_auto_sel_ratio(descriptor['sel']) rcut = descriptor['rcut'] tmp_sel = get_sel(jdata, rcut) - descriptor['sel'] = [int(ii * ratio) for ii in tmp_sel] + descriptor['sel'] = [int(wrap_up_4(ii * ratio)) for ii in tmp_sel] return descriptor diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 5cf7c2e9c6..4dfd552d89 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -58,7 +58,9 @@ def descrpt_local_frame_args (): def descrpt_se_a_args(): - doc_sel = 'A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.' + doc_sel = 'This parameter set the number of selected neighbors for each type of atom. It can be:\n\n\ + - `List[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.\n\n\ + - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' doc_rcut = 'The cut-off radius.' doc_rcut_smth = 'Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`' doc_neuron = 'Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built.' @@ -90,7 +92,9 @@ def descrpt_se_a_args(): def descrpt_se_t_args(): - doc_sel = 'A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.' + doc_sel = 'This parameter set the number of selected neighbors for each type of atom. It can be:\n\n\ + - `List[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.\n\n\ + - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' doc_rcut = 'The cut-off radius.' doc_rcut_smth = 'Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`' doc_neuron = 'Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built.' @@ -129,7 +133,9 @@ def descrpt_se_a_tpe_args(): def descrpt_se_r_args(): - doc_sel = 'A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.' + doc_sel = 'This parameter set the number of selected neighbors for each type of atom. It can be:\n\n\ + - `List[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.\n\n\ + - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' doc_rcut = 'The cut-off radius.' doc_rcut_smth = 'Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`' doc_neuron = 'Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built.' diff --git a/doc/train-input-auto.rst b/doc/train-input-auto.rst index e51a96c6a2..cf9063d7ca 100644 --- a/doc/train-input-auto.rst +++ b/doc/train-input-auto.rst @@ -221,10 +221,14 @@ model: .. _`model/descriptor[se_e2_a]/sel`: sel: - | type: ``list`` + | type: ``list`` | ``str``, optional, default: ``auto`` | argument path: ``model/descriptor[se_e2_a]/sel`` - A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment. + This parameter set the number of selected neighbors for each type of atom. It can be: + + - `List[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment. + + - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1". .. _`model/descriptor[se_e2_a]/rcut`: @@ -330,10 +334,14 @@ model: .. _`model/descriptor[se_e2_r]/sel`: sel: - | type: ``list`` + | type: ``list`` | ``str``, optional, default: ``auto`` | argument path: ``model/descriptor[se_e2_r]/sel`` - A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment. + This parameter set the number of selected neighbors for each type of atom. It can be: + + - `List[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment. + + - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1". .. _`model/descriptor[se_e2_r]/rcut`: @@ -431,10 +439,14 @@ model: .. _`model/descriptor[se_e3]/sel`: sel: - | type: ``list`` + | type: ``list`` | ``str``, optional, default: ``auto`` | argument path: ``model/descriptor[se_e3]/sel`` - A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment. + This parameter set the number of selected neighbors for each type of atom. It can be: + + - `List[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment. + + - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1". .. _`model/descriptor[se_e3]/rcut`: @@ -516,10 +528,14 @@ model: .. _`model/descriptor[se_a_tpe]/sel`: sel: - | type: ``list`` + | type: ``list`` | ``str``, optional, default: ``auto`` | argument path: ``model/descriptor[se_a_tpe]/sel`` - A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment. + This parameter set the number of selected neighbors for each type of atom. It can be: + + - `List[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment. + + - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1". .. _`model/descriptor[se_a_tpe]/rcut`: @@ -748,7 +764,7 @@ model: .. _`model/fitting_net[ener]/trainable`: trainable: - | type: ``bool`` | ``list``, optional, default: ``True`` + | type: ``list`` | ``bool``, optional, default: ``True`` | argument path: ``model/fitting_net[ener]/trainable`` Whether the parameters in the fitting net are trainable. This option can be @@ -821,7 +837,7 @@ model: .. _`model/fitting_net[dipole]/sel_type`: sel_type: - | type: ``int`` | ``NoneType`` | ``list``, optional + | type: ``list`` | ``int`` | ``NoneType``, optional | argument path: ``model/fitting_net[dipole]/sel_type`` The atom types for which the atomic dipole will be provided. If not set, all types will be selected. @@ -882,7 +898,7 @@ model: .. _`model/fitting_net[polar]/scale`: scale: - | type: ``float`` | ``list``, optional, default: ``1.0`` + | type: ``list`` | ``float``, optional, default: ``1.0`` | argument path: ``model/fitting_net[polar]/scale`` The output of the fitting net (polarizability matrix) will be scaled by ``scale`` @@ -898,7 +914,7 @@ model: .. _`model/fitting_net[polar]/sel_type`: sel_type: - | type: ``int`` | ``NoneType`` | ``list``, optional + | type: ``list`` | ``int`` | ``NoneType``, optional | argument path: ``model/fitting_net[polar]/sel_type`` The atom types for which the atomic polarizability will be provided. If not set, all types will be selected. @@ -1078,7 +1094,7 @@ loss: .. _`loss[ener]/start_pref_e`: start_pref_e: - | type: ``float`` | ``int``, optional, default: ``0.02`` + | type: ``int`` | ``float``, optional, default: ``0.02`` | argument path: ``loss[ener]/start_pref_e`` The prefactor of energy loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the energy label should be provided by file energy.npy in each data system. If both start_pref_energy and limit_pref_energy are set to 0, then the energy will be ignored. @@ -1086,7 +1102,7 @@ loss: .. _`loss[ener]/limit_pref_e`: limit_pref_e: - | type: ``float`` | ``int``, optional, default: ``1.0`` + | type: ``int`` | ``float``, optional, default: ``1.0`` | argument path: ``loss[ener]/limit_pref_e`` The prefactor of energy loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity. @@ -1094,7 +1110,7 @@ loss: .. _`loss[ener]/start_pref_f`: start_pref_f: - | type: ``float`` | ``int``, optional, default: ``1000`` + | type: ``int`` | ``float``, optional, default: ``1000`` | argument path: ``loss[ener]/start_pref_f`` The prefactor of force loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the force label should be provided by file force.npy in each data system. If both start_pref_force and limit_pref_force are set to 0, then the force will be ignored. @@ -1102,7 +1118,7 @@ loss: .. _`loss[ener]/limit_pref_f`: limit_pref_f: - | type: ``float`` | ``int``, optional, default: ``1.0`` + | type: ``int`` | ``float``, optional, default: ``1.0`` | argument path: ``loss[ener]/limit_pref_f`` The prefactor of force loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity. @@ -1110,7 +1126,7 @@ loss: .. _`loss[ener]/start_pref_v`: start_pref_v: - | type: ``float`` | ``int``, optional, default: ``0.0`` + | type: ``int`` | ``float``, optional, default: ``0.0`` | argument path: ``loss[ener]/start_pref_v`` The prefactor of virial loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the virial label should be provided by file virial.npy in each data system. If both start_pref_virial and limit_pref_virial are set to 0, then the virial will be ignored. @@ -1118,7 +1134,7 @@ loss: .. _`loss[ener]/limit_pref_v`: limit_pref_v: - | type: ``float`` | ``int``, optional, default: ``0.0`` + | type: ``int`` | ``float``, optional, default: ``0.0`` | argument path: ``loss[ener]/limit_pref_v`` The prefactor of virial loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity. @@ -1126,7 +1142,7 @@ loss: .. _`loss[ener]/start_pref_ae`: start_pref_ae: - | type: ``float`` | ``int``, optional, default: ``0.0`` + | type: ``int`` | ``float``, optional, default: ``0.0`` | argument path: ``loss[ener]/start_pref_ae`` The prefactor of atom_ener loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the atom_ener label should be provided by file atom_ener.npy in each data system. If both start_pref_atom_ener and limit_pref_atom_ener are set to 0, then the atom_ener will be ignored. @@ -1134,7 +1150,7 @@ loss: .. _`loss[ener]/limit_pref_ae`: limit_pref_ae: - | type: ``float`` | ``int``, optional, default: ``0.0`` + | type: ``int`` | ``float``, optional, default: ``0.0`` | argument path: ``loss[ener]/limit_pref_ae`` The prefactor of atom_ener loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity. @@ -1155,7 +1171,7 @@ loss: .. _`loss[tensor]/pref`: pref: - | type: ``float`` | ``int`` + | type: ``int`` | ``float`` | argument path: ``loss[tensor]/pref`` The prefactor of the weight of global loss. It should be larger than or equal to 0. If controls the weight of loss corresponding to global label, i.e. 'polarizability.npy` or `dipole.npy`, whose shape should be #frames x [9 or 3]. If it's larger than 0.0, this npy should be included. @@ -1163,7 +1179,7 @@ loss: .. _`loss[tensor]/pref_atomic`: pref_atomic: - | type: ``float`` | ``int`` + | type: ``int`` | ``float`` | argument path: ``loss[tensor]/pref_atomic`` The prefactor of the weight of atomic loss. It should be larger than or equal to 0. If controls the weight of loss corresponding to atomic label, i.e. `atomic_polarizability.npy` or `atomic_dipole.npy`, whose shape should be #frames x ([9 or 3] x #selected atoms). If it's larger than 0.0, this npy should be included. Both `pref` and `pref_atomic` should be provided, and either can be set to 0.0. @@ -1260,7 +1276,7 @@ training: .. _`training/training_data/batch_size`: batch_size: - | type: ``int`` | ``list`` | ``str``, optional, default: ``auto`` + | type: ``list`` | ``str`` | ``int``, optional, default: ``auto`` | argument path: ``training/training_data/batch_size`` This key can be @@ -1290,7 +1306,7 @@ training: .. _`training/training_data/sys_probs`: sys_probs: - | type: ``NoneType`` | ``list``, optional, default: ``None``, alias: *sys_weights* + | type: ``list`` | ``NoneType``, optional, default: ``None``, alias: *sys_weights* | argument path: ``training/training_data/sys_probs`` A list of float if specified. Should be of the same length as `systems`, specifying the probability of each system. @@ -1298,7 +1314,7 @@ training: .. _`training/validation_data`: validation_data: - | type: ``NoneType`` | ``dict``, optional, default: ``None`` + | type: ``dict`` | ``NoneType``, optional, default: ``None`` | argument path: ``training/validation_data`` Configurations of validation data. Similar to that of training data, except that a `numb_btch` argument may be configured @@ -1322,7 +1338,7 @@ training: .. _`training/validation_data/batch_size`: batch_size: - | type: ``int`` | ``list`` | ``str``, optional, default: ``auto`` + | type: ``list`` | ``str`` | ``int``, optional, default: ``auto`` | argument path: ``training/validation_data/batch_size`` This key can be @@ -1352,7 +1368,7 @@ training: .. _`training/validation_data/sys_probs`: sys_probs: - | type: ``NoneType`` | ``list``, optional, default: ``None``, alias: *sys_weights* + | type: ``list`` | ``NoneType``, optional, default: ``None``, alias: *sys_weights* | argument path: ``training/validation_data/sys_probs`` A list of float if specified. Should be of the same length as `systems`, specifying the probability of each system. @@ -1400,7 +1416,7 @@ training: .. _`training/numb_test`: numb_test: - | type: ``int`` | ``list`` | ``str``, optional, default: ``1`` + | type: ``list`` | ``str`` | ``int``, optional, default: ``1`` | argument path: ``training/numb_test`` Number of frames used for the test during training. diff --git a/source/tests/test_train.py b/source/tests/test_train.py new file mode 100644 index 0000000000..21892fb8fe --- /dev/null +++ b/source/tests/test_train.py @@ -0,0 +1,120 @@ +import unittest +from unittest.mock import patch, MagicMock + +import deepmd +from deepmd.entrypoints.train import parse_auto_sel, parse_auto_sel_ratio, wrap_up_4, update_one_sel, update_sel + +class TestTrain (unittest.TestCase) : + def test_train_parse_auto_sel (self) : + self.assertTrue(parse_auto_sel("auto")) + self.assertTrue(parse_auto_sel("auto:12")) + self.assertTrue(parse_auto_sel("auto:12:13")) + self.assertFalse(parse_auto_sel([1,2])) + self.assertFalse(parse_auto_sel("abc:12:13")) + + + def test_train_parse_auto_sel_ratio (self) : + self.assertEqual(parse_auto_sel_ratio("auto"), 1.1) + self.assertEqual(parse_auto_sel_ratio("auto:1.2"), 1.2) + with self.assertRaises(RuntimeError): + parse_auto_sel_ratio("auto:1.2:1.3") + with self.assertRaises(RuntimeError): + parse_auto_sel_ratio("abc") + with self.assertRaises(RuntimeError): + parse_auto_sel_ratio([1,2,3]) + + + @patch("deepmd.entrypoints.train.get_sel") + def test_update_one_sel(self, sel_mock): + sel_mock.return_value = [10,20] + jdata = {} + descriptor = { + 'rcut': 6, + 'sel': "auto" + } + descriptor = update_one_sel(jdata, descriptor) + # self.assertEqual(descriptor['sel'], [11,22]) + self.assertEqual(descriptor['sel'], [12,24]) + descriptor = { + 'rcut': 6, + 'sel': "auto:1.5" + } + descriptor = update_one_sel(jdata, descriptor) + # self.assertEqual(descriptor['sel'], [15,30]) + self.assertEqual(descriptor['sel'], [16,32]) + + + @patch("deepmd.entrypoints.train.get_sel") + def test_update_sel_hybrid(self, sel_mock): + sel_mock.return_value = [10,20] + jdata = { + 'model' : { + 'descriptor': { + 'type' : 'hybrid', + 'list' : [ + { + 'rcut': 6, + 'sel': "auto" + }, + { + 'rcut': 6, + 'sel': "auto:1.5" + } + ] + } + } + } + expected_out = { + 'model' : { + 'descriptor': { + 'type' : 'hybrid', + 'list' : [ + { + 'rcut': 6, + 'sel': [12,24] + }, + { + 'rcut': 6, + 'sel': [16,32] + } + ] + } + } + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + + @patch("deepmd.entrypoints.train.get_sel") + def test_update_sel(self, sel_mock): + sel_mock.return_value = [10,20] + jdata = { + 'model' : { + 'descriptor': { + 'type' : 'se_e2_a', + 'rcut': 6, + 'sel': "auto" + } + } + } + expected_out = { + 'model' : { + 'descriptor': { + 'type' : 'se_e2_a', + 'rcut': 6, + 'sel': [12,24] + } + } + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + + def test_wrap_up_4(self): + self.assertEqual(wrap_up_4(12), 3 * 4) + self.assertEqual(wrap_up_4(13), 4 * 4) + self.assertEqual(wrap_up_4(14), 4 * 4) + self.assertEqual(wrap_up_4(15), 4 * 4) + self.assertEqual(wrap_up_4(16), 4 * 4) + self.assertEqual(wrap_up_4(17), 5 * 4) +