dpa1 compatible with torch.export

Author: Han Wang

deepmd/pt/model/descriptor/se_atten.py

@@ -686,7 +686,6 @@ def forward(
                     self.filter_neuron[-1],
                     self.is_sorted,
                 )[0]
-                # to make torchscript happy
                 gg = torch.empty(
                     nframes,
                     nloc,

source/tests/pt/model/test_dpa1.py

@@ -108,7 +108,7 @@ def test_consistency(
                 err_msg=err_msg,
             )
 
-    def test_jit(
+    def test_export(
         self,
     ) -> None:
         rng = np.random.default_rng(GLOBAL_SEED)
@@ -132,8 +132,6 @@ def test_jit(
             [False, True],  # use_econf_tebd
         ):
             dtype = PRECISION_DICT[prec]
-            rtol, atol = get_tols(prec)
-            err_msg = f"idt={idt} prec={prec}"
             # dpa1 new impl
             dd0 = DescrptDPA1(
                 self.rcut,
@@ -151,6 +149,9 @@ def test_jit(
             )
             dd0.se_atten.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE)
             dd0.se_atten.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE)
-            # dd1 = DescrptDPA1.deserialize(dd0.serialize())
-            model = torch.jit.script(dd0)
-            # model = torch.jit.script(dd1)
+
+            coord_ext = torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE)
+            atype_ext = torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE)
+            nlist = torch.tensor(self.nlist, dtype=int, device=env.DEVICE)
+
+            _ = torch.export.export(dd0, (coord_ext, atype_ext, nlist))

source/tests/pt/model/test_export.py

@@ -30,6 +30,7 @@
 )
 
 from .test_permutation import (
+    model_dpa1,
     model_se_e2_a,
 )
 
@@ -167,6 +168,22 @@ def tearDown(self) -> None:
                 os.remove(f)
 
 
+class TestEnergyModelDPA1IntegrationExport(unittest.TestCase, ExportIntegrationTest):
+    def setUp(self) -> None:
+        input_json = str(Path(__file__).parent / "water/se_atten.json")
+        with open(input_json) as f:
+            self.config = json.load(f)
+        data_file = [str(Path(__file__).parent / "water/data/data_0")]
+        self.config["training"]["training_data"]["systems"] = data_file
+        self.config["training"]["validation_data"]["systems"] = data_file
+        self.config["model"] = deepcopy(model_dpa1)
+        self.config["training"]["numb_steps"] = 2
+        self.config["training"]["save_freq"] = 2
+
+    def tearDown(self) -> None:
+        ExportIntegrationTest.tearDown(self)
+
+
 class TestEnergyModelSeAIntegrationExport(unittest.TestCase, ExportIntegrationTest):
     def setUp(self) -> None:
         input_json = str(Path(__file__).parent / "water/se_atten.json")

source/tests/pt/model/test_jit.py

@@ -21,7 +21,6 @@
 
 from .test_permutation import (
     model_dos,
-    model_dpa1,
     model_dpa2,
     model_hybrid,
 )
@@ -62,22 +61,6 @@ def tearDown(self) -> None:
         JITTest.tearDown(self)
 
 
-class TestEnergyModelDPA1(unittest.TestCase, JITTest):
-    def setUp(self) -> None:
-        input_json = str(Path(__file__).parent / "water/se_atten.json")
-        with open(input_json) as f:
-            self.config = json.load(f)
-        data_file = [str(Path(__file__).parent / "water/data/data_0")]
-        self.config["training"]["training_data"]["systems"] = data_file
-        self.config["training"]["validation_data"]["systems"] = data_file
-        self.config["model"] = deepcopy(model_dpa1)
-        self.config["training"]["numb_steps"] = 10
-        self.config["training"]["save_freq"] = 10
-
-    def tearDown(self) -> None:
-        JITTest.tearDown(self)
-
-
 class TestEnergyModelDPA2(unittest.TestCase, JITTest):
     def setUp(self) -> None:
         input_json = str(Path(__file__).parent / "water/se_atten.json")

source/tests/pt/model/test_se_atten_v2.py

@@ -104,7 +104,7 @@ def test_consistency(
                 err_msg=err_msg,
             )
 
-    def test_jit(
+    def test_export(
         self,
     ) -> None:
         rng = np.random.default_rng()
@@ -140,5 +140,10 @@ def test_jit(
                 seed=GLOBAL_SEED,
             )
             dd0.se_atten.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE)
-            dd0.se_atten.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE)
-            _ = torch.jit.script(dd0)
+            dd0.se_atten.stddev = torch.tensor(dstd, dtype=dtype, device=env.DEVICE)
+
+            coord_ext = torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE)
+            atype_ext = torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE)
+            nlist = torch.tensor(self.nlist, dtype=int, device=env.DEVICE)
+
+            _ = torch.export.export(dd0, (coord_ext, atype_ext, nlist))