feat(pt): Add command to check the available model branches in multi-task pre-trained model(Issue #3742)

deepmd/main.py

@@ -274,6 +274,11 @@ def main_parser() -> argparse.ArgumentParser:
         default="",
         help="(Supported backend: PyTorch) Model branch chosen for fine-tuning if multi-task. If not specified, it will re-init the fitting net.",
     )
+    parser_train.add_argument(
+        "--list-model-branch",
+        action="store_true",
+        help="(Supported backend: PyTorch) List model branches of a multitask pretrained model.",
+    )
     parser_train.add_argument(
         "--force-load",
         action="store_true",

deepmd/pt/entrypoints/main.py

@@ -44,6 +44,9 @@
 from deepmd.pt.train import (
     training,
 )
+from deepmd.pt.utils import (
+    env,
+)
 from deepmd.pt.utils.dataloader import (
     DpLoaderSet,
 )
@@ -309,7 +312,21 @@
     log.info("DeepMD version: %s", __version__)
 
     if FLAGS.command == "train":
-        train(FLAGS)
+        if FLAGS.list_model_branch:
+            assert (
+                FLAGS.finetune is not None
+            ), "Error: The '--list-model-branch' option requires the '--finetune' argument to specify the model."
+            state_dict = torch.load(FLAGS.finetune, map_location=env.DEVICE)
+            if "model" in state_dict:
+                state_dict = state_dict["model"]
+            model_params = state_dict["_extra_state"]["model_params"]
+            finetune_from_multi_task = "model_dict" in model_params
+            #  Pretrained model must be multitask mode
+            assert finetune_from_multi_task, "Error: The '--list-model-branch' option requires a multitask pretrained model. The provided model does not meet this criterion."
+            model_branch = list(model_params["model_dict"].keys())
+            log.info(f"Available model branches are {model_branch}")
+        else:
+            train(FLAGS)
     elif FLAGS.command == "freeze":
         if Path(FLAGS.checkpoint_folder).is_dir():
             checkpoint_path = Path(FLAGS.checkpoint_folder)

doc/train/finetuning.md

@@ -99,11 +99,12 @@ $ dp --pt train input.json --finetune multitask_pretrained.pt --model-branch CHO
 ```
 
 :::{note}
-To check the available model branches, you can typically refer to the documentation of the pre-trained model.
-If you're still unsure about the available branches, you can try inputting an arbitrary branch name.
-This will prompt an error message that displays a list of all the available model branches.
+One can check the available model branches in multi-task pre-trained model by refering to the documentation of the pre-trained model or by using the following command:
+
+```bash
+$ dp --pt train input.json --finetune multitask_pretrained.pt --list-model-branch
+```
 
-Please note that this feature will be improved in the upcoming version to provide a more user-friendly experience.
 :::
 
 This command will start fine-tuning based on the pre-trained model's descriptor and the selected branch's fitting net.