Skip to content

Commit 7c34c58

Browse files
LuLu
Lu
authored and
Lu
committed
fix bugs of gelu activation function
1 parent 5bf6f16 commit 7c34c58

File tree

10 files changed

+39
-53
lines changed

10 files changed

+39
-53
lines changed

‎source/CMakeLists.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,11 @@ include_directories(${TensorFlow_INCLUDE_DIRS})
184184
if (BUILD_CPP_IF)
185185
set (LIB_DEEPMD "deepmd")
186186
set (LIB_DEEPMD_OP "deepmd_op")
187+
if (USE_CUDA_TOOLKIT)
188+
set (LIB_DEEPMD_OP_CUDA "deepmd_op_cuda")
189+
else ()
190+
set (LIB_DEEPMD_OP_CUDA "deepmd_op")
191+
endif()
187192
if (CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 4.9)
188193
set (LIB_DEEPMD_NATIVE "deepmd_native_md")
189194
set (LIB_DEEPMD_IPI "deepmd_ipi")

‎source/lmp/env.sh.in‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,4 +8,4 @@ TF_RPATH=`echo $TENSORFLOW_LIBRARY_PATH | sed "s/;/ -Wl,-rpath=/g"`
88

99
NNP_INC=" -std=c++11 @PREC_DEF@ @TTM_DEF@ @OLD_LMP_PPPM_DEF@ -I$TF_INCLUDE_DIRS -I$DEEPMD_ROOT/include/deepmd "
1010
NNP_PATH=" -L$TF_LIBRARY_PATH -L$DEEPMD_ROOT/lib"
11-
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"
11+
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‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ if (BUILD_CPP_IF)
2424
endif (BUILD_CPP_IF)
2525

2626
if (BUILD_PY_IF)
27+
set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
28+
set(CMAKE_INSTALL_RPATH DESTINATION ${CMAKE_BINARY_DIR}/op/cuda)
2729
if (USE_CUDA_TOOLKIT)
2830
add_library(op_abi SHARED ${OP_PY_CUDA_SRC} ${OP_LIB})
2931
add_library(op_grads SHARED ${OP_GRADS_SRC})
@@ -33,11 +35,11 @@ if (BUILD_PY_IF)
3335
set (EXTRA_LIBS ${EXTRA_LIBS} deepmd_op_cuda)
3436
target_link_libraries (op_abi ${EXTRA_LIBS})
3537
target_link_libraries (op_grads ${EXTRA_LIBS})
36-
message(STATUS ${TensorFlowFramework_LIBRARY})
3738
else (USE_CUDA_TOOLKIT)
3839
add_library(op_abi SHARED ${OP_SRC} ${OP_LIB})
3940
add_library(op_grads SHARED ${OP_GRADS_SRC})
4041
endif(USE_CUDA_TOOLKIT)
42+
message(STATUS ${TensorFlowFramework_LIBRARY})
4143
target_link_libraries(
4244
op_abi ${TensorFlowFramework_LIBRARY}
4345
)

‎source/op/_gelu.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,9 @@
77
from deepmd.env import op_module
88

99
@ops.RegisterGradient("Gelu")
10-
def gelu_cc (op, dy) :
10+
def _gelu_cc (op, dy) :
1111
return op_module.gelu_grad(dy, op.inputs[0])
1212

1313
@ops.RegisterGradient("GeluGrad")
14-
def gelu_grad_cc (op, dy) :
15-
return [None, op_module.gelu_grad_grad(dy, op.inputs[0], op.inputs[1])]
14+
def _gelu_grad_cc (op, dy) :
15+
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‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ set (SOURCE_FILES
8383
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
8484
)
8585

86-
cuda_add_library(deepmd_op_cuda STATIC ${SOURCE_FILES})
86+
cuda_add_library(deepmd_op_cuda SHARED ${SOURCE_FILES})
8787

8888
if (BUILD_CPP_IF)
8989
install(TARGETS deepmd_op_cuda DESTINATION lib/)

‎source/op/gelu.cc‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,32 +11,20 @@ using GPUDevice = Eigen::GpuDevice;
1111
REGISTER_OP("Gelu")
1212
.Attr("T: {float, double}")
1313
.Input("x: T")
14-
.Output("output: T")
15-
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
16-
c->set_output(0, c->input(0));
17-
return Status::OK();
18-
});
14+
.Output("output: T");
1915

2016
REGISTER_OP("GeluGrad")
2117
.Attr("T: {float, double}")
2218
.Input("dy: T")
2319
.Input("x: T")
24-
.Output("output: T")
25-
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
26-
c->set_output(0, c->input(1));
27-
return Status::OK();
28-
});
20+
.Output("output: T");
2921

3022
REGISTER_OP("GeluGradGrad")
3123
.Attr("T: {float, double}")
3224
.Input("dy: T")
3325
.Input("dy_: T")
3426
.Input("x: T")
35-
.Output("output: T")
36-
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
37-
c->set_output(0, c->input(2));
38-
return Status::OK();
39-
});
27+
.Output("output: T");
4028

4129
template <typename Device, typename T>
4230
struct GeluFunctor {

‎source/op/gelu_gpu.cc‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -10,32 +10,20 @@ using GPUDevice = Eigen::GpuDevice;
1010
REGISTER_OP("Gelu")
1111
.Attr("T: {float, double}")
1212
.Input("x: T")
13-
.Output("output: T")
14-
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
15-
c->set_output(0, c->input(0));
16-
return Status::OK();
17-
});
13+
.Output("output: T");
1814

1915
REGISTER_OP("GeluGrad")
2016
.Attr("T: {float, double}")
2117
.Input("dy: T")
2218
.Input("x: T")
23-
.Output("output: T")
24-
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
25-
c->set_output(0, c->input(1));
26-
return Status::OK();
27-
});
19+
.Output("output: T");
2820

2921
REGISTER_OP("GeluGradGrad")
3022
.Attr("T: {float, double}")
3123
.Input("dy: T")
3224
.Input("dy_: T")
3325
.Input("x: T")
34-
.Output("output: T")
35-
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
36-
c->set_output(0, c->input(2));
37-
return Status::OK();
38-
});
26+
.Output("output: T");
3927

4028
// maybe instead use cudnn activation forward
4129
void GeluLauncher(const float * in, float * out, int const size);

‎source/train/DescrptSeA.py‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -353,6 +353,7 @@ def _filter(self,
353353
self.filter_precision,
354354
tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed),
355355
trainable = trainable)
356+
hidden = tf.reshape(activation_fn(tf.matmul(xyz_scatter, w) + b), [-1, outputs_size[ii]])
356357
if self.filter_resnet_dt :
357358
idt = tf.get_variable('idt_'+str(ii)+'_'+str(type_i),
358359
[1, outputs_size[ii]],
@@ -361,16 +362,16 @@ def _filter(self,
361362
trainable = trainable)
362363
if outputs_size[ii] == outputs_size[ii-1]:
363364
if self.filter_resnet_dt :
364-
xyz_scatter += activation_fn(tf.matmul(xyz_scatter, w) + b) * idt
365+
xyz_scatter += hidden * idt
365366
else :
366-
xyz_scatter += activation_fn(tf.matmul(xyz_scatter, w) + b)
367+
xyz_scatter += hidden
367368
elif outputs_size[ii] == outputs_size[ii-1] * 2:
368369
if self.filter_resnet_dt :
369-
xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + activation_fn(tf.matmul(xyz_scatter, w) + b) * idt
370+
xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden * idt
370371
else :
371-
xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + activation_fn(tf.matmul(xyz_scatter, w) + b)
372+
xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden
372373
else:
373-
xyz_scatter = activation_fn(tf.matmul(xyz_scatter, w) + b)
374+
xyz_scatter = hidden
374375
else:
375376
w = tf.zeros((outputs_size[0], outputs_size[-1]), dtype=global_tf_float_precision)
376377
xyz_scatter = tf.matmul(xyz_scatter, w)
@@ -440,6 +441,7 @@ def _filter_type_ext(self,
440441
self.filter_precision,
441442
tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed),
442443
trainable = trainable)
444+
hidden = tf.reshape(activation_fn(tf.matmul(xyz_scatter, w) + b), [-1, outputs_size[ii]])
443445
if self.filter_resnet_dt :
444446
idt = tf.get_variable('idt_'+str(ii)+'_'+str(type_i),
445447
[1, outputs_size[ii]],
@@ -448,16 +450,16 @@ def _filter_type_ext(self,
448450
trainable = trainable)
449451
if outputs_size[ii] == outputs_size[ii-1]:
450452
if self.filter_resnet_dt :
451-
xyz_scatter += activation_fn(tf.matmul(xyz_scatter, w) + b) * idt
453+
xyz_scatter += hidden * idt
452454
else :
453-
xyz_scatter += activation_fn(tf.matmul(xyz_scatter, w) + b)
455+
xyz_scatter += hidden
454456
elif outputs_size[ii] == outputs_size[ii-1] * 2:
455457
if self.filter_resnet_dt :
456-
xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + activation_fn(tf.matmul(xyz_scatter, w) + b) * idt
458+
xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden * idt
457459
else :
458-
xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + activation_fn(tf.matmul(xyz_scatter, w) + b)
460+
xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden
459461
else:
460-
xyz_scatter = activation_fn(tf.matmul(xyz_scatter, w) + b)
462+
xyz_scatter = hidden
461463
# natom x nei_type_i x out_size
462464
xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1]//4, outputs_size[-1]))
463465
# natom x nei_type_i x 4

‎source/train/DescrptSeR.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -298,6 +298,7 @@ def _filter_r(self,
298298
self.filter_precision,
299299
tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed),
300300
trainable = trainable)
301+
hidden = tf.reshape(activation_fn(tf.matmul(xyz_scatter, w) + b), [-1, outputs_size[ii]])
301302
if self.filter_resnet_dt :
302303
idt = tf.get_variable('idt_'+str(ii)+'_'+str(type_i),
303304
[1, outputs_size[ii]],
@@ -306,16 +307,16 @@ def _filter_r(self,
306307
trainable = trainable)
307308
if outputs_size[ii] == outputs_size[ii-1]:
308309
if self.filter_resnet_dt :
309-
xyz_scatter += activation_fn(tf.matmul(xyz_scatter, w) + b) * idt
310+
xyz_scatter += hidden * idt
310311
else :
311-
xyz_scatter += activation_fn(tf.matmul(xyz_scatter, w) + b)
312+
xyz_scatter += hidden
312313
elif outputs_size[ii] == outputs_size[ii-1] * 2:
313314
if self.filter_resnet_dt :
314-
xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + activation_fn(tf.matmul(xyz_scatter, w) + b) * idt
315+
xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden * idt
315316
else :
316-
xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + activation_fn(tf.matmul(xyz_scatter, w) + b)
317+
xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden
317318
else:
318-
xyz_scatter = activation_fn(tf.matmul(xyz_scatter, w) + b)
319+
xyz_scatter = hidden
319320
else:
320321
w = tf.zeros((outputs_size[0], outputs_size[-1]), dtype=global_tf_float_precision)
321322
xyz_scatter = tf.matmul(xyz_scatter, w)

‎source/train/Network.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,9 @@ def one_layer(inputs,
4141
# return activation_fn(hidden_bn)
4242
else:
4343
if use_timestep :
44-
return activation_fn(hidden) * idt
44+
return tf.reshape(activation_fn(hidden), [-1, outputs_size]) * idt
4545
else :
46-
return activation_fn(hidden)
46+
return tf.reshape(activation_fn(hidden), [-1, outputs_size])
4747
else:
4848
if useBN:
4949
None

0 commit comments

Comments
 (0)