feat(tf2): add training workflow

.github/workflows/test_python.yml

@@ -60,7 +60,7 @@ jobs:
           restore-keys: |
             test2-durations-combined-${{ matrix.python }}-${{ github.sha }}
             test2-durations-combined-${{ matrix.python }}
-      - run: pytest --cov=deepmd source/tests  --splits 12 --group ${{ matrix.group }} --store-durations --clean-durations --durations-path=.test_durations --splitting-algorithm least_duration
+      - run: pytest --cov=deepmd source/tests --ignore=source/tests/tf2 --splits 12 --group ${{ matrix.group }} --store-durations --clean-durations --durations-path=.test_durations --splitting-algorithm least_duration
         env:
           NUM_WORKERS: 0
           DP_CI_IMPORT_PADDLE_BEFORE_TF: 1
@@ -81,6 +81,10 @@ jobs:
             return "$status"
           }
 
+          run_pytest_allow_no_tests --cov=deepmd --cov-append \
+            source/tests/tf2 \
+            --splits 12 \
+            --group ${{ matrix.group }}
           run_pytest_allow_no_tests --cov=deepmd --cov-append \
             source/tests/consistent/io/test_io.py \
             source/jax2tf_tests \

deepmd/backend/tf2.py

@@ -33,7 +33,12 @@ class TensorFlow2Backend(Backend):
     """TensorFlow 2 eager backend."""
 
     name = "TensorFlow2"
-    features: ClassVar[Backend.Feature] = Backend.Feature.DEEP_EVAL | Backend.Feature.IO
+    features: ClassVar[Backend.Feature] = (
+        Backend.Feature.ENTRY_POINT
+        | Backend.Feature.DEEP_EVAL
+        | Backend.Feature.NEIGHBOR_STAT
+        | Backend.Feature.IO
+    )
     suffixes: ClassVar[list[str]] = [".savedmodeltf"]
 
     @classmethod
@@ -45,7 +50,11 @@ def is_available(self) -> bool:
 
     @property
     def entry_point_hook(self) -> Callable[["Namespace"], None]:
-        raise NotImplementedError("Training entry point is not implemented for TF2")
+        from deepmd.tf2.entrypoints.main import (
+            main,
+        )
+
+        return main
 
     @property
     def deep_eval(self) -> type["DeepEvalBackend"]:
@@ -57,7 +66,11 @@ def deep_eval(self) -> type["DeepEvalBackend"]:
 
     @property
     def neighbor_stat(self) -> type["NeighborStat"]:
-        raise NotImplementedError("Neighbor statistics are not implemented for TF2")
+        from deepmd.dpmodel.utils.neighbor_stat import (
+            NeighborStat,
+        )
+
+        return NeighborStat
 
     @property
     def serialize_hook(self) -> Callable[[str], dict]:

deepmd/dpmodel/common.py

@@ -90,6 +90,18 @@ def get_xp_precision(
         raise ValueError(f"unsupported precision {precision} for {xp}")
 
 
+def to_numpy_dtype(dtype: Any) -> np.dtype:
+    """Normalize backend dtype objects to a NumPy dtype."""
+    dtype = getattr(dtype, "as_numpy_dtype", dtype)
+    try:
+        return np.dtype(dtype)
+    except TypeError:
+        dtype_name = getattr(dtype, "name", None)
+        if dtype_name is not None:
+            return np.dtype(dtype_name)
+        raise
+
+
 class NativeOP(ABC):
     """The unit operation of a native model."""
 

deepmd/dpmodel/descriptor/dpa1.py

@@ -27,6 +27,7 @@
 from deepmd.dpmodel.common import (
     cast_precision,
     to_numpy_array,
+    to_numpy_dtype,
 )
 from deepmd.dpmodel.utils import (
     EmbeddingNet,
@@ -1409,7 +1410,7 @@ def enable_compression(
     ) -> None:
         """Store tabulated geometric embedding-net data."""
         net = "filter_net"
-        dtype = self.mean.dtype
+        dtype = to_numpy_dtype(self.mean.dtype)
         self.compress_info = [
             np.asarray(
                 [

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -22,6 +22,7 @@
 from deepmd.dpmodel.common import (
     cast_precision,
     to_numpy_array,
+    to_numpy_dtype,
 )
 from deepmd.dpmodel.utils import (
     EmbeddingNet,
@@ -431,7 +432,7 @@ def _store_compress_data(
         """Store tabulated embedding-net data in the descriptor state."""
         compress_data = []
         compress_info = []
-        dtype = self.davg.dtype
+        dtype = to_numpy_dtype(self.davg.dtype)
         ndim = 1 if self.type_one_side else 2
         for embedding_idx in range(self.ntypes**ndim):
             if self.type_one_side:

deepmd/dpmodel/descriptor/se_r.py

@@ -22,6 +22,7 @@
     cast_precision,
     get_xp_precision,
     to_numpy_array,
+    to_numpy_dtype,
 )
 from deepmd.dpmodel.utils import (
     EmbeddingNet,
@@ -410,7 +411,7 @@ def _store_compress_data(
         """Store tabulated embedding-net data in the descriptor state."""
         compress_data = []
         compress_info = []
-        dtype = self.davg.dtype
+        dtype = to_numpy_dtype(self.davg.dtype)
         for embedding_idx in range(self.ntypes):
             net = "filter_-1_net_" + str(embedding_idx)
             if net not in table_data:

deepmd/main.py

@@ -327,7 +327,7 @@ def main_parser() -> argparse.ArgumentParser:
         "--output",
         type=str,
         default="frozen_model",
-        help="Filename (prefix) of the output model file. TensorFlow backend: suffix is .pb; PyTorch backend: suffix is .pth; Paddle backend: suffix is .json and .pdiparams",
+        help="Filename (prefix) of the output model file. TensorFlow backend: suffix is .pb; TensorFlow2 backend: suffix is .savedmodeltf; PyTorch backend: suffix is .pth; Paddle backend: suffix is .json and .pdiparams",
     )
     parser_frz.add_argument(
         "-n",
@@ -605,14 +605,14 @@ def main_parser() -> argparse.ArgumentParser:
         "--input",
         default="frozen_model",
         type=str,
-        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; DPModel backend: suffix is .dp; JAX backend: suffix is .hlo or .jax",
+        help="The original frozen model or checkpoint, which will be compressed by the code. Filename (prefix) of the input model file. TensorFlow backend: suffix is .pb; TensorFlow2 backend: .tf2 checkpoint directory or checkpoint prefix; PyTorch backend: suffix is .pth; DPModel backend: suffix is .dp; JAX backend: suffix is .hlo or .jax",
     )
     parser_compress.add_argument(
         "-o",
         "--output",
         default="frozen_model_compressed",
         type=str,
-        help="The compressed model. Filename (prefix) of the output model file. TensorFlow backend: suffix is .pb; PyTorch backend: suffix is .pth; DPModel backend: suffix is .dp; JAX backend: suffix is .hlo or .jax",
+        help="The compressed model. Filename (prefix) of the output model file. TensorFlow backend: suffix is .pb; TensorFlow2 backend: suffix is .savedmodeltf; PyTorch backend: suffix is .pth; DPModel backend: suffix is .dp; JAX backend: suffix is .hlo or .jax",
     )
     parser_compress.add_argument(
         "-s",
@@ -659,6 +659,13 @@ def main_parser() -> argparse.ArgumentParser:
         default=None,
         help="The training script of the input frozen model",
     )
+    parser_compress.add_argument(
+        "--head",
+        "--model-branch",
+        default=None,
+        type=str,
+        help="Task head (alias: model branch) to compress if in multi-task mode.",
+    )
 
     # * print docs script **************************************************************
     parsers_doc = subparsers.add_parser(

deepmd/tf2/atomic_model/dp_atomic_model.py

@@ -11,6 +11,7 @@
 )
 from deepmd.tf2.env import (
     stop_gradient,
+    tf,
     xp,
 )
 from deepmd.tf2.fitting.base_fitting import (
@@ -41,6 +42,13 @@ class tf2_atomic_model(dpmodel_atomic_model):
         base_fitting_cls = BaseFitting
         """The base fitting class."""
 
+        @tf.autograph.experimental.do_not_convert
+        def make_atom_mask(
+            self,
+            atype: xp.ndarray,
+        ) -> xp.ndarray:
+            return atype >= 0
+
         def forward_common_atomic(
             self,
             extended_coord: xp.ndarray,

deepmd/tf2/entrypoints/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Entry points for the TensorFlow 2 backend."""
+
+from deepmd.tf2.entrypoints.compress import (
+    enable_compression,
+)
+from deepmd.tf2.entrypoints.freeze import (
+    freeze,
+)
+from deepmd.tf2.entrypoints.train import (
+    train,
+)
+
+__all__ = ["enable_compression", "freeze", "train"]

deepmd/tf2/entrypoints/compress.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Compress TensorFlow 2 checkpoints by tabulating embedding networks."""
+
+from __future__ import (
+    annotations,
+)
+
+import logging
+from typing import (
+    Any,
+)
+
+from deepmd.backend.suffix import (
+    format_model_suffix,
+)
+from deepmd.dpmodel.entrypoints.compress_common import (
+    enable_model_compression,
+    resolve_min_nbor_dist,
+)
+from deepmd.dpmodel.utils.update_sel import (
+    UpdateSel,
+)
+from deepmd.tf2.entrypoints.freeze import (
+    select_model_branch,
+)
+from deepmd.tf2.model.base_model import (
+    BaseModel,
+)
+from deepmd.tf2.utils.serialization import (
+    deserialize_to_file,
+    serialize_from_file,
+)
+
+log = logging.getLogger(__name__)
+
+
+def enable_compression(
+    input_file: str,
+    output: str,
+    stride: float = 0.01,
+    extrapolate: int = 5,
+    check_frequency: int = -1,
+    training_script: str | None = None,
+    head: str | None = None,
+    **kwargs: Any,
+) -> None:
+    """Compress a TF2 training checkpoint and export a SavedModel."""
+    del kwargs
+    output = format_model_suffix(
+        output,
+        preferred_backend="tf2",
+        strict_prefer=True,
+    )
+    data = serialize_from_file(input_file)
+    data = select_model_branch(data, head=head)
+    model = BaseModel.deserialize(data["model"])
+    min_nbor_dist = resolve_min_nbor_dist(
+        model,
+        [data],
+        training_script,
+        UpdateSel,
+    )
+    enable_model_compression(
+        model,
+        min_nbor_dist,
+        stride,
+        extrapolate,
+        check_frequency,
+    )
+
+    compressed_data = data.copy()
+    compressed_data["model"] = model.serialize()
+    compressed_data["min_nbor_dist"] = float(min_nbor_dist)
+    deserialize_to_file(output, compressed_data)
+    log.info("Compressed TF2 model saved to %s", output)

deepmd/tf2/entrypoints/freeze.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Freeze utilities for the TensorFlow 2 backend."""
+
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    Any,
+)
+
+from deepmd.backend.suffix import (
+    format_model_suffix,
+)
+from deepmd.dpmodel.train import (
+    DEFAULT_TASK_KEY,
+)
+from deepmd.tf2.utils.serialization import (
+    deserialize_to_file,
+    serialize_from_file,
+)
+from deepmd.utils.model_branch_dict import (
+    get_model_dict,
+)
+
+
+def freeze(
+    *,
+    checkpoint_folder: str,
+    output: str,
+    head: str | None = None,
+    **kwargs: Any,
+) -> None:
+    """Freeze a TF2 training checkpoint into a TensorFlow SavedModel."""
+    del kwargs
+    output = format_model_suffix(
+        output,
+        preferred_backend="tf2",
+        strict_prefer=True,
+    )
+    data = serialize_from_file(checkpoint_folder)
+    data = select_model_branch(data, head=head)
+    deserialize_to_file(output, data)
+
+
+def select_model_branch(
+    data: dict[str, Any], head: str | None = None
+) -> dict[str, Any]:
+    """Select one branch from a single-task or multi-task serialized payload."""
+    model_def_script = data["model_def_script"]
+    if "model_dict" not in model_def_script:
+        if head not in (None, "", DEFAULT_TASK_KEY):
+            raise ValueError(
+                f"Single-task TF2 checkpoints do not have a head named {head!r}."
+            )
+        return data
+
+    if head in (None, ""):
+        raise ValueError(
+            "Multi-task TF2 checkpoints require --head/--model-branch to select "
+            "which model branch to freeze."
+        )
+    model_alias_dict, _ = get_model_dict(model_def_script["model_dict"])
+    if head not in model_alias_dict:
+        raise ValueError(
+            f"No model branch or alias named {head!r}. Available branches are "
+            f"{list(model_def_script['model_dict'])}."
+        )
+    resolved_head = model_alias_dict[head]
+    selected = data.copy()
+    selected["model"] = data["model"]["model_dict"][resolved_head]
+    selected["model_def_script"] = model_def_script["model_dict"][resolved_head]
+    min_nbor_dist = data.get("min_nbor_dist")
+    if isinstance(min_nbor_dist, dict):
+        selected["min_nbor_dist"] = min_nbor_dist.get(resolved_head)
+    return selected