diff options
| author |  Benoit Steiner <benoit.steiner.goog@gmail.com> | 2016-01-07 13:27:01 -0800 |
|---|---|---|
| committer |  Vijay Vasudevan <vrv@google.com> | 2016-01-07 13:27:01 -0800 |
| commit | d38fecedf54d405270377a096f58413101068792 (patch) | |
| tree | f00dc43ca235838b2c7ab566f99aedbdba9366f4 | |
| parent | 339db86e2b70a1961e877c59f3a33fa015d5a7a8 (diff) | |
Further improved the performance of the contrast adjustment code by optimizing the broadcast of scalars
Change: 111631992
| -rw-r--r-- | tensorflow/core/kernels/adjust_contrast_op.h | 9 |
1 files changed, 2 insertions, 7 deletions
diff --git a/tensorflow/core/kernels/adjust_contrast_op.h b/tensorflow/core/kernels/adjust_contrast_op.h index 2fdb2dd9a6..46dbf0ce01 100644 --- a/tensorflow/core/kernels/adjust_contrast_op.h +++ b/tensorflow/core/kernels/adjust_contrast_op.h @@ -38,14 +38,11 @@ struct AdjustContrast { Eigen::array<int, 4> scalar_broadcast{{batch, height, width, channels}}; #if !defined(EIGEN_HAS_INDEX_LIST) Eigen::array<int, 2> reduction_axis{{1, 2}}; - Eigen::array<int, 4> scalar{{1, 1, 1, 1}}; Eigen::array<int, 4> broadcast_dims{{1, height, width, 1}}; Eigen::Tensor<int, 4>::Dimensions reshape_dims{{batch, 1, 1, channels}}; #else Eigen::IndexList<Eigen::type2index<1>, Eigen::type2index<2> > reduction_axis; - Eigen::IndexList<Eigen::type2index<1>, Eigen::type2index<1>, - Eigen::type2index<1>, Eigen::type2index<1> > scalar; Eigen::IndexList<Eigen::type2index<1>, int, int, Eigen::type2index<1> > broadcast_dims; broadcast_dims.set(1, height); @@ -55,6 +52,7 @@ struct AdjustContrast { reshape_dims.set(0, batch); reshape_dims.set(3, channels); #endif + Eigen::Sizes<1, 1, 1, 1> scalar; float num_reduced_coeffs = height * width; mean_values.device(d) = (input.template cast<float>().sum(reduction_axis).eval() / @@ -88,16 +86,12 @@ struct AdjustContrastv2 { Eigen::array<int, 4> scalar_broadcast{{batch, height, width, channels}}; #if !defined(EIGEN_HAS_INDEX_LIST) Eigen::array<int, 2> reduction_axis{{0, 1}}; - Eigen::array<int, 4> scalar{{1, 1, 1, 1}}; Eigen::array<int, 4> broadcast_dims{{1, height, width, 1}}; Eigen::Tensor<int, 4>::Dimensions reshape_dims{{batch, 1, 1, channels}}; Eigen::array<int, 4> reduced_dims_first{{1, 2, 0, 3}}; #else Eigen::IndexList<Eigen::type2index<0>, Eigen::type2index<1> > reduction_axis; - Eigen::IndexList<Eigen::type2index<1>, Eigen::type2index<1>, - Eigen::type2index<1>, Eigen::type2index<1> > - scalar; Eigen::IndexList<Eigen::type2index<1>, int, int, Eigen::type2index<1> > broadcast_dims; broadcast_dims.set(1, height); @@ -110,6 +104,7 @@ struct AdjustContrastv2 { Eigen::type2index<0>, Eigen::type2index<3> > reduced_dims_first; #endif + Eigen::Sizes<1, 1, 1, 1> scalar; float num_reduced_coeffs = height * width; output.device(d) = (input.shuffle(reduced_dims_first).sum(reduction_axis).eval() / |