diff --git a/deepmd/common.py b/deepmd/common.py index 5d23986b44..07232e8b47 100644 --- a/deepmd/common.py +++ b/deepmd/common.py @@ -578,3 +578,10 @@ def wrapper(self, *args, **kwargs): else: return safe_cast_tensor(returned_tensor, self.precision, GLOBAL_TF_FLOAT_PRECISION) return wrapper + + +def clear_session(): + """Reset all state generated by DeePMD-kit.""" + tf.reset_default_graph() + # TODO: remove this line when data_requirement is not a global variable + data_requirement.clear() diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index 2820f2ecfe..845f3940c7 100644 --- a/deepmd/entrypoints/main.py +++ b/deepmd/entrypoints/main.py @@ -7,6 +7,7 @@ from typing import Dict, List, Optional from deepmd import __version__ +from deepmd.common import clear_session from deepmd.entrypoints import ( compress, config, @@ -548,15 +549,24 @@ def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace: return parsed_args -def main(): +def main(args: Optional[List[str]] = None): """DeePMD-Kit entry point. + Parameters + ---------- + args: List[str], optional + list of command line arguments, used to avoid calling from the subprocess, + as it is quite slow to import tensorflow + Raises ------ RuntimeError if no command was input """ - args = parse_args() + if args is not None: + clear_session() + + args = parse_args(args=args) # do not set log handles for None, it is useless # log handles for train will be set separatelly @@ -592,3 +602,6 @@ def main(): pass else: raise RuntimeError(f"unknown command {args.command}") + + if args is not None: + clear_session() diff --git a/source/tests/common.py b/source/tests/common.py index 38c5bc1aee..4b86ef477d 100644 --- a/source/tests/common.py +++ b/source/tests/common.py @@ -6,6 +6,7 @@ from deepmd.env import GLOBAL_NP_FLOAT_PRECISION from deepmd.common import j_loader as dp_j_loader from deepmd.utils import random as dp_random +from deepmd.entrypoints.main import main if GLOBAL_NP_FLOAT_PRECISION == np.float32 : global_default_fv_hh = 1e-2 @@ -407,3 +408,28 @@ def strerch_box(old_coord, old_box, new_box): nbox = new_box.reshape(3,3) ncoord = ocoord @ np.linalg.inv(obox) @ nbox return ncoord.reshape(old_coord.shape) + + +def run_dp(cmd: str) -> int: + """Run DP directly from the entry point instead of the subprocess. + + It is quite slow to start DeePMD-kit with subprocess. + + Parameters + ---------- + cmd : str + The command to run. + + Returns + ------- + int + Always returns 0. + """ + cmds = cmd.split() + if cmds[0] == 'dp': + cmds = cmds[1:] + else: + raise RuntimeError('The command is not dp') + + main(cmds) + return 0 diff --git a/source/tests/test_adjust_sel.py b/source/tests/test_adjust_sel.py index d16cffe6ce..3f3494f040 100644 --- a/source/tests/test_adjust_sel.py +++ b/source/tests/test_adjust_sel.py @@ -6,7 +6,7 @@ from deepmd.infer import DeepPot from deepmd.env import MODEL_VERSION # from deepmd.entrypoints.compress import compress -from common import j_loader, tests_path +from common import j_loader, tests_path, run_dp from deepmd.env import GLOBAL_NP_FLOAT_PRECISION if GLOBAL_NP_FLOAT_PRECISION == np.float32 : @@ -44,26 +44,26 @@ def _init_models(): with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) - ret = _subprocess_run("dp train " + INPUT + " --skip-neighbor-stat") + ret = run_dp("dp train " + INPUT + " --skip-neighbor-stat") np.testing.assert_equal(ret, 0, 'DP train failed!') - ret = _subprocess_run("dp freeze -o " + frozen_model) + ret = run_dp("dp freeze -o " + frozen_model) np.testing.assert_equal(ret, 0, 'DP freeze failed!') jdata["training"]["numb_steps"] = 0 jdata["model"]["descriptor"]["sel"] = [2, 4] # equal to data with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) - ret = _subprocess_run("dp train " + INPUT + " -f " + frozen_model + " --skip-neighbor-stat") + ret = run_dp("dp train " + INPUT + " -f " + frozen_model + " --skip-neighbor-stat") np.testing.assert_equal(ret, 0, 'DP model adjust sel failed!') - ret = _subprocess_run("dp freeze -o " + decreased_model) + ret = run_dp("dp freeze -o " + decreased_model) np.testing.assert_equal(ret, 0, 'DP freeze failed!') jdata["model"]["descriptor"]["sel"] = [300, 300] # equal to data with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) - ret = _subprocess_run("dp train " + INPUT + " -f " + frozen_model + " --skip-neighbor-stat") + ret = run_dp("dp train " + INPUT + " -f " + frozen_model + " --skip-neighbor-stat") np.testing.assert_equal(ret, 0, 'DP model adjust sel failed!') - ret = _subprocess_run("dp freeze -o " + increased_model) + ret = run_dp("dp freeze -o " + increased_model) np.testing.assert_equal(ret, 0, 'DP freeze failed!') return INPUT, frozen_model, decreased_model, increased_model diff --git a/source/tests/test_finetune_se_atten.py b/source/tests/test_finetune_se_atten.py index 54eb663a9d..a3124c726c 100644 --- a/source/tests/test_finetune_se_atten.py +++ b/source/tests/test_finetune_se_atten.py @@ -2,7 +2,7 @@ import numpy as np import unittest import subprocess as sp -from common import j_loader, tests_path +from common import j_loader, tests_path, run_dp from deepmd.utils.graph import get_tensor_by_name from deepmd.utils.argcheck import normalize from deepmd.utils.compat import update_deepmd_input @@ -62,17 +62,17 @@ def _init_models(): with open(INPUT_FINETUNE_MIX, "w") as fp: json.dump(jdata_finetune, fp, indent=4) - ret = _subprocess_run("dp train " + INPUT_PRE) + ret = run_dp("dp train " + INPUT_PRE) np.testing.assert_equal(ret, 0, 'DP train failed!') - ret = _subprocess_run("dp freeze -o " + pretrained_model) + ret = run_dp("dp freeze -o " + pretrained_model) np.testing.assert_equal(ret, 0, 'DP freeze failed!') - ret = _subprocess_run("dp train " + INPUT_FINETUNE + " -t " + pretrained_model) + ret = run_dp("dp train " + INPUT_FINETUNE + " -t " + pretrained_model) np.testing.assert_equal(ret, 0, 'DP finetune failed!') - ret = _subprocess_run("dp freeze -o " + finetuned_model) + ret = run_dp("dp freeze -o " + finetuned_model) np.testing.assert_equal(ret, 0, 'DP freeze failed!') - ret = _subprocess_run("dp train " + INPUT_FINETUNE_MIX + " -t " + pretrained_model) + ret = run_dp("dp train " + INPUT_FINETUNE_MIX + " -t " + pretrained_model) np.testing.assert_equal(ret, 0, 'DP finetune failed!') - ret = _subprocess_run("dp freeze -o " + finetuned_model_mixed_type) + ret = run_dp("dp freeze -o " + finetuned_model_mixed_type) np.testing.assert_equal(ret, 0, 'DP freeze failed!') jdata_pre = update_deepmd_input(jdata_pre, warning=True, dump="input_v2_compat.json") diff --git a/source/tests/test_init_frz_model_se_a.py b/source/tests/test_init_frz_model_se_a.py index cad93dbeca..1de2996524 100644 --- a/source/tests/test_init_frz_model_se_a.py +++ b/source/tests/test_init_frz_model_se_a.py @@ -2,7 +2,7 @@ import numpy as np import unittest import subprocess as sp -from common import j_loader, tests_path +from common import j_loader, tests_path, run_dp from deepmd.train.trainer import DPTrainer from deepmd.train.run_options import RunOptions from deepmd.utils.argcheck import normalize @@ -48,9 +48,9 @@ def _init_models(): jdata["training"]["save_ckpt"] = ckpt with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) - ret = _subprocess_run("dp train " + INPUT) + ret = run_dp("dp train " + INPUT) np.testing.assert_equal(ret, 0, 'DP train failed!') - ret = _subprocess_run("dp freeze -c " + str(tests_path) + " -o " + frozen_model) + ret = run_dp("dp freeze -c " + str(tests_path) + " -o " + frozen_model) np.testing.assert_equal(ret, 0, 'DP freeze failed!') jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") diff --git a/source/tests/test_init_frz_model_se_a_type.py b/source/tests/test_init_frz_model_se_a_type.py index 2e7a693e3f..7c66922d6b 100644 --- a/source/tests/test_init_frz_model_se_a_type.py +++ b/source/tests/test_init_frz_model_se_a_type.py @@ -2,7 +2,7 @@ import numpy as np import unittest import subprocess as sp -from common import j_loader, tests_path +from common import j_loader, tests_path, run_dp from deepmd.train.trainer import DPTrainer from deepmd.train.run_options import RunOptions from deepmd.utils.argcheck import normalize @@ -52,9 +52,9 @@ def _init_models(): jdata['model']["type_embedding"] = type_embed with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) - ret = _subprocess_run("dp train " + INPUT) + ret = run_dp("dp train " + INPUT) np.testing.assert_equal(ret, 0, 'DP train failed!') - ret = _subprocess_run("dp freeze -c " + str(tests_path) + " -o " + frozen_model) + ret = run_dp("dp freeze -c " + str(tests_path) + " -o " + frozen_model) np.testing.assert_equal(ret, 0, 'DP freeze failed!') jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") diff --git a/source/tests/test_init_frz_model_se_atten.py b/source/tests/test_init_frz_model_se_atten.py index 0decb022e9..d6b0b33260 100644 --- a/source/tests/test_init_frz_model_se_atten.py +++ b/source/tests/test_init_frz_model_se_atten.py @@ -2,7 +2,7 @@ import numpy as np import unittest import subprocess as sp -from common import j_loader, tests_path +from common import j_loader, tests_path, run_dp from deepmd.train.trainer import DPTrainer from deepmd.train.run_options import RunOptions from deepmd.utils.argcheck import normalize @@ -51,9 +51,9 @@ def _init_models(): jdata['model']["descriptor"]['sel'] = 120 with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) - ret = _subprocess_run("dp train " + INPUT) + ret = run_dp("dp train " + INPUT) np.testing.assert_equal(ret, 0, 'DP train failed!') - ret = _subprocess_run("dp freeze -c " + str(tests_path) + " -o " + frozen_model) + ret = run_dp("dp freeze -c " + str(tests_path) + " -o " + frozen_model) np.testing.assert_equal(ret, 0, 'DP freeze failed!') jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") diff --git a/source/tests/test_init_frz_model_se_r.py b/source/tests/test_init_frz_model_se_r.py index f74a0ea222..20d27b99da 100644 --- a/source/tests/test_init_frz_model_se_r.py +++ b/source/tests/test_init_frz_model_se_r.py @@ -2,7 +2,7 @@ import numpy as np import unittest import subprocess as sp -from common import j_loader, tests_path +from common import j_loader, tests_path, run_dp from deepmd.train.trainer import DPTrainer from deepmd.train.run_options import RunOptions from deepmd.utils.argcheck import normalize @@ -56,9 +56,9 @@ def _init_models(): jdata["training"]["save_ckpt"] = ckpt with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) - ret = _subprocess_run("dp train " + INPUT) + ret = run_dp("dp train " + INPUT) np.testing.assert_equal(ret, 0, 'DP train failed!') - ret = _subprocess_run("dp freeze -c " + str(tests_path) + " -o " + frozen_model) + ret = run_dp("dp freeze -c " + str(tests_path) + " -o " + frozen_model) np.testing.assert_equal(ret, 0, 'DP freeze failed!') jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") diff --git a/source/tests/test_mixed_prec_training.py b/source/tests/test_mixed_prec_training.py index 28a1d485f7..e51c64aaaf 100644 --- a/source/tests/test_mixed_prec_training.py +++ b/source/tests/test_mixed_prec_training.py @@ -6,7 +6,7 @@ from deepmd.infer import DeepPot # from deepmd.entrypoints.compress import compress -from common import j_loader, tests_path +from common import j_loader, tests_path, run_dp from deepmd.env import TF_VERSION @@ -43,7 +43,7 @@ def test_training(self): _TF_VERSION = Version(TF_VERSION) # check the TF_VERSION, when TF < 1.12, mixed precision is not allowed if _TF_VERSION >= Version('1.14.0'): - ret = _subprocess_run("dp train " + self.INPUT) + ret = run_dp("dp train " + self.INPUT) np.testing.assert_equal(ret, 0, 'DP train failed!') def tearDown(self): diff --git a/source/tests/test_model_compression_se_a.py b/source/tests/test_model_compression_se_a.py index 1c13a9d610..4db6e8a12f 100644 --- a/source/tests/test_model_compression_se_a.py +++ b/source/tests/test_model_compression_se_a.py @@ -6,7 +6,7 @@ from deepmd.infer import DeepPot from deepmd.env import MODEL_VERSION # from deepmd.entrypoints.compress import compress -from common import j_loader, tests_path +from common import j_loader, tests_path, run_dp from deepmd.env import GLOBAL_NP_FLOAT_PRECISION if GLOBAL_NP_FLOAT_PRECISION == np.float32 : @@ -41,11 +41,11 @@ def _init_models(): with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) - ret = _subprocess_run("dp train " + INPUT) + ret = run_dp("dp train " + INPUT) np.testing.assert_equal(ret, 0, 'DP train failed!') - ret = _subprocess_run("dp freeze -o " + frozen_model) + ret = run_dp("dp freeze -o " + frozen_model) np.testing.assert_equal(ret, 0, 'DP freeze failed!') - ret = _subprocess_run("dp compress " + " -i " + frozen_model + " -o " + compressed_model) + ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model) np.testing.assert_equal(ret, 0, 'DP model compression failed!') return INPUT, frozen_model, compressed_model diff --git a/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py b/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py index b223742d2b..ab575283ab 100644 --- a/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py +++ b/source/tests/test_model_compression_se_a_type_one_side_exclude_types.py @@ -6,7 +6,7 @@ from deepmd.infer import DeepPot from deepmd.env import MODEL_VERSION # from deepmd.entrypoints.compress import compress -from common import j_loader, tests_path +from common import j_loader, tests_path, run_dp from deepmd.env import GLOBAL_NP_FLOAT_PRECISION if GLOBAL_NP_FLOAT_PRECISION == np.float32 : @@ -43,11 +43,11 @@ def _init_models(): with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) - ret = _subprocess_run("dp train " + INPUT) + ret = run_dp("dp train " + INPUT) np.testing.assert_equal(ret, 0, 'DP train failed!') - ret = _subprocess_run("dp freeze -o " + frozen_model) + ret = run_dp("dp freeze -o " + frozen_model) np.testing.assert_equal(ret, 0, 'DP freeze failed!') - ret = _subprocess_run("dp compress " + " -i " + frozen_model + " -o " + compressed_model) + ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model) np.testing.assert_equal(ret, 0, 'DP model compression failed!') return INPUT, frozen_model, compressed_model diff --git a/source/tests/test_model_compression_se_r.py b/source/tests/test_model_compression_se_r.py index 9b9218d175..9433c93017 100644 --- a/source/tests/test_model_compression_se_r.py +++ b/source/tests/test_model_compression_se_r.py @@ -6,7 +6,7 @@ from deepmd.infer import DeepPot from deepmd.env import MODEL_VERSION # from deepmd.entrypoints.compress import compress -from common import j_loader, tests_path +from common import j_loader, tests_path, run_dp from deepmd.env import GLOBAL_NP_FLOAT_PRECISION if GLOBAL_NP_FLOAT_PRECISION == np.float32 : @@ -49,11 +49,11 @@ def _init_models(): with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) - ret = _subprocess_run("dp train " + INPUT) + ret = run_dp("dp train " + INPUT) np.testing.assert_equal(ret, 0, 'DP train failed!') - ret = _subprocess_run("dp freeze -o " + frozen_model) + ret = run_dp("dp freeze -o " + frozen_model) np.testing.assert_equal(ret, 0, 'DP freeze failed!') - ret = _subprocess_run("dp compress " + " -i " + frozen_model + " -o " + compressed_model) + ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model) np.testing.assert_equal(ret, 0, 'DP model compression failed!') return INPUT, frozen_model, compressed_model diff --git a/source/tests/test_model_compression_se_t.py b/source/tests/test_model_compression_se_t.py index 4c77a392c7..71e58fd212 100644 --- a/source/tests/test_model_compression_se_t.py +++ b/source/tests/test_model_compression_se_t.py @@ -6,7 +6,7 @@ from deepmd.infer import DeepPot from deepmd.env import MODEL_VERSION # from deepmd.entrypoints.compress import compress -from common import j_loader, tests_path +from common import j_loader, tests_path, run_dp from deepmd.env import GLOBAL_NP_FLOAT_PRECISION if GLOBAL_NP_FLOAT_PRECISION == np.float32 : @@ -49,11 +49,11 @@ def _init_models(): with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) - ret = _subprocess_run("dp train " + INPUT) + ret = run_dp("dp train " + INPUT) np.testing.assert_equal(ret, 0, 'DP train failed!') - ret = _subprocess_run("dp freeze -o " + frozen_model) + ret = run_dp("dp freeze -o " + frozen_model) np.testing.assert_equal(ret, 0, 'DP freeze failed!') - ret = _subprocess_run("dp compress " + " -i " + frozen_model + " -o " + compressed_model) + ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model) np.testing.assert_equal(ret, 0, 'DP model compression failed!') return INPUT, frozen_model, compressed_model diff --git a/source/tests/test_transfer.py b/source/tests/test_transfer.py index 7b5884b2b4..7322be71fd 100644 --- a/source/tests/test_transfer.py +++ b/source/tests/test_transfer.py @@ -5,7 +5,7 @@ from deepmd.env import tf from deepmd.infer import DeepPot -from common import j_loader, tests_path +from common import j_loader, tests_path, run_dp from infer.convert2pb import convert_pbtxt_to_pb from deepmd.entrypoints.transfer import load_graph, transform_graph @@ -37,7 +37,7 @@ def setUpClass(self): self.new_model = str(tests_path / "dp-new.pb") convert_pbtxt_to_pb(str(tests_path / os.path.join("infer","deeppot.pbtxt")), self.old_model) convert_pbtxt_to_pb(str(tests_path / os.path.join("infer","deeppot-1.pbtxt")), self.raw_model) - ret = _subprocess_run("dp transfer -O " + self.old_model + " -r " + self.raw_model + " -o " + self.new_model) + ret = run_dp("dp transfer -O " + self.old_model + " -r " + self.raw_model + " -o " + self.new_model) np.testing.assert_equal(ret, 0, 'DP transfer failed!') self.dp = DeepPot(self.new_model)