deepmd.pt.utils.multi_task

deepmd.pt.utils.multi_task#

Functions#

`_cascade_top_level_defaults`(→ None)	In-place: lower model-wide `model.<key>` entries into each branch.
`preprocess_shared_params`(→ tuple[dict[str, Any], ...)	Preprocess the model params for multitask model, and generate the links dict for further sharing.
`get_class_name`(→ type)

Module Contents#

deepmd.pt.utils.multi_task._cascade_top_level_defaults(model_config: dict[str, Any]) → None[source]#

In-place: lower model-wide model.<key> entries into each branch.

Any key at the top of model other than model_dict / shared_dict is setdefault-copied into every model_dict entry (explicit branch values win) and then removed from the top level so the multi-task argcheck, which only accepts model_dict / shared_dict there, does not reject it as an unknown field.

deepmd.pt.utils.multi_task.preprocess_shared_params(model_config: dict[str, Any]) → tuple[dict[str, Any], dict[str, Any]][source]#

Preprocess the model params for multitask model, and generate the links dict for further sharing.

Parameters:

model_config – Model params of multitask model.

Returns:

model_config: Preprocessed model params of multitask model.

Those string names are replaced with real params in shared_dict of model params.

shared_links: Dict of link infos for further sharing.

Each item, whose key must be in shared_dict, is a dict with following keys: - “type”: The real class type of this item. - “links”: List of shared settings, each sub-item is a dict with following keys:

“model_key”: Model key in the model_dict to share this item.
“shared_type”: Type of this shard item.
“shared_level”: Shared level (int) of this item in this model.
Lower for more params to share, 0 means to share all params in this item.
This list are sorted by “shared_level”.

For example, if one has model_config like this:

“model”: {

“shared_dict”: {: “my_type_map”: [“foo”, “bar”], “my_des1”: {

“type”: “se_e2_a”, “neuron”: [10, 20, 40] },

}, “model_dict”: {

“model_1”: {
“type_map”: “my_type_map”, “descriptor”: “my_des1”, “fitting_net”: {

“neuron”: [100, 100, 100]

}
}, “model_2”: {
“type_map”: “my_type_map”, “descriptor”: “my_des1”, “fitting_net”: {
“neuron”: [100, 100, 100]
}
} “model_3”: {
“type_map”: “my_type_map”, “descriptor”: “my_des1:1”, “fitting_net”: {
“neuron”: [100, 100, 100]
}
}

}

}

The above config will init three model branches named model_1 and model_2 and model_3,

in which:

model_2 and model_3 will have the same type_map as that in model_1.
model_2 will share all the parameters of descriptor with model_1,

while model_3 will share part of parameters of descriptor with model_1 on human-defined share-level 1 (default is 0, meaning share all the parameters). - model_1, model_2 and model_3 have three different `fitting_net`s.

The returned model_config will automatically fulfill the input model_config as if there’s no sharing,

and the shared_links will keep all the sharing information with looking:

{

‘my_des1’: {

‘type’: ‘DescrptSeA’, ‘links’: [

{‘model_key’: ‘model_1’, ‘shared_type’: ‘descriptor’, ‘shared_level’: 0}, {‘model_key’: ‘model_2’, ‘shared_type’: ‘descriptor’, ‘shared_level’: 0}, {‘model_key’: ‘model_3’, ‘shared_type’: ‘descriptor’, ‘shared_level’: 1} ]

}

}
Any key placed directly under model other than model_dict /
shared_dict is lowered into every branch via _cascade_top_level_defaults
(explicit branch values win), so model-wide switches can be written
once at the top level.

deepmd.pt.utils.multi_task.get_class_name(item_key: str, item_params: dict[str, Any]) → type[source]#