diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index f0dc571142..1e2cdcd40c 100644 --- a/deepmd/utils/tabulate.py +++ b/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 diff --git a/doc/freeze/compress.md b/doc/freeze/compress.md index ed55db0c4b..362ae22840 100644 --- a/doc/freeze/compress.md +++ b/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 + + diff --git a/source/op/unaggregated_grad.cc b/source/op/unaggregated_grad.cc index 89c14a84fb..c5f872ab42 100644 --- a/source/op/unaggregated_grad.cc +++ b/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; }