Add pt compress commad line

deepmd/main.py

@@ -424,29 +424,30 @@ def main_parser() -> argparse.ArgumentParser:
     parser_compress = subparsers.add_parser(
         "compress",
         parents=[parser_log, parser_mpi_log],
-        help="(Supported backend: TensorFlow) compress a model",
+        help="Compress a model",
         formatter_class=RawTextArgumentDefaultsHelpFormatter,
         epilog=textwrap.dedent(
             """\
         examples:
             dp compress
             dp compress -i graph.pb -o compressed.pb
+            dp --pt compress -i model.pth -o compressed.pth
         """
         ),
     )
     parser_compress.add_argument(
         "-i",
         "--input",
-        default="frozen_model.pb",
+        default="frozen_model",
         type=str,
-        help="The original frozen model, which will be compressed by the code",
+        help="The original frozen model, which will be compressed by the code. Filename (prefix) of the input model file. TensorFlow backend: suffix is .pb; PyTorch backend: suffix is .pth",
     )
     parser_compress.add_argument(
         "-o",
         "--output",
-        default="frozen_model_compressed.pb",
+        default="frozen_model_compressed",
         type=str,
-        help="The compressed model",
+        help="The compressed model. Filename (prefix) of the output model file. TensorFlow backend: suffix is .pb; PyTorch backend: suffix is .pth",
     )
     parser_compress.add_argument(
         "-s",

deepmd/pt/entrypoints/compress.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import json
+
+import torch
+
+from deepmd.pt.model.model import (
+    get_model,
+)
+
+
+def enable_compression(
+    input_file: str,
+    output: str,
+    stride: float = 0.01,
+    extrapolate: int = 5,
+    check_frequency: int = -1,
+):
+    saved_model = torch.jit.load(input_file, map_location="cpu")
+    model_def_script = json.loads(saved_model.model_def_script)
+    model = get_model(model_def_script)
+    model.load_state_dict(saved_model.state_dict())
+
+    model.enable_compression(
+        extrapolate,
+        stride,
+        stride * 10,
+        check_frequency,
+    )
+
+    model = torch.jit.script(model)
+    torch.jit.save(model, output)

deepmd/pt/entrypoints/main.py

@@ -38,6 +38,9 @@
 from deepmd.pt.cxx_op import (
     ENABLE_CUSTOMIZED_OP,
 )
+from deepmd.pt.entrypoints.compress import (
+    enable_compression,
+)
 from deepmd.pt.infer import (
     inference,
 )
@@ -549,6 +552,16 @@ def main(args: Optional[Union[list[str], argparse.Namespace]] = None):
             model_branch=FLAGS.model_branch,
             output=FLAGS.output,
         )
+    elif FLAGS.command == "compress":
+        FLAGS.input = str(Path(FLAGS.input).with_suffix(".pth"))
+        FLAGS.output = str(Path(FLAGS.output).with_suffix(".pth"))
+        enable_compression(
+            input_file=FLAGS.input,
+            output=FLAGS.output,
+            stride=FLAGS.step,
+            extrapolate=FLAGS.extrapolate,
+            check_frequency=FLAGS.frequency,
+        )
     else:
         raise RuntimeError(f"Invalid command {FLAGS.command}!")
 

deepmd/tf/entrypoints/main.py

@@ -77,6 +77,12 @@ def main(args: Optional[Union[list[str], argparse.Namespace]] = None):
     elif args.command == "transfer":
         transfer(**dict_args)
     elif args.command == "compress":
+        dict_args["input"] = format_model_suffix(
+            dict_args["input"], preferred_backend=args.backend, strict_prefer=True
+        )
+        dict_args["output"] = format_model_suffix(
+            dict_args["output"], preferred_backend=args.backend, strict_prefer=True
+        )
         compress(**dict_args)
     elif args.command == "convert-from":
         convert(**dict_args)

source/tests/pt/common.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import pathlib
 from typing import (
     Optional,
     Union,
@@ -7,6 +8,7 @@
 import numpy as np
 import torch
 
+from deepmd.common import j_loader as dp_j_loader
 from deepmd.main import (
     main,
 )
@@ -15,6 +17,12 @@
     GLOBAL_PT_FLOAT_PRECISION,
 )
 
+tests_path = pathlib.Path(__file__).parent.absolute()
+
+
+def j_loader(filename):
+    return dp_j_loader(tests_path / filename)
+
 
 def run_dp(cmd: str) -> int:
     """Run DP directly from the entry point instead of the subprocess.

source/tests/pt/model_compression/data/type.raw

@@ -0,0 +1,6 @@
+0
+1
+1
+0
+1
+1

source/tests/pt/model_compression/data/type_map.raw

@@ -0,0 +1,2 @@
+O
+H

source/tests/pt/model_compression/input.json

@@ -0,0 +1,85 @@
+{
+  "_comment1": " model parameters",
+  "model": {
+    "type_map": [
+      "O",
+      "H"
+    ],
+    "descriptor": {
+      "type": "se_e2_a",
+      "sel": [
+        46,
+        92
+      ],
+      "rcut_smth": 0.50,
+      "rcut": 6.00,
+      "_comment": "N2=2N1, N2=N1, and otherwise can be tested",
+      "neuron": [
+        4,
+        8,
+        17,
+        17
+      ],
+      "resnet_dt": false,
+      "axis_neuron": 16,
+      "seed": 1,
+      "_comment2": " that's all"
+    },
+    "fitting_net": {
+      "neuron": [
+        20,
+        20,
+        20
+      ],
+      "resnet_dt": true,
+      "seed": 1,
+      "_comment3": " that's all"
+    },
+    "_comment4": " that's all"
+  },
+
+  "learning_rate": {
+    "type": "exp",
+    "decay_steps": 5000,
+    "start_lr": 0.001,
+    "stop_lr": 3.51e-8,
+    "_comment5": "that's all"
+  },
+
+  "loss": {
+    "type": "ener",
+    "start_pref_e": 0.02,
+    "limit_pref_e": 1,
+    "start_pref_f": 1000,
+    "limit_pref_f": 1,
+    "start_pref_v": 0,
+    "limit_pref_v": 0,
+    "_comment6": " that's all"
+  },
+
+  "training": {
+    "training_data": {
+      "systems": [
+        "model_compression/data"
+      ],
+      "batch_size": "auto",
+      "_comment7": "that's all"
+    },
+    "validation_data": {
+      "systems": [
+        "model_compression/data"
+      ],
+      "batch_size": 1,
+      "numb_btch": 3,
+      "_comment8": "that's all"
+    },
+    "numb_steps": 1,
+    "seed": 10,
+    "disp_file": "lcurve.out",
+    "disp_freq": 1,
+    "save_freq": 1,
+    "_comment9": "that's all"
+  },
+
+  "_comment10": "that's all"
+}