fix bugs of gelu activation function

.travis.yml

@@ -83,7 +83,7 @@ jobs:
         - CXX=g++-7
         - TENSORFLOW_VERSION=2.1
       install:
-        - python -m pip install twine cibuildwheel==1.1.0 scikit-build
+        - python -m pip install twine cibuildwheel==1.1.0 scikit-build setuptools_scm
       script:
         - python -m cibuildwheel --output-dir wheelhouse
         - python setup.py sdist

pyproject.toml

@@ -1,3 +1,3 @@
 [build-system]
-requires = ["setuptools", "wheel", "scikit-build", "cmake", "ninja", "m2r"]
+requires = ["setuptools", "setuptools_scm", "wheel", "scikit-build", "cmake", "ninja", "m2r"]
 

setup.py

@@ -1,9 +1,21 @@
 from skbuild import setup
 from skbuild.exceptions import SKBuildError
 from skbuild.cmaker import get_cmake_version
+from setuptools_scm import get_version
 from packaging.version import LegacyVersion
 from os import path, makedirs
-import imp
+import imp, sys, platform
+
+def get_dp_install_path() :
+    site_packages_path = path.join(path.dirname(path.__file__), 'site-packages')
+    dp_scm_version     = get_version(root="./", relative_to=__file__)
+    python_version     = 'py' + str(sys.version_info.major + sys.version_info.minor * 0.1)
+    os_info            = sys.platform
+    machine_info       = platform.machine()
+    dp_pip_install_path    = site_packages_path + '/deepmd'
+    dp_setup_install_path    = site_packages_path + '/deepmd_kit-' + dp_scm_version + '-' + python_version + '-' + os_info + '-' + machine_info + '.egg/deepmd'
+
+    return dp_pip_install_path, dp_setup_install_path
 
 readme_file = path.join(path.dirname(path.abspath(__file__)), 'README.md')
 try:
@@ -34,6 +46,8 @@
 except OSError:
     pass
 
+dp_pip_install_path, dp_setup_install_path = get_dp_install_path()
+
 setup(
     name="deepmd-kit",
     setup_requires=setup_requires,
@@ -56,6 +70,8 @@
                 '-DBUILD_PY_IF:BOOL=TRUE', 
                 '-DBUILD_CPP_IF:BOOL=FALSE',
                 '-DFLOAT_PREC:STRING=high',
+                '-DDP_PIP_INSTALL_PATH=%s' % dp_pip_install_path,
+                '-DDP_SETUP_INSTALL_PATH=%s' % dp_setup_install_path,
     ],
     cmake_source_dir='source',
     cmake_minimum_required_version='3.0',

source/CMakeLists.txt

@@ -184,6 +184,11 @@ include_directories(${TensorFlow_INCLUDE_DIRS})
 if (BUILD_CPP_IF)
   set (LIB_DEEPMD		"deepmd")
   set (LIB_DEEPMD_OP		"deepmd_op")
+  if (USE_CUDA_TOOLKIT)
+    set (LIB_DEEPMD_OP_CUDA		"deepmd_op_cuda")
+  else ()
+    set (LIB_DEEPMD_OP_CUDA		"deepmd_op")
+  endif()
   if (CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 4.9)
     set (LIB_DEEPMD_NATIVE	"deepmd_native_md")
     set (LIB_DEEPMD_IPI		"deepmd_ipi")

source/lmp/env.sh.in

@@ -8,4 +8,4 @@ TF_RPATH=`echo $TENSORFLOW_LIBRARY_PATH | sed "s/;/ -Wl,-rpath=/g"`
 
 NNP_INC=" -std=c++11 @PREC_DEF@ @TTM_DEF@ @OLD_LMP_PPPM_DEF@ -I$TF_INCLUDE_DIRS -I$DEEPMD_ROOT/include/deepmd "
 NNP_PATH=" -L$TF_LIBRARY_PATH -L$DEEPMD_ROOT/lib"
-NNP_LIB=" -Wl,--no-as-needed -l@LIB_DEEPMD_OP@ -l@LIB_DEEPMD@ -ltensorflow_cc -ltensorflow_framework -Wl,-rpath=$TF_RPATH -Wl,-rpath=$DEEPMD_ROOT/lib"
+NNP_LIB=" -Wl,--no-as-needed -l@LIB_DEEPMD_OP_CUDA@ -l@LIB_DEEPMD_OP@ -l@LIB_DEEPMD@ -ltensorflow_cc -ltensorflow_framework -Wl,-rpath=$TF_RPATH -Wl,-rpath=$DEEPMD_ROOT/lib"

source/op/CMakeLists.txt

@@ -24,6 +24,8 @@ if (BUILD_CPP_IF)
 endif (BUILD_CPP_IF)
 
 if (BUILD_PY_IF)
+  set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
+  set(CMAKE_INSTALL_RPATH DESTINATION ${DP_PIP_INSTALL_PATH} ${DP_SETUP_INSTALL_PATH} ${CMAKE_BINARY_DIR}/op/cuda)
   if (USE_CUDA_TOOLKIT)
     add_library(op_abi SHARED ${OP_PY_CUDA_SRC} ${OP_LIB})
     add_library(op_grads SHARED ${OP_GRADS_SRC})
@@ -33,11 +35,11 @@ if (BUILD_PY_IF)
     set (EXTRA_LIBS ${EXTRA_LIBS} deepmd_op_cuda)
     target_link_libraries (op_abi ${EXTRA_LIBS})
     target_link_libraries (op_grads ${EXTRA_LIBS})
-    message(STATUS ${TensorFlowFramework_LIBRARY})
   else (USE_CUDA_TOOLKIT)
     add_library(op_abi SHARED ${OP_SRC} ${OP_LIB})
     add_library(op_grads SHARED ${OP_GRADS_SRC})
   endif(USE_CUDA_TOOLKIT)
+  message(STATUS ${TensorFlowFramework_LIBRARY})
   target_link_libraries(
     op_abi ${TensorFlowFramework_LIBRARY}
     )

source/op/_gelu.py

@@ -7,9 +7,9 @@
 from deepmd.env import op_module
 
 @ops.RegisterGradient("Gelu")
-def gelu_cc (op, dy) :
+def _gelu_cc (op, dy) :
     return op_module.gelu_grad(dy, op.inputs[0])     
 
 @ops.RegisterGradient("GeluGrad")
-def gelu_grad_cc (op, dy) :
-    return [None, op_module.gelu_grad_grad(dy, op.inputs[0], op.inputs[1])]
+def _gelu_grad_cc (op, dy) :
+    return [op_module.gelu_grad(dy, op.inputs[1]), op_module.gelu_grad_grad(dy, op.inputs[0], op.inputs[1])]

source/op/cuda/CMakeLists.txt

@@ -83,7 +83,7 @@ set (SOURCE_FILES
     descrpt_se_a.cu descrpt_se_r.cu prod_force_se_a.cu prod_force_se_r.cu prod_virial_se_a.cu prod_virial_se_r.cu gelu.cu 
 )
 
-cuda_add_library(deepmd_op_cuda STATIC ${SOURCE_FILES})
+cuda_add_library(deepmd_op_cuda SHARED ${SOURCE_FILES})
 
 if (BUILD_CPP_IF)
     install(TARGETS deepmd_op_cuda DESTINATION lib/)

source/op/gelu.cc

@@ -11,32 +11,20 @@ using GPUDevice = Eigen::GpuDevice;
 REGISTER_OP("Gelu")
     .Attr("T: {float, double}")
     .Input("x: T")
-    .Output("output: T")
-    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
-        c->set_output(0, c->input(0));
-        return Status::OK();
-    });
+    .Output("output: T");
 
 REGISTER_OP("GeluGrad")
     .Attr("T: {float, double}")
     .Input("dy: T")
     .Input("x: T")
-    .Output("output: T")
-    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
-        c->set_output(0, c->input(1));
-        return Status::OK();
-    });
+    .Output("output: T");
 
 REGISTER_OP("GeluGradGrad")
     .Attr("T: {float, double}")
     .Input("dy: T")
     .Input("dy_: T")
     .Input("x: T")
-    .Output("output: T")
-    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
-        c->set_output(0, c->input(2));
-        return Status::OK();
-    });
+    .Output("output: T");
 
 template <typename Device, typename T>
 struct GeluFunctor {

source/op/gelu_gpu.cc

@@ -10,32 +10,20 @@ using GPUDevice = Eigen::GpuDevice;
 REGISTER_OP("Gelu")
     .Attr("T: {float, double}")
     .Input("x: T")
-    .Output("output: T")
-    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
-        c->set_output(0, c->input(0));
-        return Status::OK();
-    });
+    .Output("output: T");
 
 REGISTER_OP("GeluGrad")
     .Attr("T: {float, double}")
     .Input("dy: T")
     .Input("x: T")
-    .Output("output: T")
-    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
-        c->set_output(0, c->input(1));
-        return Status::OK();
-    });
+    .Output("output: T");
 
 REGISTER_OP("GeluGradGrad")
     .Attr("T: {float, double}")
     .Input("dy: T")
     .Input("dy_: T")
     .Input("x: T")
-    .Output("output: T")
-    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
-        c->set_output(0, c->input(2));
-        return Status::OK();
-    });
+    .Output("output: T");
 
 // maybe instead use cudnn activation forward 
 void GeluLauncher(const float * in, float * out, int const size);

source/train/DescrptSeA.py

@@ -353,6 +353,7 @@ def _filter(self,
                                   self.filter_precision,
                                   tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed), 
                                     trainable = trainable)
+                hidden = tf.reshape(activation_fn(tf.matmul(xyz_scatter, w) + b), [-1, outputs_size[ii]])
                 if self.filter_resnet_dt :
                     idt = tf.get_variable('idt_'+str(ii)+'_'+str(type_i), 
                                           [1, outputs_size[ii]], 
@@ -361,16 +362,16 @@ def _filter(self,
                                           trainable = trainable)
                 if outputs_size[ii] == outputs_size[ii-1]:
                     if self.filter_resnet_dt :
-                        xyz_scatter += activation_fn(tf.matmul(xyz_scatter, w) + b) * idt
+                        xyz_scatter += hidden * idt
                     else :
-                        xyz_scatter += activation_fn(tf.matmul(xyz_scatter, w) + b)
+                        xyz_scatter += hidden
                 elif outputs_size[ii] == outputs_size[ii-1] * 2: 
                     if self.filter_resnet_dt :
-                        xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + activation_fn(tf.matmul(xyz_scatter, w) + b) * idt
+                        xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden * idt
                     else :
-                        xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + activation_fn(tf.matmul(xyz_scatter, w) + b)
+                        xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden
                 else:
-                    xyz_scatter = activation_fn(tf.matmul(xyz_scatter, w) + b)
+                    xyz_scatter = hidden
             else:
               w = tf.zeros((outputs_size[0], outputs_size[-1]), dtype=global_tf_float_precision)
               xyz_scatter = tf.matmul(xyz_scatter, w)
@@ -440,6 +441,7 @@ def _filter_type_ext(self,
                                 self.filter_precision,
                                 tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed),
                                   trainable = trainable)
+              hidden = tf.reshape(activation_fn(tf.matmul(xyz_scatter, w) + b), [-1, outputs_size[ii]])
               if self.filter_resnet_dt :
                   idt = tf.get_variable('idt_'+str(ii)+'_'+str(type_i), 
                                         [1, outputs_size[ii]], 
@@ -448,16 +450,16 @@ def _filter_type_ext(self,
                                         trainable = trainable)
               if outputs_size[ii] == outputs_size[ii-1]:
                   if self.filter_resnet_dt :
-                      xyz_scatter += activation_fn(tf.matmul(xyz_scatter, w) + b) * idt
+                      xyz_scatter += hidden * idt
                   else :
-                      xyz_scatter += activation_fn(tf.matmul(xyz_scatter, w) + b)
+                      xyz_scatter += hidden
               elif outputs_size[ii] == outputs_size[ii-1] * 2: 
                   if self.filter_resnet_dt :
-                      xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + activation_fn(tf.matmul(xyz_scatter, w) + b) * idt
+                      xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden * idt
                   else :
-                      xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + activation_fn(tf.matmul(xyz_scatter, w) + b)
+                      xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden
               else:
-                  xyz_scatter = activation_fn(tf.matmul(xyz_scatter, w) + b)
+                  xyz_scatter = hidden
             # natom x nei_type_i x out_size
             xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1]//4, outputs_size[-1]))
             # natom x nei_type_i x 4  

source/train/DescrptSeR.py

@@ -298,6 +298,7 @@ def _filter_r(self,
                                             self.filter_precision,
                                             tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed), 
                                             trainable = trainable)
+                        hidden = tf.reshape(activation_fn(tf.matmul(xyz_scatter, w) + b), [-1, outputs_size[ii]])
                         if self.filter_resnet_dt :
                             idt = tf.get_variable('idt_'+str(ii)+'_'+str(type_i), 
                                                   [1, outputs_size[ii]], 
@@ -306,16 +307,16 @@ def _filter_r(self,
                                                   trainable = trainable)
                         if outputs_size[ii] == outputs_size[ii-1]:
                             if self.filter_resnet_dt :
-                                xyz_scatter += activation_fn(tf.matmul(xyz_scatter, w) + b) * idt
+                                xyz_scatter += hidden * idt
                             else :
-                                xyz_scatter += activation_fn(tf.matmul(xyz_scatter, w) + b)
+                                xyz_scatter += hidden
                         elif outputs_size[ii] == outputs_size[ii-1] * 2: 
                             if self.filter_resnet_dt :
-                                xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + activation_fn(tf.matmul(xyz_scatter, w) + b) * idt
+                                xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden * idt
                             else :
-                                xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + activation_fn(tf.matmul(xyz_scatter, w) + b)
+                                xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden
                         else:
-                            xyz_scatter = activation_fn(tf.matmul(xyz_scatter, w) + b)
+                            xyz_scatter = hidden
                 else:
                     w = tf.zeros((outputs_size[0], outputs_size[-1]), dtype=global_tf_float_precision)
                     xyz_scatter = tf.matmul(xyz_scatter, w)

source/train/Network.py

@@ -41,9 +41,9 @@ def one_layer(inputs,
                 # return activation_fn(hidden_bn)
             else:
                 if use_timestep :
-                    return activation_fn(hidden) * idt
+                    return tf.reshape(activation_fn(hidden), [-1, outputs_size]) * idt
                 else :
-                    return activation_fn(hidden)                    
+                    return tf.reshape(activation_fn(hidden), [-1, outputs_size])                    
         else:
             if useBN:
                 None