From d71ce3e2f4a1fb0cacbfbb1c39b6b877d3928727 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Sat, 13 Jul 2024 23:18:28 +0800
Subject: [PATCH] feat(tf): improve the activation setting in tebd (#3971)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Added logging to notify users when multiple neurons are used with a
"linear" activation function, equivalent to a single linear layer.

- **Configuration Updates**
- Introduced `"activation_function": "linear"` in various neural network
configuration files (`input.json`).
- Added `"activation_function": "tanh"` in the `input.json` file for
water model configuration.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 deepmd/tf/utils/type_embed.py              | 8 ++++++++
 examples/dprc/generalized_force/input.json | 1 +
 examples/dprc/normal/input.json            | 1 +
 examples/dprc/pairwise/input.json          | 1 +
 examples/water/se_e2_a_tebd/input.json     | 1 +
 5 files changed, 12 insertions(+)

diff --git a/deepmd/tf/utils/type_embed.py b/deepmd/tf/utils/type_embed.py
index 3f40a0cef1..c8d40b5d0b 100644
--- a/deepmd/tf/utils/type_embed.py
+++ b/deepmd/tf/utils/type_embed.py
@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import logging
 import re
 from typing import (
     List,
@@ -33,6 +34,8 @@
     check_version_compatibility,
 )
 
+log = logging.getLogger(__name__)
+
 
 def embed_atom_type(
     ntypes: int,
@@ -131,6 +134,11 @@ def __init__(
         self.filter_precision = get_precision(precision)
         self.filter_activation_fn_name = str(activation_function)
         self.filter_activation_fn = get_activation_func(activation_function)
+        if len(neuron) > 1 and self.filter_activation_fn_name == "linear":
+            log.warning(
+                "Multi-layer type embedding is being used with a 'linear' activation function. "
+                "This is equivalent to a single linear layer. Please ensure this is intentional."
+            )
         self.trainable = trainable
         self.uniform_seed = uniform_seed
         self.type_embedding_net_variables = None
diff --git a/examples/dprc/generalized_force/input.json b/examples/dprc/generalized_force/input.json
index 1bb0f687eb..4537f2b5c9 100644
--- a/examples/dprc/generalized_force/input.json
+++ b/examples/dprc/generalized_force/input.json
@@ -13,6 +13,7 @@
       "neuron": [
         8
       ],
+      "activation_function": "linear",
       "precision": "float32"
     },
     "descriptor": {
diff --git a/examples/dprc/normal/input.json b/examples/dprc/normal/input.json
index f85607f855..8c1058abc8 100644
--- a/examples/dprc/normal/input.json
+++ b/examples/dprc/normal/input.json
@@ -13,6 +13,7 @@
       "neuron": [
         8
       ],
+      "activation_function": "linear",
       "precision": "float32"
     },
     "descriptor": {
diff --git a/examples/dprc/pairwise/input.json b/examples/dprc/pairwise/input.json
index c9ebae0f43..c230d4d858 100644
--- a/examples/dprc/pairwise/input.json
+++ b/examples/dprc/pairwise/input.json
@@ -14,6 +14,7 @@
       "neuron": [
         8
       ],
+      "activation_function": "linear",
       "precision": "float32"
     },
     "qm_model": {
diff --git a/examples/water/se_e2_a_tebd/input.json b/examples/water/se_e2_a_tebd/input.json
index 1ec85a968c..7e8b8c6936 100644
--- a/examples/water/se_e2_a_tebd/input.json
+++ b/examples/water/se_e2_a_tebd/input.json
@@ -11,6 +11,7 @@
         4,
         8
       ],
+      "activation_function": "tanh",
       "resnet_dt": false,
       "precision": "float64",
       "seed": 1