From 3b2c794b646f220c992bbd3b0c349139d34f9701 Mon Sep 17 00:00:00 2001 From: denghuilu Date: Wed, 11 Aug 2021 15:16:19 +0800 Subject: [PATCH 1/5] add model compression support for models without training script --- deepmd/entrypoints/compress.py | 25 ++++++++++++++++++------- deepmd/entrypoints/main.py | 7 +++++++ deepmd/entrypoints/train.py | 10 ++++++++-- 3 files changed, 33 insertions(+), 9 deletions(-) diff --git a/deepmd/entrypoints/compress.py b/deepmd/entrypoints/compress.py index b0f0f42729..9394f704a6 100644 --- a/deepmd/entrypoints/compress.py +++ b/deepmd/entrypoints/compress.py @@ -1,5 +1,6 @@ """Compress a model, which including tabulating the embedding-net.""" +import os import json import logging from typing import Optional @@ -11,7 +12,7 @@ from deepmd.utils.errors import GraphTooLargeError, GraphWithoutTensorError from .freeze import freeze -from .train import train +from .train import train, get_rcut, get_min_nbor_dist from .transfer import transfer __all__ = ["compress"] @@ -27,6 +28,7 @@ def compress( step: float, frequency: str, checkpoint_folder: str, + training_script: str, mpi_log: str, log_path: Optional[str], log_level: int, @@ -54,6 +56,8 @@ def compress( frequency of tabulation overflow check checkpoint_folder : str trining checkpoint folder for freezing + training_script : str + training script of the input frozen model mpi_log : str mpi logging mode for training log_path : Optional[str] @@ -64,16 +68,23 @@ def compress( try: t_jdata = get_tensor_by_name(input, 'train_attr/training_script') t_min_nbor_dist = get_tensor_by_name(input, 'train_attr/min_nbor_dist') + jdata = json.loads(t_jdata) except GraphWithoutTensorError as e: - raise RuntimeError( - "The input frozen model: %s has no training script or min_nbor_dist information," - "which is not supported by the model compression program." - "Please consider using the dp convert-from interface to upgrade the model" % input - ) from e + if os.path.exists(training_script) == False: + raise RuntimeError( + "The input frozen model: %s has no training script or min_nbor_dist information, " + "which is not supported by the model compression interface. " + "Please consider using the --training-script command within the model compression interface to provide the training script of the input frozen model. " + "Note that the input training script must contain the correct path to the training data." % input + ) from e + else: + log.info("stage 0: compute the min_nbor_dist") + jdata = j_loader(training_script) + t_min_nbor_dist = get_min_nbor_dist(jdata, get_rcut(jdata)) + tf.constant(t_min_nbor_dist, name = 'train_attr/min_nbor_dist', dtype = GLOBAL_TF_FLOAT_PRECISION) - jdata = json.loads(t_jdata) jdata["model"]["compress"] = {} jdata["model"]["compress"]["type"] = 'se_e2_a' jdata["model"]["compress"]["compress"] = True diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index d28022ac1f..10ae8ba5a0 100644 --- a/deepmd/entrypoints/main.py +++ b/deepmd/entrypoints/main.py @@ -306,6 +306,13 @@ def parse_args(args: Optional[List[str]] = None): default=".", help="path to checkpoint folder", ) + parser_compress.add_argument( + "-t", + "--training-script", + type=str, + default="input.json", + help="The training script of the input frozen model", + ) # * print docs script ************************************************************** parsers_doc = subparsers.add_parser( diff --git a/deepmd/entrypoints/train.py b/deepmd/entrypoints/train.py index ba20f5f7e8..17c0136a69 100755 --- a/deepmd/entrypoints/train.py +++ b/deepmd/entrypoints/train.py @@ -239,8 +239,7 @@ def get_rcut(jdata): def get_type_map(jdata): return jdata['model'].get('type_map', None) - -def get_sel(jdata, rcut): +def get_nbor_stat(jdata, rcut): max_rcut = get_rcut(jdata) type_map = get_type_map(jdata) @@ -258,9 +257,16 @@ def get_sel(jdata, rcut): neistat = NeighborStat(ntypes, rcut) min_nbor_dist, max_nbor_size = neistat.get_stat(train_data) + return min_nbor_dist, max_nbor_size +def get_sel(jdata, rcut): + _, max_nbor_size = get_nbor_stat(jdata, rcut) return max_nbor_size +def get_min_nbor_dist(jdata, rcut): + min_nbor_dist, _ = get_nbor_stat(jdata, rcut) + return min_nbor_dist + def parse_auto_sel(sel): if type(sel) is not str: From 65e8b8f387b726b667b83a8d19bc492fd6a5ac55 Mon Sep 17 00:00:00 2001 From: denghuilu Date: Wed, 11 Aug 2021 15:29:45 +0800 Subject: [PATCH 2/5] fix line changes --- deepmd/entrypoints/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/entrypoints/train.py b/deepmd/entrypoints/train.py index 17c0136a69..3fe0466eae 100755 --- a/deepmd/entrypoints/train.py +++ b/deepmd/entrypoints/train.py @@ -239,6 +239,7 @@ def get_rcut(jdata): def get_type_map(jdata): return jdata['model'].get('type_map', None) + def get_nbor_stat(jdata, rcut): max_rcut = get_rcut(jdata) type_map = get_type_map(jdata) @@ -267,7 +268,6 @@ def get_min_nbor_dist(jdata, rcut): min_nbor_dist, _ = get_nbor_stat(jdata, rcut) return min_nbor_dist - def parse_auto_sel(sel): if type(sel) is not str: return False From ea359e2e19939b925277317df8baa0571d8a78e4 Mon Sep 17 00:00:00 2001 From: Denghui Lu Date: Wed, 11 Aug 2021 15:50:23 +0800 Subject: [PATCH 3/5] update doc for model compression --- doc/getting-started.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/doc/getting-started.md b/doc/getting-started.md index 0965b72661..846485432e 100644 --- a/doc/getting-started.md +++ b/doc/getting-started.md @@ -336,6 +336,9 @@ optional arguments: (default: -1) -c CHECKPOINT_FOLDER, --checkpoint-folder CHECKPOINT_FOLDER path to checkpoint folder (default: .) + -t TRAINING_SCRIPT, --training-script TRAINING_SCRIPT + The training script of the input frozen model + (default: input.json) ``` **Parameter explanation** From 11d50cc7569596a2381b76971c0ec7aafebd8d6b Mon Sep 17 00:00:00 2001 From: denghuilu Date: Wed, 11 Aug 2021 17:02:27 +0800 Subject: [PATCH 4/5] make the error message more reasonable --- deepmd/entrypoints/compress.py | 6 +++++- deepmd/entrypoints/main.py | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/deepmd/entrypoints/compress.py b/deepmd/entrypoints/compress.py index 9394f704a6..7755156601 100644 --- a/deepmd/entrypoints/compress.py +++ b/deepmd/entrypoints/compress.py @@ -70,13 +70,17 @@ def compress( t_min_nbor_dist = get_tensor_by_name(input, 'train_attr/min_nbor_dist') jdata = json.loads(t_jdata) except GraphWithoutTensorError as e: - if os.path.exists(training_script) == False: + if training_script == None: raise RuntimeError( "The input frozen model: %s has no training script or min_nbor_dist information, " "which is not supported by the model compression interface. " "Please consider using the --training-script command within the model compression interface to provide the training script of the input frozen model. " "Note that the input training script must contain the correct path to the training data." % input ) from e + elif os.path.exists(training_script) == False: + raise RuntimeError( + "The input training script %s does not exist! Please check the path of the training script. " % (input + "(" + os.path.abspath(input) + ")") + ) from e else: log.info("stage 0: compute the min_nbor_dist") jdata = j_loader(training_script) diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index 10ae8ba5a0..52e4e1c61e 100644 --- a/deepmd/entrypoints/main.py +++ b/deepmd/entrypoints/main.py @@ -310,7 +310,7 @@ def parse_args(args: Optional[List[str]] = None): "-t", "--training-script", type=str, - default="input.json", + default=None, help="The training script of the input frozen model", ) From 34796941cd8ce794b8f6f645791241a1a535eb3d Mon Sep 17 00:00:00 2001 From: Denghui Lu Date: Wed, 11 Aug 2021 17:07:23 +0800 Subject: [PATCH 5/5] update doc for model compression --- doc/getting-started.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/getting-started.md b/doc/getting-started.md index 846485432e..607e15ed07 100644 --- a/doc/getting-started.md +++ b/doc/getting-started.md @@ -338,7 +338,7 @@ optional arguments: path to checkpoint folder (default: .) -t TRAINING_SCRIPT, --training-script TRAINING_SCRIPT The training script of the input frozen model - (default: input.json) + (default: None) ``` **Parameter explanation**