handle extension of DP models correctly

Signed-off-by: zjgemi <[email protected]>

dpgen2/constants.py

@@ -2,8 +2,8 @@
 train_task_pattern = "task." + train_index_pattern
 train_script_name = "input.json"
 train_log_name = "train.log"
-model_name_pattern = "model.%03d.pb"
-model_name_match_pattern = r"model\.[0-9]{3,}\.pb"
+model_name_pattern = "model.%03d"
+model_name_match_pattern = r"model\.[0-9]{3,}"
 lmp_index_pattern = "%06d"
 lmp_task_pattern = "task." + lmp_index_pattern
 lmp_conf_name = "conf.lmp"

dpgen2/entrypoint/submit.py

@@ -522,7 +522,7 @@ def workflow_concurrent_learning(
 
     fp_config["inputs"] = fp_inputs
     fp_config["run"] = config["fp"]["run_config"]
-    if fp_style in ["deepmd", "deepmd_pt"]:
+    if fp_style == "deepmd":
         assert (
             "teacher_model_path" in fp_config["run"]
         ), f"Cannot find 'teacher_model_path' in config['fp']['run_config'] when fp_style == 'deepmd'"

dpgen2/fp/__init__.py

@@ -8,11 +8,6 @@
     PrepDeepmd,
     RunDeepmd,
 )
-from .deepmd_pt import (
-    DeepmdPtInputs,
-    PrepDeepmdPt,
-    RunDeepmdPt,
-)
 from .gaussian import (
     GaussianInputs,
     PrepGaussian,
@@ -40,11 +35,6 @@
         "prep": PrepDeepmd,
         "run": RunDeepmd,
     },
-    "deepmd_pt": {
-        "inputs": DeepmdPtInputs,
-        "prep": PrepDeepmdPt,
-        "run": RunDeepmdPt,
-    },
     "fpop_abacus": {
         "inputs": FpOpAbacusInputs,
         "prep": PrepFpOpAbacus,

dpgen2/fp/deepmd.py

@@ -45,9 +45,6 @@
 # global static variables
 deepmd_temp_path = "one_frame_temp"
 
-# global static variables
-deepmd_teacher_model = "teacher_model.pb"
-
 
 class DeepmdInputs:
     @staticmethod
@@ -136,6 +133,8 @@ def run_task(
     def _get_dp_model(self, teacher_model_path: BinaryFileInput):
         from deepmd.infer import DeepPot  # type: ignore
 
+        ext = os.path.splitext(teacher_model_path.file_name)[-1]
+        deepmd_teacher_model = "teacher_model" + ext
         teacher_model_path.save_as_file(deepmd_teacher_model)
         dp = DeepPot(Path(deepmd_teacher_model))
 

dpgen2/op/run_dp_train.py

@@ -179,7 +179,7 @@ def execute(
             return OPIO(
                 {
                     "script": work_dir / train_script_name,
-                    "model": work_dir / "frozen_model.pb",
+                    "model": init_model,
                     "lcurve": work_dir / "lcurve.out",
                     "log": work_dir / "train.log",
                 }
@@ -199,8 +199,8 @@ def clean_before_quit():
             # train model
             if impl == "tensorflow" and os.path.isfile("checkpoint"):
                 command = dp_command + ["train", "--restart", "model.ckpt", train_script_name]
-            elif impl == "pytorch" and len(glob.glob("model_[0-9]*.pt")) > 0:
-                checkpoint = "model_%s.pt" % max([int(f[6:-3]) for f in glob.glob("model_[0-9]*.pt")])
+            elif impl == "pytorch" and len(glob.glob("model.ckpt-[0-9]*.pt")) > 0:
+                checkpoint = "model.ckpt-%s.pt" % max([int(f[11:-3]) for f in glob.glob("model.ckpt-[0-9]*.pt")])
                 command = dp_command + ["train", "--restart", checkpoint, train_script_name]
             elif (do_init_model or finetune_mode == "train-init") and not config["init_model_with_finetune"]:
                 if impl == "tensorflow":
@@ -384,7 +384,6 @@ def skip_training(
                     f"The training is skipped.\n"
                 )
                 Path("lcurve.out").touch()
-                shutil.copy(init_model, "frozen_model.pb")
             return True
         else:
             return False