Added all activation functions for model compression.

deepmd/utils/tabulate.py

@@ -68,6 +68,14 @@ def __init__(self,
             self.functype = 1
         elif activation_fn == ACTIVATION_FN_DICT["gelu"]:
             self.functype = 2
+        elif activation_fn == ACTIVATION_FN_DICT["relu"]:
+            self.functype = 3
+        elif activation_fn == ACTIVATION_FN_DICT["relu6"]:
+            self.functype = 4
+        elif activation_fn == ACTIVATION_FN_DICT["softplus"]:
+            self.functype = 5
+        elif activation_fn == ACTIVATION_FN_DICT["sigmoid"]:
+            self.functype = 6
         else:
             raise RuntimeError("Unknown actication function type!")
         self.activation_fn = activation_fn

doc/freeze/compress.md

@@ -83,3 +83,14 @@ The model compression interface requires the version of deepmd-kit used in origi
 **Acceptable descriptor type**
 
 Note only descriptors with `se_e2_a` or `se_e3` type are supported by the model compression feature. Hybrid mixed with above descriptors is also supported.
+
+
+**Available activation functions for descriptor:**
+- tanh
+- gelu
+- relu
+- relu6
+- softplus
+- sigmoid
+
+

source/op/unaggregated_grad.cc

@@ -52,6 +52,36 @@ FPTYPE grad(const FPTYPE xbar, const FPTYPE y, const int functype)  //functype=t
             const FPTYPE var = tanh(SQRT_2_PI * (xbar + GGELU * xbar * xbar * xbar));
             return 0.5 * SQRT_2_PI * xbar * (1 - var * var) * (3 * GGELU * xbar * xbar + 1) + 0.5 * var + 0.5;
         }
+        case 3:
+        {
+            if(xbar<=0)
+            {
+                return 0;
+            }
+            else
+            {
+                return 1;
+            }
+        }
+        case 4:
+        {
+            if(xbar<=0 || xbar>=6)
+            {
+                return 0;
+            }
+            else
+            {
+                return 1;
+            }
+        }
+        case 5:
+        {
+            return 1.0-1.0/(1.0+exp(xbar));
+        }
+        case 6:
+        {
+            return y*(1-y);
+        }
         default:
             return -1;
     }
@@ -71,6 +101,22 @@ FPTYPE grad_grad(const FPTYPE xbar, const FPTYPE y, const int functype)
             const FPTYPE var2 = SQRT_2_PI * (1 - var1 * var1) * (3 * GGELU * xbar * xbar + 1);
             return  3 * GGELU * SQRT_2_PI * xbar * xbar * (1 - var1 * var1) - SQRT_2_PI * xbar * var2 * (3 * GGELU * xbar * xbar + 1) * var1 + var2;
         }
+        case 3:
+        {
+            return 0;
+        }
+        case 4:
+        {
+            return 0;
+        }
+        case 5:
+        {
+            return exp(xbar)/((1+exp(xbar))*(1+exp(xbar)));
+        }
+        case 6:
+        {
+            return y*(1-y)*(1-2*y);
+        }
         default:
             return -1;
     }