Adding an environment variable to control whether to use Cudnn with A…

…vgPool: TF_AVGPOOL_USE_CUDNN. The default is false for now. Change: 127123594
tensorflow · Jul 11, 2016 · 7f37ded · 7f37ded · cancan101 · Mar 30, 2017
1 parent fe436f7
commit 7f37ded
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 1 deletion.
diff --git a/tensorflow/core/kernels/avgpooling_op.cc b/tensorflow/core/kernels/avgpooling_op.cc
@@ -39,6 +39,7 @@ limitations under the License.
 #if GOOGLE_CUDA
 #include "tensorflow/core/kernels/maxpooling_op_gpu.h"
 #include "tensorflow/core/kernels/pooling_ops_common_gpu.h"
+#include "tensorflow/core/util/use_cudnn.h"
 #endif  // GOOGLE_CUDA
 
 namespace tensorflow {
@@ -150,7 +151,7 @@ class AvgPoolingOp<GPUDevice, T> : public UnaryOp<T> {
 
     TensorShape output_shape = params.forward_output_shape();
 
-    if (data_format_ == FORMAT_NCHW) {
+    if (internal::AvgPoolUseCudnn() || data_format_ == FORMAT_NCHW) {
       DnnPoolingOp<T>::Compute(
           context, perftools::gputools::dnn::PoolingMode::kAverage, ksize_,
           stride_, padding_, data_format_, tensor_in, output_shape);

diff --git a/tensorflow/core/util/use_cudnn.cc b/tensorflow/core/util/use_cudnn.cc
@@ -40,4 +40,11 @@ bool CudnnUseAutotune() {
   return ReadBoolFromEnvVar("TF_CUDNN_USE_AUTOTUNE", true);
 }
 
+namespace internal {
+
+bool AvgPoolUseCudnn() {
+  return ReadBoolFromEnvVar("TF_AVGPOOL_USE_CUDNN", false);
+}
+
+}  // namespace internal
 }  // namespace tensorflow
diff --git a/tensorflow/core/util/use_cudnn.h b/tensorflow/core/util/use_cudnn.h
@@ -23,6 +23,12 @@ namespace tensorflow {
 bool CanUseCudnn();
 bool CudnnUseAutotune();
 
+namespace internal {
+
+// This function is for transition only. And it may go away at any time.
+bool AvgPoolUseCudnn();
+
+}  // namespace internal
 }  // namespace tensorflow
 
 #endif  // TENSORFLOW_UTIL_USE_CUDNN_H_