From 78f32bab68386f48a336c616ac462057e3d6a3bd Mon Sep 17 00:00:00 2001 From: zhangqi Date: Mon, 11 Jul 2022 17:29:46 +0800 Subject: [PATCH] fix gamma bug --- .../ops/mindspore.ops.func_random_gamma.rst | 4 +- .../cpu/kernel/eigen/gamma_cpu_kernel.cc | 53 ++++++++++--------- .../mindspore/ops/operations/random_ops.py | 4 +- tests/st/ops/cpu/test_random_gamma_op.py | 4 +- 4 files changed, 34 insertions(+), 31 deletions(-) diff --git a/docs/api/api_python/ops/mindspore.ops.func_random_gamma.rst b/docs/api/api_python/ops/mindspore.ops.func_random_gamma.rst index 260c839205f..1705e2c1f78 100644 --- a/docs/api/api_python/ops/mindspore.ops.func_random_gamma.rst +++ b/docs/api/api_python/ops/mindspore.ops.func_random_gamma.rst @@ -1,7 +1,7 @@ -mindspore.ops.gamma +mindspore.ops.random_gamma ==================== -.. py:function:: mindspore.ops.gamma(shape, alpha, seed=None) +.. py:function:: mindspore.ops.random_gamma(shape, alpha, seed=None) 根据伽马分布产生成随机数。 diff --git a/mindspore/ccsrc/plugin/device/cpu/kernel/eigen/gamma_cpu_kernel.cc b/mindspore/ccsrc/plugin/device/cpu/kernel/eigen/gamma_cpu_kernel.cc index 37ecdd53285..808d4e657a2 100644 --- a/mindspore/ccsrc/plugin/device/cpu/kernel/eigen/gamma_cpu_kernel.cc +++ b/mindspore/ccsrc/plugin/device/cpu/kernel/eigen/gamma_cpu_kernel.cc @@ -44,12 +44,25 @@ bool GammaCpuKernelMod::Init(const BaseOperatorPtr &base_operator, const std::ve alpha_shape_ = inputs[1]->GetShapeVector(); alpha_dtype_ = inputs[1]->GetDtype(); shape_dtype_ = inputs[0]->GetDtype(); + shape_shape_ = inputs[0]->GetShapeVector(); is_need_retrieve_output_shape_ = true; return true; } +template +void GammaCpuKernelMod::InferShape(const std::vector &inputs) { + const auto *shape_value = reinterpret_cast(inputs[0]->addr); + + for (int64_t i = 0; i < shape_shape_[0]; i++) { + output_shape_.emplace_back(static_cast(shape_value[i])); + } + for (size_t i = 0; i < alpha_shape_.size(); i++) { + output_shape_.emplace_back(alpha_shape_[i]); + } +} + int GammaCpuKernelMod::Resize(const BaseOperatorPtr &base_operator, const std::vector &inputs, const std::vector &outputs, const std::map &inputsOnHost) { @@ -62,9 +75,6 @@ int GammaCpuKernelMod::Resize(const BaseOperatorPtr &base_operator, const std::v return ret; } - shape_shape_ = inputs[0]->GetShapeVector(); - alpha_shape_ = inputs[1]->GetShapeVector(); - return KRET_OK; } @@ -74,7 +84,7 @@ void GammaCpuKernelMod::Generate(const std::vector &inputs, const st const auto *alpha_flat = reinterpret_cast(inputs[1]->addr); auto *samples_flat = reinterpret_cast(outputs[0]->addr); - int64_t num_samples = std::accumulate(output_shape_.begin(), output_shape_.end(), 0); + int64_t num_samples = std::accumulate(output_shape_.begin(), output_shape_.end(), 1, std::multiplies()); if (num_samples == 0) { MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "' the sizes of output is zero."; } @@ -92,11 +102,16 @@ void GammaCpuKernelMod::Generate(const std::vector &inputs, const st static constexpr int kReservedSamplesPerOutput = 256; - int64 num_alphas = std::accumulate(alpha_shape_.begin(), alpha_shape_.end(), 0); + int64 num_alphas = std::accumulate(alpha_shape_.begin(), alpha_shape_.end(), 1, std::multiplies()); + if (num_alphas == 0) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "' the sizes of alpha is zero."; + } + int64_t sample_shape_per_al = num_samples / num_alphas; - PhiloxRandom rng = generator_.ReserveRandomOutputs(num_samples * num_alphas, kReservedSamplesPerOutput); + PhiloxRandom rng = generator_.ReserveRandomOutputs(num_samples, kReservedSamplesPerOutput); - auto DoWork = [num_samples, num_alphas, &rng, samples_flat, alpha_flat](int64 start_output, int64 limit_output) { + auto DoWork = [sample_shape_per_al, num_alphas, &rng, samples_flat, alpha_flat](int64 start_output, + int64 limit_output) { using Eigen::numext::exp; using Eigen::numext::log; using Eigen::numext::pow; @@ -107,7 +122,7 @@ void GammaCpuKernelMod::Generate(const std::vector &inputs, const st typename Uniform::ResultType uniform_res; for (int64 output_idx = start_output; output_idx < limit_output;) { - int64 alpha_idx = output_idx / num_samples; + int64 alpha_idx = output_idx / sample_shape_per_al; T *const samples_alpha_offset = samples_flat + alpha_idx; const double alpha_value = static_cast(alpha_flat[alpha_idx]); @@ -115,8 +130,8 @@ void GammaCpuKernelMod::Generate(const std::vector &inputs, const st if (alpha_value == static_cast(1.0)) { // ENABLE_FLOAT_EQUALITY_WARNING // Sample from an exponential distribution. - for (int64 sample_idx = output_idx % num_samples; sample_idx < num_samples && output_idx < limit_output; - sample_idx++, output_idx++) { + for (int64 sample_idx = output_idx % sample_shape_per_al; + sample_idx < sample_shape_per_al && output_idx < limit_output; sample_idx++, output_idx++) { PhiloxRandom gen = rng; gen.Skip(kReservedSamplesPerOutput * output_idx); int16 uniform_remaining = 0; @@ -132,8 +147,8 @@ void GammaCpuKernelMod::Generate(const std::vector &inputs, const st const double cut = 1.0 / 3 / sqrt(su); // Compute the rest of the samples for the current alpha value. - for (int64 sample_idx = output_idx % num_samples; sample_idx < num_samples && output_idx < limit_output; - sample_idx++, output_idx++) { + for (int64 sample_idx = output_idx % sample_shape_per_al; + sample_idx < sample_shape_per_al && output_idx < limit_output; sample_idx++, output_idx++) { PhiloxRandom gen = rng; gen.Skip(kReservedSamplesPerOutput * output_idx); int16 norm_remaining = 0; @@ -171,19 +186,7 @@ void GammaCpuKernelMod::Generate(const std::vector &inputs, const st } }; #undef UNIFORM - ParallelLaunchAutoSearch(DoWork, num_alphas * num_samples, this, ¶llel_search_info_); -} - -template -void GammaCpuKernelMod::InferShape(const std::vector &inputs) { - const auto *shape_value = reinterpret_cast(inputs[0]->addr); - - for (int64_t i = 0; i < shape_shape_[0]; i++) { - output_shape_.emplace_back(static_cast(shape_value[i])); - } - for (size_t i = 0; i < alpha_shape_.size(); i++) { - output_shape_.emplace_back(alpha_shape_[i]); - } + ParallelLaunchAutoSearch(DoWork, num_alphas * sample_shape_per_al, this, ¶llel_search_info_); } bool GammaCpuKernelMod::Launch(const std::vector &inputs, const std::vector &workspace, diff --git a/mindspore/python/mindspore/ops/operations/random_ops.py b/mindspore/python/mindspore/ops/operations/random_ops.py index ef74a975407..3112d16629e 100644 --- a/mindspore/python/mindspore/ops/operations/random_ops.py +++ b/mindspore/python/mindspore/ops/operations/random_ops.py @@ -321,8 +321,8 @@ class Gamma(PrimitiveWithInfer): >>> shape = (3, 1, 2) >>> alpha = Tensor(np.array([[3, 4], [5, 6]]), mstype.float32) >>> beta = Tensor(np.array([1.0]), mstype.float32) - >>> random_gamma = ops.RandomGamma(seed=3) - >>> output = random_gamma(shape, alpha, beta) + >>> gamma = ops.Gamma(seed=3) + >>> output = gamma(shape, alpha, beta) >>> result = output.shape >>> print(result) (3, 2, 2) diff --git a/tests/st/ops/cpu/test_random_gamma_op.py b/tests/st/ops/cpu/test_random_gamma_op.py index 3685382cbef..9b546efc627 100644 --- a/tests/st/ops/cpu/test_random_gamma_op.py +++ b/tests/st/ops/cpu/test_random_gamma_op.py @@ -71,10 +71,10 @@ def test_random_gamma_functional(dtype): """ ms.set_context(mode=ms.GRAPH_MODE, device_target='CPU') - shape = Tensor(np.array([3, 2]), ms.int32) + shape = Tensor(np.array([10, 10]), ms.int32) alpha = Tensor(np.array([[3, 4], [5, 6]]), dtype) output = F.random_gamma(shape=shape, alpha=alpha, seed=2) - expect = np.array([3, 2, 2, 2]) + expect = np.array([10, 10, 2, 2]) print(output) assert (output.shape == expect).all()