fix global pool with default alg

NVIDIA · Dec 27, 2020 · 2d31312 · 2d31312
1 parent 8f53bf9
commit 2d31312
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 6 deletions.
diff --git a/src/global_pooling_cpu.cpp b/src/global_pooling_cpu.cpp
@@ -225,15 +225,19 @@ GlobalPoolingBackwardCPU(at::Tensor const &in_feat, at::Tensor &grad_out_feat,
   // if (pooling_mode == 0)
   //   pooling_mode = in_feat.size(0) / batch_size > 100 ? 1 : 2;
 
-  if (pooling_mode == PoolingMode::GLOBAL_SUM_POOLING_KERNEL ||
+  if (pooling_mode == PoolingMode::GLOBAL_SUM_POOLING_DEFAULT ||
+      pooling_mode == PoolingMode::GLOBAL_AVG_POOLING_DEFAULT ||
+      pooling_mode == PoolingMode::GLOBAL_SUM_POOLING_KERNEL ||
       pooling_mode == PoolingMode::GLOBAL_AVG_POOLING_KERNEL ||
       pooling_mode == PoolingMode::GLOBAL_SUM_POOLING_PYTORCH_INDEX ||
       pooling_mode == PoolingMode::GLOBAL_AVG_POOLING_PYTORCH_INDEX) {
 
+    LOG_DEBUG("GLOBAL_POOLING");
     if (batch_size == 1) {
-      if (use_avg)
+      if (use_avg) {
+        LOG_DEBUG("Copying grad_out_feat. size:", in_feat.size(0));
         grad_in_feat.copy_(grad_out_feat / in_feat.size(0));
-      else
+      } else
         grad_in_feat.copy_(grad_out_feat);
     } else {
       const auto &in_outs = p_map_manager->origin_map(p_in_map_key);

diff --git a/src/global_pooling_gpu.cu b/src/global_pooling_gpu.cu
@@ -242,15 +242,19 @@ at::Tensor GlobalPoolingBackwardGPU(
   //   pooling_mode = in_feat.size(0) / batch_size > 100 ? 1 : 2;
 
   cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
-  if (pooling_mode == PoolingMode::GLOBAL_SUM_POOLING_KERNEL ||
+  if (pooling_mode == PoolingMode::GLOBAL_SUM_POOLING_DEFAULT ||
+      pooling_mode == PoolingMode::GLOBAL_AVG_POOLING_DEFAULT ||
+      pooling_mode == PoolingMode::GLOBAL_SUM_POOLING_KERNEL ||
       pooling_mode == PoolingMode::GLOBAL_AVG_POOLING_KERNEL ||
       pooling_mode == PoolingMode::GLOBAL_SUM_POOLING_PYTORCH_INDEX ||
       pooling_mode == PoolingMode::GLOBAL_AVG_POOLING_PYTORCH_INDEX) {
 
+    LOG_DEBUG("GLOBAL_POOLING");
     if (batch_size == 1) {
-      if (use_avg)
+      if (use_avg) {
+        LOG_DEBUG("Copying grad_out_feat. size:", in_feat.size(0));
         grad_in_feat.copy_(grad_out_feat / in_feat.size(0));
-      else
+      } else
         grad_in_feat.copy_(grad_out_feat);
     } else {
       const auto &in_outs = p_map_manager->origin_map(p_in_map_key);