From ec2ffb68a5ae5c4e0e1a77e5752c43ab3517bd21 Mon Sep 17 00:00:00 2001
From: niuliling123 <51102941+niuliling123@users.noreply.github.com>
Date: Fri, 9 Apr 2021 17:44:34 +0800
Subject: [PATCH] make high precision for avg_pool and adaptive_avg_pool when
 data_type is float16 (#31887)

* make high precision for avg_pool
---
 paddle/fluid/operators/math/pooling.h | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/operators/math/pooling.h b/paddle/fluid/operators/math/pooling.h
index 21d588cc01..3547de0a4d 100644
--- a/paddle/fluid/operators/math/pooling.h
+++ b/paddle/fluid/operators/math/pooling.h
@@ -18,6 +18,7 @@ limitations under the License. */
 
 #include "paddle/fluid/framework/eigen.h"
 #include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/operators/amp/fp16_type_traits.h"
 #include "paddle/fluid/platform/device_context.h"
 #include "paddle/fluid/platform/hostdevice.h"
 #include "paddle/fluid/platform/macros.h"
@@ -46,10 +47,22 @@ class MaxPool {
 
 template <class T>
 class AvgPool {
+  using MT = typename details::MPTypeTrait<T>::Type;
+  MT intermediate_res;
+
  public:
-  DEVICE inline T initial() { return static_cast<T>(0); }
-  DEVICE inline void compute(const T& x, T* y) { *y += x; }
-  DEVICE inline void finalize(const T& pool_field, T* y) { *y /= pool_field; }
+  DEVICE inline T initial() {
+    intermediate_res = static_cast<MT>(0.0f);
+    return static_cast<T>(0);
+  }
+
+  DEVICE inline void compute(const T& x, T* y) {
+    intermediate_res += static_cast<MT>(x);
+  }
+
+  DEVICE inline void finalize(const T& pool_field, T* y) {
+    *y = static_cast<T>(intermediate_res / (static_cast<MT>(pool_field)));
+  }
 };
 
 template <class T>
-- 
GitLab