tf: support checkpoint path (instead of directory) in dp freeze (#3254)

To have the same behavior between TF and PT.

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

deepmd/main.py

@@ -275,7 +275,7 @@ def main_parser() -> argparse.ArgumentParser:
         "--checkpoint",
         type=str,
         default=".",
-        help="Path to checkpoint. TensorFlow backend: a folder; PyTorch backend: either a folder containing checkpoint, or a pt file",
+        help="Path to checkpoint, either a folder containing checkpoint or the checkpoint prefix",
     )
     parser_frz.add_argument(
         "-o",

deepmd/tf/entrypoints/freeze.py

@@ -12,6 +12,9 @@
 from os.path import (
     abspath,
 )
+from pathlib import (
+    Path,
+)
 from typing import (
     List,
     Optional,
@@ -479,7 +482,7 @@ def freeze(
     Parameters
     ----------
     checkpoint_folder : str
-        location of the folder with model
+        location of either the folder with checkpoint or the checkpoint prefix
     output : str
         output file name
     node_names : Optional[str], optional
@@ -492,8 +495,11 @@ def freeze(
         other arguments
     """
     # We retrieve our checkpoint fullpath
-    checkpoint = tf.train.get_checkpoint_state(checkpoint_folder)
-    input_checkpoint = checkpoint.model_checkpoint_path
+    if Path(checkpoint_folder).is_dir():
+        checkpoint = tf.train.get_checkpoint_state(checkpoint_folder)
+        input_checkpoint = checkpoint.model_checkpoint_path
+    else:
+        input_checkpoint = checkpoint_folder
 
     # expand the output file to full path
     output_graph = abspath(output)