From 0e2e79b21ee084529f752bc18f4d014817ac6986 Mon Sep 17 00:00:00 2001 From: denghuilu Date: Thu, 12 Aug 2021 18:34:45 +0800 Subject: [PATCH] speedup model compression UTs --- source/tests/test_model_compression.py | 167 +++++++------------------ 1 file changed, 44 insertions(+), 123 deletions(-) diff --git a/source/tests/test_model_compression.py b/source/tests/test_model_compression.py index 78dfbb26ac..8465c5a05e 100644 --- a/source/tests/test_model_compression.py +++ b/source/tests/test_model_compression.py @@ -17,27 +17,32 @@ def _file_delete(file) : if os.path.exists(file): os.remove(file) -class TestDeepPotAPBC(unittest.TestCase) : - def setUp(self): - self.data_file = str(tests_path / os.path.join("model_compression", "data")) - self.frozen_model = str(tests_path / "dp-original.pb") - self.compressed_model = str(tests_path / "dp-compressed.pb") - self.INPUT = str(tests_path / "input.json") - jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) - jdata["training"]["training_data"]["systems"] = self.data_file - jdata["training"]["validation_data"]["systems"] = self.data_file - with open(self.INPUT, "w") as fp: - json.dump(jdata, fp, indent=4) +def _init_models(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original.pb") + compressed_model = str(tests_path / "dp-compressed.pb") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + jdata["training"]["training_data"]["systems"] = data_file + jdata["training"]["validation_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = os.system("dp train " + INPUT) + assert(ret == 0), "DP train error!" + ret = os.system("dp freeze -o " + frozen_model) + assert(ret == 0), "DP freeze error!" + ret = os.system("dp compress " + " -i " + frozen_model + " -o " + compressed_model) + assert(ret == 0), "DP model compression error!" + return INPUT, frozen_model, compressed_model + +INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() - ret = os.system("dp train " + self.INPUT) - assert(ret == 0), "DP train error!" - ret = os.system("dp freeze -o " + self.frozen_model) - assert(ret == 0), "DP freeze error!" - ret = os.system("dp compress " + " -i " + self.frozen_model + " -o " + self.compressed_model) - assert(ret == 0), "DP model compression error!" - - self.dp_original = DeepPot(self.frozen_model) - self.dp_compressed = DeepPot(self.compressed_model) +class TestDeepPotAPBC(unittest.TestCase) : + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(COMPRESSED_MODEL) self.coords = np.array([12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 00.25, 3.32, 1.68, @@ -47,18 +52,6 @@ def setUp(self): self.atype = [0, 1, 1, 0, 1, 1] self.box = np.array([13., 0., 0., 0., 13., 0., 0., 0., 13.]) - def tearDown(self): - _file_delete(self.INPUT) - _file_delete(self.frozen_model) - _file_delete(self.compressed_model) - _file_delete("out.json") - _file_delete("compress.json") - _file_delete("checkpoint") - _file_delete("lcurve.out") - _file_delete("model.ckpt.meta") - _file_delete("model.ckpt.index") - _file_delete("model.ckpt.data-00000-of-00001") - def test_attrs(self): self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places = default_places) @@ -153,26 +146,10 @@ def test_2frame_atm(self): class TestDeepPotANoPBC(unittest.TestCase) : - def setUp(self): - self.data_file = str(tests_path / os.path.join("model_compression", "data")) - self.frozen_model = str(tests_path / "dp-original.pb") - self.compressed_model = str(tests_path / "dp-compressed.pb") - self.INPUT = str(tests_path / "input.json") - jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) - jdata["training"]["training_data"]["systems"] = self.data_file - jdata["training"]["validation_data"]["systems"] = self.data_file - with open(self.INPUT, "w") as fp: - json.dump(jdata, fp, indent=4) - - ret = os.system("dp train " + self.INPUT) - assert(ret == 0), "DP train error!" - ret = os.system("dp freeze -o " + self.frozen_model) - assert(ret == 0), "DP freeze error!" - ret = os.system("dp compress " + " -i " + self.frozen_model + " -o " + self.compressed_model) - assert(ret == 0), "DP model compression error!" - - self.dp_original = DeepPot(self.frozen_model) - self.dp_compressed = DeepPot(self.compressed_model) + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(COMPRESSED_MODEL) self.coords = np.array([12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 00.25, 3.32, 1.68, @@ -182,18 +159,6 @@ def setUp(self): self.atype = [0, 1, 1, 0, 1, 1] self.box = None - def tearDown(self): - _file_delete(self.INPUT) - _file_delete(self.frozen_model) - _file_delete(self.compressed_model) - _file_delete("out.json") - _file_delete("compress.json") - _file_delete("checkpoint") - _file_delete("lcurve.out") - _file_delete("model.ckpt.meta") - _file_delete("model.ckpt.index") - _file_delete("model.ckpt.data-00000-of-00001") - def test_1frame(self): ee0, ff0, vv0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = False) ee1, ff1, vv1 = self.dp_compressed.eval(self.coords, self.box, self.atype, atomic = False) @@ -274,26 +239,10 @@ def test_2frame_atm(self): class TestDeepPotALargeBoxNoPBC(unittest.TestCase) : - def setUp(self): - self.data_file = str(tests_path / os.path.join("model_compression", "data")) - self.frozen_model = str(tests_path / "dp-original.pb") - self.compressed_model = str(tests_path / "dp-compressed.pb") - self.INPUT = str(tests_path / "input.json") - jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) - jdata["training"]["training_data"]["systems"] = self.data_file - jdata["training"]["validation_data"]["systems"] = self.data_file - with open(self.INPUT, "w") as fp: - json.dump(jdata, fp, indent=4) - - ret = os.system("dp train " + self.INPUT) - assert(ret == 0), "DP train error!" - ret = os.system("dp freeze -o " + self.frozen_model) - assert(ret == 0), "DP freeze error!" - ret = os.system("dp compress " + " -i " + self.frozen_model + " -o " + self.compressed_model) - assert(ret == 0), "DP model compression error!" - - self.dp_original = DeepPot(self.frozen_model) - self.dp_compressed = DeepPot(self.compressed_model) + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(COMPRESSED_MODEL) self.coords = np.array([12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 00.25, 3.32, 1.68, @@ -303,18 +252,6 @@ def setUp(self): self.atype = [0, 1, 1, 0, 1, 1] self.box = np.array([19., 0., 0., 0., 13., 0., 0., 0., 13.]) - def tearDown(self): - _file_delete(self.INPUT) - _file_delete(self.frozen_model) - _file_delete(self.compressed_model) - _file_delete("out.json") - _file_delete("compress.json") - _file_delete("checkpoint") - _file_delete("lcurve.out") - _file_delete("model.ckpt.meta") - _file_delete("model.ckpt.index") - _file_delete("model.ckpt.data-00000-of-00001") - def test_1frame(self): ee0, ff0, vv0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = False) ee1, ff1, vv1 = self.dp_compressed.eval(self.coords, self.box, self.atype, atomic = False) @@ -369,11 +306,11 @@ def test_ase(self): water0 = Atoms('OHHOHH', positions=self.coords.reshape((-1,3)), cell=self.box.reshape((3,3)), - calculator=DP(self.frozen_model)) + calculator=DP(FROZEN_MODEL)) water1 = Atoms('OHHOHH', positions=self.coords.reshape((-1,3)), cell=self.box.reshape((3,3)), - calculator=DP(self.compressed_model)) + calculator=DP(COMPRESSED_MODEL)) ee0 = water0.get_potential_energy() ff0 = water0.get_forces() ee1 = water1.get_potential_energy() @@ -385,27 +322,10 @@ def test_ase(self): self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) class TestDeepPotAPBCExcludeTypes(unittest.TestCase) : - def setUp(self): - self.data_file = str(tests_path / os.path.join("model_compression", "data")) - self.frozen_model = str(tests_path / "dp-original.pb") - self.compressed_model = str(tests_path / "dp-compressed.pb") - self.INPUT = str(tests_path / "input.json") - jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) - jdata["training"]["training_data"]["systems"] = self.data_file - jdata["training"]["validation_data"]["systems"] = self.data_file - jdata["model"]["descriptor"]["exclude_types"] = [[0, 1]] - with open(self.INPUT, "w") as fp: - json.dump(jdata, fp, indent=4) - - ret = os.system("dp train " + self.INPUT) - assert(ret == 0), "DP train error!" - ret = os.system("dp freeze -o " + self.frozen_model) - assert(ret == 0), "DP freeze error!" - ret = os.system("dp compress " + " -i " + self.frozen_model + " -o " + self.compressed_model) - assert(ret == 0), "DP model compression error!" - - self.dp_original = DeepPot(self.frozen_model) - self.dp_compressed = DeepPot(self.compressed_model) + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(COMPRESSED_MODEL) self.coords = np.array([12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 00.25, 3.32, 1.68, @@ -415,10 +335,11 @@ def setUp(self): self.atype = [0, 1, 1, 0, 1, 1] self.box = np.array([13., 0., 0., 0., 13., 0., 0., 0., 13.]) - def tearDown(self): - _file_delete(self.INPUT) - _file_delete(self.frozen_model) - _file_delete(self.compressed_model) + @classmethod + def tearDownClass(self): + _file_delete(INPUT) + _file_delete(FROZEN_MODEL) + _file_delete(COMPRESSED_MODEL) _file_delete("out.json") _file_delete("compress.json") _file_delete("checkpoint")