add cudnn_pool3d_op

7 years ago · 7ba3d1e4bd
parent a93a59ec7d
commit 7ba3d1e4bd
2 changed files with 11 additions and 0 deletions
--- a/paddle/operators/pool_cudnn_op.cc
+++ b/paddle/operators/pool_cudnn_op.cc
@ -23,3 +23,11 @@ REGISTER_OP_CPU_KERNEL(pool2d_cudnn,
                       ops::PoolKernel<paddle::platform::CPUPlace, float>);
 REGISTER_OP_CPU_KERNEL(pool2d_cudnn_grad,
                       ops::PoolGradKernel<paddle::platform::CPUPlace, float>)
+
+REGISTER_OP(pool3d_cudnn, ops::PoolOp, ops::Pool3dOpMaker, pool3d_cudnn_grad,
+            ops::PoolOpGrad);
+
+REGISTER_OP_CPU_KERNEL(pool3d_cudnn,
+                       ops::PoolKernel<paddle::platform::CPUPlace, float>);
+REGISTER_OP_CPU_KERNEL(pool3d_cudnn_grad,
+                       ops::PoolGradKernel<paddle::platform::CPUPlace, float>)
--- a/paddle/operators/pool_cudnn_op.cu
+++ b/paddle/operators/pool_cudnn_op.cu
@ -153,3 +153,6 @@ namespace ops = paddle::operators;

 REGISTER_OP_GPU_KERNEL(pool2d_cudnn, ops::PoolCudnnOpKernel<float>);
 REGISTER_OP_GPU_KERNEL(pool2d_cudnn_grad, ops::PoolCudnnGradOpKernel<float>);
+
+REGISTER_OP_GPU_KERNEL(pool3d_cudnn, ops::PoolCudnnOpKernel<float>);
+REGISTER_OP_GPU_KERNEL(pool3d_cudnn_grad, ops::PoolCudnnGradOpKernel<float>);