diff options
author | Benoit Steiner <benoit.steiner.goog@gmail.com> | 2016-01-06 12:56:08 -0800 |
---|---|---|
committer | Vijay Vasudevan <vrv@google.com> | 2016-01-06 12:56:08 -0800 |
commit | b6acce62fb3e426dc4521fbf366762bb4c02d28f (patch) | |
tree | cf71b433e21a68cd957d0fe56ff63e78e720306d | |
parent | 5b2e90ae734c1c42c7b7c42db56df53b4c48c6a8 (diff) |
Improved the performance of the contrast adjustment op by manually computing
the mean.
Change: 111535650
-rw-r--r-- | tensorflow/core/kernels/adjust_contrast_op.h | 9 |
1 files changed, 5 insertions, 4 deletions
diff --git a/tensorflow/core/kernels/adjust_contrast_op.h b/tensorflow/core/kernels/adjust_contrast_op.h index 7141dd1b9c..514373a0b6 100644 --- a/tensorflow/core/kernels/adjust_contrast_op.h +++ b/tensorflow/core/kernels/adjust_contrast_op.h @@ -106,10 +106,11 @@ struct AdjustContrastv2 { reshape_dims.set(0, batch); reshape_dims.set(3, channels); #endif - mean_values.device(d) = input.mean(reduction_axis) - .eval() - .reshape(reshape_dims) - .broadcast(broadcast_dims); + float num_reduced_coeffs = height * width; + mean_values.device(d) = + (input.sum(reduction_axis).eval() / num_reduced_coeffs) + .reshape(reshape_dims) + .broadcast(broadcast_dims); auto contrast_factor_tensor = contrast_factor.reshape(scalar).broadcast(scalar_broadcast); auto adjusted = |