add option to select backends TF/PT (#1545)

reopen PR #1541 due to branch is deleted

add a new key in `param.json` file

```
"train_backend": "pytorch"/"tensorflow",
```
relate to this issue #1462

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit


- **New Features**
- Improved model management by dynamically generating model suffixes
based on the selected backend, enhancing compatibility.
  
- **Enhancements**
- Updated model-related functions to incorporate backend-specific model
suffixes for accurate file handling during training processes.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: C. Thang Nguyen <[email protected]>
Signed-off-by: Jinzhe Zeng <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jinzhe Zeng <[email protected]>

Author: 3 people

dpgen/generator/arginfo.py

@@ -87,6 +87,9 @@ def training_args() -> list[Argument]:
     list[dargs.Argument]
         List of training arguments.
     """
+    doc_train_backend = (
+        "The backend of the training. Currently only support tensorflow and pytorch."
+    )
     doc_numb_models = "Number of models to be trained in 00.train. 4 is recommend."
     doc_training_iter0_model_path = "The model used to init the first iter training. Number of element should be equal to numb_models."
     doc_training_init_model = "Iteration > 0, the model parameters will be initilized from the model trained at the previous iteration. Iteration == 0, the model parameters will be initialized from training_iter0_model_path."
@@ -123,6 +126,13 @@ def training_args() -> list[Argument]:
     doc_training_finetune_model = "At interation 0, finetune the model parameters from the given frozen models. Number of element should be equal to numb_models."
 
     return [
+        Argument(
+            "train_backend",
+            str,
+            optional=True,
+            default="tensorflow",
+            doc=doc_train_backend,
+        ),
         Argument("numb_models", int, optional=False, doc=doc_numb_models),
         Argument(
             "training_iter0_model_path",

dpgen/generator/run.py

@@ -125,6 +125,19 @@
 run_opt_file = os.path.join(ROOT_PATH, "generator/lib/calypso_run_opt.py")
 
 
+def _get_model_suffix(jdata) -> str:
+    """Return the model suffix based on the backend."""
+    suffix_map = {"tensorflow": ".pb", "pytorch": ".pth"}
+    backend = jdata.get("train_backend", "tensorflow")
+    if backend in suffix_map:
+        suffix = suffix_map[backend]
+    else:
+        raise ValueError(
+            f"The backend {backend} is not available. Supported backends are: 'tensorflow', 'pytorch'."
+        )
+    return suffix
+
+
 def get_job_names(jdata):
     jobkeys = []
     for ii in jdata.keys():
@@ -172,7 +185,7 @@ def _check_empty_iter(iter_index, max_v=0):
     return all(empty_sys)
 
 
-def copy_model(numb_model, prv_iter_index, cur_iter_index):
+def copy_model(numb_model, prv_iter_index, cur_iter_index, suffix=".pb"):
     cwd = os.getcwd()
     prv_train_path = os.path.join(make_iter_name(prv_iter_index), train_name)
     cur_train_path = os.path.join(make_iter_name(cur_iter_index), train_name)
@@ -184,7 +197,8 @@ def copy_model(numb_model, prv_iter_index, cur_iter_index):
         os.chdir(cur_train_path)
         os.symlink(os.path.relpath(prv_train_task), train_task_fmt % ii)
         os.symlink(
-            os.path.join(train_task_fmt % ii, "frozen_model.pb"), "graph.%03d.pb" % ii
+            os.path.join(train_task_fmt % ii, f"frozen_model{suffix}"),
+            "graph.%03d%s" % (ii, suffix),
         )
         os.chdir(cwd)
     with open(os.path.join(cur_train_path, "copied"), "w") as fp:
@@ -315,18 +329,19 @@ def make_train(iter_index, jdata, mdata):
         number_old_frames = 0
         number_new_frames = 0
 
+    suffix = _get_model_suffix(jdata)
     model_devi_engine = jdata.get("model_devi_engine", "lammps")
     if iter_index > 0 and _check_empty_iter(iter_index - 1, fp_task_min):
         log_task("prev data is empty, copy prev model")
-        copy_model(numb_models, iter_index - 1, iter_index)
+        copy_model(numb_models, iter_index - 1, iter_index, suffix)
         return
     elif (
         model_devi_engine != "calypso"
         and iter_index > 0
         and _check_skip_train(model_devi_jobs[iter_index - 1])
     ):
         log_task("skip training at step %d " % (iter_index - 1))
-        copy_model(numb_models, iter_index - 1, iter_index)
+        copy_model(numb_models, iter_index - 1, iter_index, suffix)
         return
     else:
         iter_name = make_iter_name(iter_index)
@@ -647,7 +662,9 @@ def make_train(iter_index, jdata, mdata):
     )
     if copied_models is not None:
         for ii in range(len(copied_models)):
-            _link_old_models(work_path, [copied_models[ii]], ii, basename="init.pb")
+            _link_old_models(
+                work_path, [copied_models[ii]], ii, basename=f"init{suffix}"
+            )
     # Copy user defined forward files
     symlink_user_forward_files(mdata=mdata, task_type="train", work_path=work_path)
     # HDF5 format for training data
@@ -699,6 +716,7 @@ def run_train(iter_index, jdata, mdata):
     # print("debug:run_train:mdata", mdata)
     # load json param
     numb_models = jdata["numb_models"]
+    suffix = _get_model_suffix(jdata)
     # train_param = jdata['train_param']
     train_input_file = default_train_input_file
     training_reuse_iter = jdata.get("training_reuse_iter")
@@ -730,7 +748,11 @@ def run_train(iter_index, jdata, mdata):
             "training_init_model, training_init_frozen_model, and training_finetune_model are mutually exclusive."
         )
 
-    train_command = mdata.get("train_command", "dp")
+    train_command = mdata.get("train_command", "dp").strip()
+    # assert train_command == "dp", "The 'train_command' should be 'dp'"     # the tests should be updated to run this command
+    if suffix == ".pth":
+        train_command += " --pt"
+
     train_resources = mdata["train_resources"]
 
     # paths
@@ -761,9 +783,9 @@ def run_train(iter_index, jdata, mdata):
         if training_init_model:
             init_flag = " --init-model old/model.ckpt"
         elif training_init_frozen_model is not None:
-            init_flag = " --init-frz-model old/init.pb"
+            init_flag = f" --init-frz-model old/init{suffix}"
         elif training_finetune_model is not None:
-            init_flag = " --finetune old/init.pb"
+            init_flag = f" --finetune old/init{suffix}"
         command = f"{train_command} train {train_input_file}{extra_flags}"
         command = f"{{ if [ ! -f model.ckpt.index ]; then {command}{init_flag}; else {command} --restart model.ckpt; fi }}"
         command = f"/bin/sh -c {shlex.quote(command)}"
@@ -792,23 +814,35 @@ def run_train(iter_index, jdata, mdata):
     if "srtab_file_path" in jdata.keys():
         forward_files.append(zbl_file)
     if training_init_model:
-        forward_files += [
-            os.path.join("old", "model.ckpt.meta"),
-            os.path.join("old", "model.ckpt.index"),
-            os.path.join("old", "model.ckpt.data-00000-of-00001"),
-        ]
+        if suffix == ".pb":
+            forward_files += [
+                os.path.join("old", "model.ckpt.meta"),
+                os.path.join("old", "model.ckpt.index"),
+                os.path.join("old", "model.ckpt.data-00000-of-00001"),
+            ]
+        elif suffix == ".pth":
+            forward_files += [os.path.join("old", "model.ckpt.pt")]
     elif training_init_frozen_model is not None or training_finetune_model is not None:
-        forward_files.append(os.path.join("old", "init.pb"))
+        forward_files.append(os.path.join("old", f"init{suffix}"))
 
-    backward_files = ["frozen_model.pb", "lcurve.out", "train.log"]
-    backward_files += [
-        "model.ckpt.meta",
-        "model.ckpt.index",
-        "model.ckpt.data-00000-of-00001",
+    backward_files = [
+        f"frozen_model{suffix}",
+        "lcurve.out",
+        "train.log",
         "checkpoint",
     ]
     if jdata.get("dp_compress", False):
-        backward_files.append("frozen_model_compressed.pb")
+        backward_files.append(f"frozen_model_compressed{suffix}")
+
+    if suffix == ".pb":
+        backward_files += [
+            "model.ckpt.meta",
+            "model.ckpt.index",
+            "model.ckpt.data-00000-of-00001",
+        ]
+    elif suffix == ".pth":
+        backward_files += ["model.ckpt.pt"]
+
     if not jdata.get("one_h5", False):
         init_data_sys_ = jdata["init_data_sys"]
         init_data_sys = []
@@ -879,13 +913,14 @@ def post_train(iter_index, jdata, mdata):
         log_task("copied model, do not post train")
         return
     # symlink models
+    suffix = _get_model_suffix(jdata)
     for ii in range(numb_models):
-        if not jdata.get("dp_compress", False):
-            model_name = "frozen_model.pb"
-        else:
-            model_name = "frozen_model_compressed.pb"
+        model_name = f"frozen_model{suffix}"
+        if jdata.get("dp_compress", False):
+            model_name = f"frozen_model_compressed{suffix}"
+
+        ofile = os.path.join(work_path, "graph.%03d%s" % (ii, suffix))
         task_file = os.path.join(train_task_fmt % ii, model_name)
-        ofile = os.path.join(work_path, "graph.%03d.pb" % ii)
         if os.path.isfile(ofile):
             os.remove(ofile)
         os.symlink(task_file, ofile)
@@ -1124,7 +1159,8 @@ def make_model_devi(iter_index, jdata, mdata):
     iter_name = make_iter_name(iter_index)
     train_path = os.path.join(iter_name, train_name)
     train_path = os.path.abspath(train_path)
-    models = sorted(glob.glob(os.path.join(train_path, "graph*pb")))
+    suffix = _get_model_suffix(jdata)
+    models = sorted(glob.glob(os.path.join(train_path, f"graph*{suffix}")))
     work_path = os.path.join(iter_name, model_devi_name)
     create_path(work_path)
     if model_devi_engine == "calypso":
@@ -1305,7 +1341,8 @@ def _make_model_devi_revmat(iter_index, jdata, mdata, conf_systems):
     iter_name = make_iter_name(iter_index)
     train_path = os.path.join(iter_name, train_name)
     train_path = os.path.abspath(train_path)
-    models = sorted(glob.glob(os.path.join(train_path, "graph*pb")))
+    suffix = _get_model_suffix(jdata)
+    models = sorted(glob.glob(os.path.join(train_path, f"graph*{suffix}")))
     task_model_list = []
     for ii in models:
         task_model_list.append(os.path.join("..", os.path.basename(ii)))
@@ -1502,7 +1539,8 @@ def _make_model_devi_native(iter_index, jdata, mdata, conf_systems):
     iter_name = make_iter_name(iter_index)
     train_path = os.path.join(iter_name, train_name)
     train_path = os.path.abspath(train_path)
-    models = glob.glob(os.path.join(train_path, "graph*pb"))
+    suffix = _get_model_suffix(jdata)
+    models = sorted(glob.glob(os.path.join(train_path, f"graph*{suffix}")))
     task_model_list = []
     for ii in models:
         task_model_list.append(os.path.join("..", os.path.basename(ii)))
@@ -1644,7 +1682,8 @@ def _make_model_devi_native_gromacs(iter_index, jdata, mdata, conf_systems):
     iter_name = make_iter_name(iter_index)
     train_path = os.path.join(iter_name, train_name)
     train_path = os.path.abspath(train_path)
-    models = glob.glob(os.path.join(train_path, "graph*pb"))
+    suffix = _get_model_suffix(jdata)
+    models = sorted(glob.glob(os.path.join(train_path, f"graph*{suffix}")))
     task_model_list = []
     for ii in models:
         task_model_list.append(os.path.join("..", os.path.basename(ii)))
@@ -1827,7 +1866,8 @@ def _make_model_devi_amber(
                 .replace("@qm_theory@", jdata["low_level"])
                 .replace("@rcut@", str(jdata["cutoff"]))
             )
-            models = sorted(glob.glob(os.path.join(train_path, "graph.*.pb")))
+            suffix = _get_model_suffix(jdata)
+            models = sorted(glob.glob(os.path.join(train_path, f"graph.*{suffix}")))
             task_model_list = []
             for ii in models:
                 task_model_list.append(os.path.join("..", os.path.basename(ii)))
@@ -1935,7 +1975,9 @@ def run_md_model_devi(iter_index, jdata, mdata):
     run_tasks = [os.path.basename(ii) for ii in run_tasks_]
     # dlog.info("all_task is ", all_task)
     # dlog.info("run_tasks in run_model_deviation",run_tasks_)
-    all_models = glob.glob(os.path.join(work_path, "graph*pb"))
+
+    suffix = _get_model_suffix(jdata)
+    all_models = glob.glob(os.path.join(work_path, f"graph*{suffix}"))
     model_names = [os.path.basename(ii) for ii in all_models]
 
     model_devi_engine = jdata.get("model_devi_engine", "lammps")

dpgen/simplify/simplify.py

@@ -31,6 +31,7 @@
     record_iter,
 )
 from dpgen.generator.run import (
+    _get_model_suffix,
     data_system_fmt,
     fp_name,
     fp_task_fmt,
@@ -186,7 +187,9 @@ def make_model_devi(iter_index, jdata, mdata):
     # link the model
     train_path = os.path.join(iter_name, train_name)
     train_path = os.path.abspath(train_path)
-    models = glob.glob(os.path.join(train_path, "graph*pb"))
+    suffix = _get_model_suffix(jdata)
+    models = glob.glob(os.path.join(train_path, f"graph*{suffix}"))
+
     for mm in models:
         model_name = os.path.basename(mm)
         os.symlink(mm, os.path.join(work_path, model_name))