mean accuracy fix

Author: NKNaN

paddle/phi/kernels/reduce_mean_kernel.cc

@@ -16,6 +16,7 @@
 
 #include "paddle/phi/backends/all_context.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/cast_kernel.h"
 #include "paddle/phi/kernels/reduce_kernel_impl.h"
 
 namespace phi {
@@ -27,7 +28,24 @@ void MeanKernel(const Context& dev_ctx,
  bool keep_dim,
  DenseTensor* out) {
  bool reduce_all = recompute_reduce_all(x, dims);
- MeanRawKernel<T>(dev_ctx, x, dims, keep_dim, reduce_all, out);
+ if (std::is_same<T, int>::value || std::is_same<T, int64_t>::value ||
+ std::is_same<T, bool>::value) {
+ using Type =
+ typename std::conditional<std::is_same<T, int>::value ||
+ std::is_same<T, int64_t>::value ||
+ std::is_same<T, bool>::value,
+ float,
+ T>::type;
+ DenseTensor x_float =
+ phi::Cast<T, Context>(dev_ctx, x, phi::DataType::FLOAT32);
+ DenseTensor* out_float = new DenseTensor();
+ out_float->Resize(out->dims());
+ MeanRawKernel<Type>(
+ dev_ctx, x_float, dims, keep_dim, reduce_all, out_float);
+ phi::CastKernel<Type, Context>(dev_ctx, *out_float, x.dtype(), out);
+ } else {
+ MeanRawKernel<T>(dev_ctx, x, dims, keep_dim, reduce_all, out);
+ }
 }
 
 } // namespace phi

test/legacy_test/test_mean_op.py

@@ -806,6 +806,64 @@ def test_errors(self):
  )
 
 
+class TestMeanAPIInt32(unittest.TestCase):
+ def setUp(self):
+ self.x_shape = [2, 3, 4, 5]
+ self.dtype = "int32"
+ self.x_np = np.random.randint(-1, 10000, self.x_shape).astype(
+ self.dtype
+ )
+ self.places = [paddle.CPUPlace()]
+ if core.is_compiled_with_cuda():
+ self.places.append(paddle.CUDAPlace(0))
+
+ def test_dygraph(self):
+ for place in self.places:
+ with base.dygraph.guard(place):
+ x = paddle.to_tensor(self.x_np)
+ out = paddle.mean(x=x)
+ np.testing.assert_equal(
+ out.numpy(),
+ np.mean(self.x_np.astype("float32")).astype(self.dtype),
+ )
+
+ def test_static(self):
+ paddle.enable_static()
+ for place in self.places:
+ with base.program_guard(base.Program(), base.Program()):
+ x = paddle.static.data(
+ "x", shape=self.x_shape, dtype=self.dtype
+ )
+ out = paddle.mean(x=x)
+ exe = base.Executor(place)
+ res = exe.run(feed={"x": self.x_np}, fetch_list=[out])
+ np.testing.assert_equal(
+ res[0], np.mean(self.x_np.astype("float32")).astype(self.dtype)
+ )
+
+
+class TestMeanAPIInt64(TestMeanAPIInt32):
+ def setUp(self):
+ self.x_shape = [2, 3, 4, 5]
+ self.dtype = "int64"
+ self.x_np = np.random.randint(-1, 10000, self.x_shape).astype(
+ self.dtype
+ )
+ self.places = [paddle.CPUPlace()]
+ if core.is_compiled_with_cuda():
+ self.places.append(paddle.CUDAPlace(0))
+
+
+class TestMeanAPIBool(TestMeanAPIInt32):
+ def setUp(self):
+ self.x_shape = [2, 3, 4, 5]
+ self.dtype = "bool"
+ self.x_np = np.random.uniform(-1, 1, self.x_shape).astype(self.dtype)
+ self.places = [paddle.CPUPlace()]
+ if core.is_compiled_with_cuda():
+ self.places.append(paddle.CUDAPlace(0))
+
+
 class TestMeanWithTensorAxis1(TestReduceOPTensorAxisBase):
  def init_data(self):
  self.pd_api = paddle.mean