pt: support --init-frz-model

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/main.py

@@ -226,7 +226,7 @@ def main_parser() -> argparse.ArgumentParser:
         "--init-frz-model",
         type=str,
         default=None,
-        help="(Supported backend: TensorFlow) Initialize the training from the frozen model.",
+        help="Initialize the training from the frozen model.",
     )
     parser_train_subgroup.add_argument(
         "-t",

deepmd/pt/entrypoints/main.py

@@ -65,6 +65,7 @@ def get_trainer(
     finetune_model=None,
     model_branch="",
     force_load=False,
+    init_frz_model=None,
 ):
     # Initialize DDP
     local_rank = os.environ.get("LOCAL_RANK")
@@ -200,6 +201,7 @@ def prepare_trainer_input_single(
         finetune_model=finetune_model,
         force_load=force_load,
         shared_links=shared_links,
+        init_frz_model=init_frz_model,
     )
     return trainer
 
@@ -243,6 +245,7 @@ def train(FLAGS):
         FLAGS.finetune,
         FLAGS.model_branch,
         FLAGS.force_load,
+        FLAGS.init_frz_model,
     )
     trainer.run()
 

deepmd/pt/train/training.py

@@ -75,6 +75,7 @@ def __init__(
         finetune_model=None,
         force_load=False,
         shared_links=None,
+        init_frz_model=None,
     ):
         """Construct a DeePMD trainer.
 
@@ -394,6 +395,9 @@ def get_loss(loss_params, start_lr, _ntypes):
                         ntest=ntest,
                         bias_shift=model_params.get("bias_shift", "delta"),
                     )
+        if init_frz_model is not None:
+            frz_model = torch.jit.load(init_frz_model, map_location=DEVICE)
+            self.model.load_state_dict(frz_model.state_dict())
 
         # Set trainable params
         self.wrapper.set_trainable_params()