From befc26cd708f8997c78bbde5087846484ca4cdce Mon Sep 17 00:00:00 2001 From: fuzhiye Date: Thu, 13 Aug 2020 12:00:05 +0800 Subject: [PATCH] fix overflow bug when input data is big in softmax op --- .../lite/src/runtime/kernel/arm/nnacl/fp32/softmax.c | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/mindspore/lite/src/runtime/kernel/arm/nnacl/fp32/softmax.c b/mindspore/lite/src/runtime/kernel/arm/nnacl/fp32/softmax.c index 6bf03d35c4..2f79552399 100644 --- a/mindspore/lite/src/runtime/kernel/arm/nnacl/fp32/softmax.c +++ b/mindspore/lite/src/runtime/kernel/arm/nnacl/fp32/softmax.c @@ -16,6 +16,7 @@ #include "nnacl/fp32/softmax.h" #include +#include // output = exp(input) / reduce_sum(exp(input), axis) void Softmax(const float *input_ptr, float *output_ptr, float *sum_data, SoftmaxParameter *parameter) { @@ -24,8 +25,13 @@ void Softmax(const float *input_ptr, float *output_ptr, float *sum_data, Softmax int ele_size = parameter->element_size_; int *input_shape = parameter->input_shape_; + float max_data = -FLT_MAX; for (int i = 0; i < ele_size; i++) { - output_ptr[i] = exp(input_ptr[i]); + max_data = max_data > input_ptr[i] ? max_data : input_ptr[i]; + } + + for (int i = 0; i < ele_size; i++) { + output_ptr[i] = exp(input_ptr[i] - max_data); } int inner_size = 1, outter_size = 1; for (int i = 0; i < axis; i++) {