softmax.h

00001 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
00002 
00003 Licensed under the Apache License, Version 2.0 (the "License");
00004 you may not use this file except in compliance with the License.
00005 You may obtain a copy of the License at
00006 
00007     http://www.apache.org/licenses/LICENSE-2.0
00008 
00009 Unless required by applicable law or agreed to in writing, software
00010 distributed under the License is distributed on an "AS IS" BASIS,
00011 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
00012 See the License for the specific language governing permissions and
00013 limitations under the License.
00014 ==============================================================================*/
00015 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_SOFTMAX_H_
00016 #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_SOFTMAX_H_
00017 
00018 #include "fixedpoint/fixedpoint.h"
00019 #include "tensorflow/lite/kernels/internal/common.h"
00020 #include "tensorflow/lite/kernels/internal/quantization_util.h"
00021 #include "tensorflow/lite/kernels/internal/round.h"
00022 #include "tensorflow/lite/kernels/internal/types.h"
00023 #include "tensorflow/lite/kernels/op_macros.h"
00024 
00025 namespace tflite {
00026 namespace reference_ops {
00027 
00028 inline void Softmax(const SoftmaxParams& params,
00029                     const RuntimeShape& input_shape, const float* input_data,
00030                     const RuntimeShape& output_shape, float* output_data) {
00031   const int trailing_dim = input_shape.DimensionsCount() - 1;
00032   const int outer_size =
00033       MatchingFlatSizeSkipDim(input_shape, trailing_dim, output_shape);
00034   const int depth =
00035       MatchingDim(input_shape, trailing_dim, output_shape, trailing_dim);
00036 
00037   for (int i = 0; i < outer_size; ++i) {
00038     // Find max element value which we'll use to ensure numerical stability
00039     // taking advantage of the following equality:
00040     // exp(x[i])/sum(exp(x[i])) == exp(x[i]+C)/sum(exp(x[i]+C))
00041     float max = std::numeric_limits<float>::lowest();
00042     for (int c = 0; c < depth; ++c) {
00043       max = std::max(max, input_data[i * depth + c]);
00044     }
00045 
00046     // Compute sum.
00047     float sum = 0.f;
00048     for (int c = 0; c < depth; ++c) {
00049       sum += std::exp((input_data[i * depth + c] - max) * params.beta);
00050     }
00051 
00052     // Compute result.
00053     for (int c = 0; c < depth; ++c) {
00054       output_data[i * depth + c] =
00055           std::exp((input_data[i * depth + c] - max) * params.beta) / sum;
00056     }
00057   }
00058 }
00059 
00060 inline void Softmax(const SoftmaxParams& params,
00061                     const RuntimeShape& input_shape, const uint8* input_data,
00062                     const RuntimeShape& output_shape, uint8* output_data) {
00063   const int32 input_beta_multiplier = params.input_multiplier;
00064   const int32 input_beta_left_shift = params.input_left_shift;
00065   const int diff_min = params.diff_min;
00066   // The representation chosen for the input to the exp() function is Q5.26.
00067   // We need to leave extra space since values that we skip might be as large as
00068   // -32 before multiplying by input_beta_multiplier, and therefore as large as
00069   // -16 afterwards.  Note that exp(-8) is definitely not insignificant to
00070   // accumulation, but exp(-16) definitely is.
00071   static const int kScaledDiffIntegerBits = 5;
00072   static const int kAccumulationIntegerBits = 12;
00073   using FixedPointScaledDiff =
00074       gemmlowp::FixedPoint<int32, kScaledDiffIntegerBits>;
00075   using FixedPointAccum = gemmlowp::FixedPoint<int32, kAccumulationIntegerBits>;
00076   using FixedPoint0 = gemmlowp::FixedPoint<int32, 0>;
00077 
00078   const int trailing_dim = input_shape.DimensionsCount() - 1;
00079   const int outer_size =
00080       MatchingFlatSizeSkipDim(input_shape, trailing_dim, output_shape);
00081   const int depth =
00082       MatchingDim(input_shape, trailing_dim, output_shape, trailing_dim);
00083 
00084   for (int i = 0; i < outer_size; ++i) {
00085     uint8 max_in_row = 0;
00086     for (int c = 0; c < depth; ++c) {
00087       max_in_row = std::max(max_in_row, input_data[i * depth + c]);
00088     }
00089 
00090     FixedPointAccum sum_of_exps = FixedPointAccum::Zero();
00091     for (int c = 0; c < depth; ++c) {
00092       int32 input_diff =
00093           static_cast<int32>(input_data[i * depth + c]) - max_in_row;
00094       if (input_diff >= diff_min) {
00095         const int32 input_diff_rescaled =
00096             MultiplyByQuantizedMultiplierGreaterThanOne(
00097                 input_diff, input_beta_multiplier, input_beta_left_shift);
00098         const FixedPointScaledDiff scaled_diff_f8 =
00099             FixedPointScaledDiff::FromRaw(input_diff_rescaled);
00100         sum_of_exps = sum_of_exps + gemmlowp::Rescale<kAccumulationIntegerBits>(
00101                                         exp_on_negative_values(scaled_diff_f8));
00102       }
00103     }
00104 
00105     int num_bits_over_unit;
00106     FixedPoint0 shifted_scale = FixedPoint0::FromRaw(GetReciprocal(
00107         sum_of_exps.raw(), kAccumulationIntegerBits, &num_bits_over_unit));
00108 
00109     for (int c = 0; c < depth; ++c) {
00110       int32 input_diff =
00111           static_cast<int32>(input_data[i * depth + c]) - max_in_row;
00112       if (input_diff >= diff_min) {
00113         const int32 input_diff_rescaled =
00114             MultiplyByQuantizedMultiplierGreaterThanOne(
00115                 input_diff, input_beta_multiplier, input_beta_left_shift);
00116         const FixedPointScaledDiff scaled_diff_f8 =
00117             FixedPointScaledDiff::FromRaw(input_diff_rescaled);
00118 
00119         FixedPoint0 exp_in_0 = exp_on_negative_values(scaled_diff_f8);
00120         int32 unsat_output = gemmlowp::RoundingDivideByPOT(
00121             (shifted_scale * exp_in_0).raw(), num_bits_over_unit + 31 - 8);
00122 
00123         output_data[i * depth + c] = static_cast<uint8>(
00124             std::max(std::min(unsat_output, static_cast<int32>(255)),
00125                      static_cast<int32>(0)));
00126 
00127       } else {
00128         output_data[i * depth + c] = 0;
00129       }
00130     }
00131   }
00132 }
00133 
00134 // Performs softmax along the input of size (input_size * batch_size).
00135 inline void Softmax(const float* in, const int input_size, const int batch_size,
00136                     const float beta, float* out) {
00137   //  TF_LITE_ASSERT(input_size > 0);
00138 
00139   // For each batch
00140   for (int b = 0; b < batch_size; b++) {
00141     // Find the max coeff.
00142     float max_coeff = in[0];
00143     for (int i = 1; i < input_size; i++) {
00144       if (in[i] > max_coeff) max_coeff = in[i];
00145     }
00146 
00147     // Compute the normalized sum of exps.
00148     float exp_sum = 0.0;
00149     for (int i = 0; i < input_size; i++) {
00150       out[i] = std::exp((in[i] - max_coeff) * beta);
00151       exp_sum += out[i];
00152     }
00153 
00154     // Divide by the sum of exps.
00155     float reciprocal_sum_exp = 1.f / exp_sum;
00156     for (int i = 0; i < input_size; i++) {
00157       out[i] *= reciprocal_sum_exp;
00158     }
00159 
00160     // Advance in and out pointers for the next batch.
00161     in += input_size;
00162     out += input_size;
00163   }
00164 }
00165 
00166 }  // namespace reference_ops
00167 }  // namespace tflite
00168 
00169 #endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_SOFTMAX_H_