replace spatial with per_activation mode for bn op to improve perf (#33887)

3629bf4f · Li Min · GitHub · eae31856 · 3629bf4f
显示空白变更内容
内联并排

Showing with 14 addition and 2 deletion

paddle/fluid/operators/batch_norm_op.cu paddle/fluid/operators/batch_norm_op.cu +14 -2

未找到文件。
--- a/paddle/fluid/operators/batch_norm_op.cu
+++ b/paddle/fluid/operators/batch_norm_op.cu
@@ -225,11 +225,17 @@ class BatchNormKernel<platform::CUDADeviceContext, T>
 #elif CUDNN_VERSION_MIN(7, 0, 1)
    if (FLAGS_cudnn_batchnorm_spatial_persistent) {
      mode_ = CUDNN_BATCHNORM_SPATIAL_PERSISTENT;
+    } else if (H == 1 && W == 1) {
+      mode_ = CUDNN_BATCHNORM_PER_ACTIVATION;
    } else {
      mode_ = CUDNN_BATCHNORM_SPATIAL;
    }
 #else
+    if (H == 1 && W == 1) {
+      mode_ = CUDNN_BATCHNORM_PER_ACTIVATION;
+    } else {
      mode_ = CUDNN_BATCHNORM_SPATIAL;
+    }
 #endif  // CUDNN_VERSION_MIN(7, 0, 1)

    VLOG(3) << "Setting descriptors.";
@@ -989,11 +995,17 @@ class BatchNormGradKernel<platform::CUDADeviceContext, T>
 #elif CUDNN_VERSION_MIN(7, 0, 1)
      if (FLAGS_cudnn_batchnorm_spatial_persistent) {
        mode_ = CUDNN_BATCHNORM_SPATIAL_PERSISTENT;
+      } else if (H == 1 && W == 1) {
+        mode_ = CUDNN_BATCHNORM_PER_ACTIVATION;
      } else {
        mode_ = CUDNN_BATCHNORM_SPATIAL;
      }
 #else
+      if (H == 1 && W == 1) {
+        mode_ = CUDNN_BATCHNORM_PER_ACTIVATION;
+      } else {
        mode_ = CUDNN_BATCHNORM_SPATIAL;
+      }
 #endif  // CUDNN_VERSION_MIN(7, 0, 1)

 #ifdef PADDLE_WITH_HIP