feat(pt): support fine-tuning from random fitting

deepmd/pt/utils/finetune.py

@@ -30,9 +30,12 @@ def get_finetune_rule_single(
 
     if not from_multitask:
         single_config_chosen = deepcopy(_model_param_pretrained)
+        if model_branch_from == "RANDOM":
+            # not ["", "RANDOM"], because single-from-single finetune uses pretrained fitting in default
+            new_fitting = True
     else:
         model_dict_params = _model_param_pretrained["model_dict"]
-        if model_branch_from == "":
+        if model_branch_from in ["", "RANDOM"]:
             model_branch_chosen = next(iter(model_dict_params.keys()))
             new_fitting = True
             log.warning(
@@ -164,21 +167,27 @@ def get_finetune_rules(
             pretrained_keys = last_model_params["model_dict"].keys()
         for model_key in target_keys:
             resuming = False
-            if "finetune_head" in model_config["model_dict"][model_key]:
+            if (
+                "finetune_head" in model_config["model_dict"][model_key]
+                and model_config["model_dict"][model_key]["finetune_head"] != "RANDOM"
+            ):
                 pretrained_key = model_config["model_dict"][model_key]["finetune_head"]
                 assert pretrained_key in pretrained_keys, (
                     f"'{pretrained_key}' head chosen to finetune not exist in the pretrained model!"
                     f"Available heads are: {list(pretrained_keys)}"
                 )
                 model_branch_from = pretrained_key
-            elif model_key in pretrained_keys:
+            elif (
+                "finetune_head" not in model_config["model_dict"][model_key]
+                and model_key in pretrained_keys
+            ):
                 # not do anything if not defined "finetune_head" in heads that exist in the pretrained model
                 # this will just do resuming
                 model_branch_from = model_key
                 resuming = True
             else:
-                # if not defined "finetune_head" in new heads, the fitting net will bre randomly initialized
-                model_branch_from = ""
+                # if not defined "finetune_head" in new heads or "finetune_head" is "RANDOM", the fitting net will bre randomly initialized
+                model_branch_from = "RANDOM"
             model_config["model_dict"][model_key], finetune_rule = (
                 get_finetune_rule_single(
                     model_config["model_dict"][model_key],

source/tests/pt/test_training.py

@@ -34,6 +34,7 @@ def test_dp_train(self):
         # test training from scratch
         trainer = get_trainer(deepcopy(self.config))
         trainer.run()
+        state_dict_trained = trainer.wrapper.model.state_dict()
 
         # test fine-tuning using same input
         finetune_model = self.config["training"].get("save_ckpt", "model.ckpt") + ".pt"
@@ -46,7 +47,6 @@ def test_dp_train(self):
             finetune_model=finetune_model,
             finetune_links=finetune_links,
         )
-        trainer_finetune.run()
 
         # test fine-tuning using empty input
         self.config_empty = deepcopy(self.config)
@@ -64,7 +64,41 @@ def test_dp_train(self):
             finetune_model=finetune_model,
             finetune_links=finetune_links,
         )
+
+        # test fine-tuning using random fitting
+        self.config["model"], finetune_links = get_finetune_rules(
+            finetune_model, self.config["model"], model_branch="RANDOM"
+        )
+        trainer_finetune_random = get_trainer(
+            deepcopy(self.config_empty),
+            finetune_model=finetune_model,
+            finetune_links=finetune_links,
+        )
+
+        # check parameters
+        state_dict_finetuned = trainer_finetune.wrapper.model.state_dict()
+        state_dict_finetuned_empty = trainer_finetune_empty.wrapper.model.state_dict()
+        state_dict_finetuned_random = trainer_finetune_random.wrapper.model.state_dict()
+        for state_key in state_dict_finetuned:
+            if "out_bias" not in state_key and "out_std" not in state_key:
+                torch.testing.assert_close(
+                    state_dict_trained[state_key],
+                    state_dict_finetuned[state_key],
+                )
+                torch.testing.assert_close(
+                    state_dict_trained[state_key],
+                    state_dict_finetuned_empty[state_key],
+                )
+                if "fitting_net" not in state_key:
+                    torch.testing.assert_close(
+                        state_dict_trained[state_key],
+                        state_dict_finetuned_random[state_key],
+                    )
+
+        # check running
+        trainer_finetune.run()
         trainer_finetune_empty.run()
+        trainer_finetune_random.run()
 
         self.tearDown()