deepmd.pt.utils.finetune

Module Contents

Functions

change_finetune_model_params_single(...[, ...])

change_finetune_model_params(finetune_model, model_config)

Load model_params according to the pretrained one.

Attributes

log

deepmd.pt.utils.finetune.log[source]
deepmd.pt.utils.finetune.change_finetune_model_params_single(_single_param_target, _model_param_pretrained, from_multitask=False, model_branch='Default', model_branch_from='')[source]
deepmd.pt.utils.finetune.change_finetune_model_params(finetune_model, model_config, model_branch='')[source]

Load model_params according to the pretrained one. This function modifies the fine-tuning input in different modes as follows: 1. Single-task fine-tuning from a single-task pretrained model:

  • Updates the model parameters based on the pretrained model.

  1. Single-task fine-tuning from a multi-task pretrained model:

    • Updates the model parameters based on the selected branch in the pretrained model.

    • The chosen branch can be defined from the command-line or finetune_head input parameter.

    • If not defined, model parameters in the fitting network will be randomly initialized.

  2. Multi-task fine-tuning from a single-task pretrained model:

    • Updates model parameters in each branch based on the single branch (‘Default’) in the pretrained model.

    • If finetune_head is not set to ‘Default’, model parameters in the fitting network of the branch will be randomly initialized.

  3. Multi-task fine-tuning from a multi-task pretrained model:

    • Updates model parameters in each branch based on the selected branch in the pretrained model.

    • The chosen branches can be defined from the finetune_head input parameter of each model.

    • If finetune_head is not defined and the model_key is the same as in the pretrained model, it will resume from the model_key branch without fine-tuning.

    • If finetune_head is not defined and a new model_key is used, model parameters in the fitting network of the branch will be randomly initialized.

Parameters:
finetune_model

The pretrained model.

model_config

The fine-tuning input parameters.

model_branch

The model branch chosen in command-line mode, only for single-task fine-tuning.

Returns:
model_config:

Updated model parameters.

finetune_links:

Fine-tuning rules in a dict format, with model_branch: model_branch_from pairs. If model_key is not in this dict, it will do just resuming instead of fine-tuning.