Remove Softmax dimension checks: all implementations work on last dimension, other dimensions are just used as batches.

PiperOrigin-RevId: 302719938
Change-Id: I56e5d6715b58f9c41ea902d04fe2e0d7d33df88c
This commit is contained in:
Robert David 2020-03-24 12:17:36 -07:00 committed by TensorFlower Gardener
parent 8ed969c39d
commit cd661aab96
5 changed files with 61 additions and 303 deletions

View File

@ -544,8 +544,7 @@ TfLiteStatus SoftmaxPrepare(TfLiteContext* context, TfLiteNode* node) {
TfLiteTensor* output = GetOutput(context, node, 0);
TF_LITE_ENSURE_EQ(context, input->type, output->type);
const int num_dims = NumDimensions(input);
TF_LITE_ENSURE(context, num_dims >= 1 && num_dims <= 4);
TF_LITE_ENSURE(context, NumDimensions(input) >= 1);
if (input->type == kTfLiteUInt8 || input->type == kTfLiteInt8) {
data->params.table = data->table;
@ -915,41 +914,22 @@ TfLiteStatus SigmoidEval(TfLiteContext* context, TfLiteNode* node) {
TfLiteStatus SoftmaxFloat(TfLiteContext* context, const TfLiteTensor* input,
TfLiteTensor* output, TfLiteSoftmaxParams* params) {
switch (NumDimensions(input)) {
case 1:
case 2:
case 3:
case 4:
SoftmaxParams op_params;
op_params.beta = params->beta;
optimized_ops::Softmax(
op_params, GetTensorShape(input), GetTensorData<float>(input),
GetTensorShape(output), GetTensorData<float>(output),
CpuBackendContext::GetFromContext(context));
return kTfLiteOk;
default:
TF_LITE_KERNEL_LOG(
context,
"Only 1D, 2D, 3D and 4D tensors supported currently, got %dD.",
NumDimensions(input));
return kTfLiteError;
}
SoftmaxParams op_params;
op_params.beta = params->beta;
optimized_ops::Softmax(op_params, GetTensorShape(input),
GetTensorData<float>(input), GetTensorShape(output),
GetTensorData<float>(output),
CpuBackendContext::GetFromContext(context));
return kTfLiteOk;
}
template <typename T>
TfLiteStatus SoftmaxQuantized(TfLiteContext* context, const TfLiteTensor* input,
TfLiteTensor* output, SoftmaxOpData* data) {
if (NumDimensions(input) >= 1 && NumDimensions(input) <= 4) {
optimized_ops::Softmax(data->params, GetTensorShape(input),
GetTensorData<T>(input), GetTensorShape(output),
GetTensorData<T>(output));
return kTfLiteOk;
} else {
TF_LITE_KERNEL_LOG(
context, "Only 1D, 2D, 3D and 4D tensors supported currently, got %dD.",
NumDimensions(input));
return kTfLiteError;
}
optimized_ops::Softmax(data->params, GetTensorShape(input),
GetTensorData<T>(input), GetTensorShape(output),
GetTensorData<T>(output));
return kTfLiteOk;
}
TfLiteStatus SoftmaxEval(TfLiteContext* context, TfLiteNode* node) {
@ -959,8 +939,6 @@ TfLiteStatus SoftmaxEval(TfLiteContext* context, TfLiteNode* node) {
const TfLiteTensor* input = GetInput(context, node, 0);
TfLiteTensor* output = GetOutput(context, node, 0);
// TODO(ahentz): consider an implementation that works for many (all?)
// dimensions.
switch (input->type) {
case kTfLiteFloat32: {
return SoftmaxFloat(context, input, output, params);

View File

@ -142,38 +142,6 @@ inline void Softmax(const SoftmaxParams& params,
}
}
// Performs softmax along the input of size (input_size * batch_size).
inline void Softmax(const float* in, const int input_size, const int batch_size,
const float beta, float* out) {
// TF_LITE_ASSERT(input_size > 0);
// For each batch
for (int b = 0; b < batch_size; b++) {
// Find the max coeff.
float max_coeff = in[0];
for (int i = 1; i < input_size; i++) {
if (in[i] > max_coeff) max_coeff = in[i];
}
// Compute the normalized sum of exps.
float exp_sum = 0.0;
for (int i = 0; i < input_size; i++) {
out[i] = std::exp((in[i] - max_coeff) * beta);
exp_sum += out[i];
}
// Divide by the sum of exps.
float reciprocal_sum_exp = 1.f / exp_sum;
for (int i = 0; i < input_size; i++) {
out[i] *= reciprocal_sum_exp;
}
// Advance in and out pointers for the next batch.
in += input_size;
out += input_size;
}
}
} // namespace reference_ops
} // namespace tflite

View File

@ -67,70 +67,17 @@ void* Init(TfLiteContext* context, const char* buffer, size_t length) {
void Free(TfLiteContext* context, void* buffer) {}
TfLiteStatus SoftmaxPrepare(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
const TfLiteTensor* input = GetInput(context, node, 0);
TF_LITE_ENSURE(context, NumDimensions(input) >= 1);
return kTfLiteOk;
}
// Takes a 1D tensor and performs softmax along it.
void Softmax1DFloat(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params) {
const int input_size = input->dims->data[0];
tflite::reference_ops::Softmax(input->data.f, input_size, 1, params->beta,
output->data.f);
}
// Takes a 2D tensor and perform softmax along the last dimension.
void Softmax2DFloat(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params) {
const int batch_size = input->dims->data[0];
const int input_size = input->dims->data[1];
tflite::reference_ops::Softmax(input->data.f, input_size, batch_size,
params->beta, output->data.f);
}
void Softmax1DQuantized(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params, OpData* data) {
const int input_size = input->dims->data[0];
const int32_t shape_data[4] = {1, 1, 1, input_size};
RuntimeShape shape(4, shape_data);
SoftmaxParams op_params;
op_params.input_multiplier = data->input_multiplier;
op_params.input_left_shift = data->input_left_shift;
op_params.diff_min = data->diff_min;
if (input->type == kTfLiteUInt8) {
tflite::reference_ops::Softmax(op_params, shape,
GetTensorData<uint8_t>(input), shape,
GetTensorData<uint8_t>(output));
} else {
arm_softmax_s8(GetTensorData<int8_t>(input), shape_data[0], shape_data[3],
op_params.input_multiplier, op_params.input_left_shift,
op_params.diff_min, GetTensorData<int8_t>(output));
}
}
void Softmax2DQuantized(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params, OpData* data) {
const int batch_size = input->dims->data[0];
const int input_size = input->dims->data[1];
const int32_t shape_data[4] = {batch_size, 1, 1, input_size};
RuntimeShape shape(4, shape_data);
SoftmaxParams op_params;
op_params.input_multiplier = data->input_multiplier;
op_params.input_left_shift = data->input_left_shift;
op_params.diff_min = data->diff_min;
if (input->type == kTfLiteUInt8) {
tflite::reference_ops::Softmax(op_params, shape,
GetTensorData<uint8_t>(input), shape,
GetTensorData<uint8_t>(output));
} else {
arm_softmax_s8(GetTensorData<int8_t>(input), shape_data[0], shape_data[3],
op_params.input_multiplier, op_params.input_left_shift,
op_params.diff_min, GetTensorData<int8_t>(output));
}
}
// Takes a 4D tensor and perform softmax along the forth dimension.
void Softmax4DFloat(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params) {
void SoftmaxFloat(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params) {
SoftmaxParams op_params;
op_params.beta = params->beta;
tflite::reference_ops::Softmax(
@ -152,13 +99,15 @@ void SoftmaxQuantized(const TfLiteTensor* input, TfLiteTensor* output,
} else {
const unsigned int num_dims = NumDimensions(input);
arm_softmax_s8(GetTensorData<int8_t>(input),
(num_dims == 4 ? input->dims->data[0] : 1) *
input->dims->data[num_dims - 3] *
input->dims->data[num_dims - 2],
input->dims->data[num_dims - 1], op_params.input_multiplier,
op_params.input_left_shift, op_params.diff_min,
GetTensorData<int8_t>(output));
const int trailing_dim = input_shape.DimensionsCount() - 1;
const int outer_size =
MatchingFlatSizeSkipDim(input_shape, trailing_dim, output_shape);
const int depth =
MatchingDim(input_shape, trailing_dim, output_shape, trailing_dim);
arm_softmax_s8(GetTensorData<int8_t>(input), outer_size, depth,
op_params.input_multiplier, op_params.input_left_shift,
op_params.diff_min, GetTensorData<int8_t>(output));
}
}
@ -175,47 +124,18 @@ TfLiteStatus SoftmaxEval(TfLiteContext* context, TfLiteNode* node) {
switch (input->type) {
case kTfLiteFloat32: {
if (NumDimensions(input) == 1) {
Softmax1DFloat(input, output, params);
return kTfLiteOk;
}
if (NumDimensions(input) == 2) {
Softmax2DFloat(input, output, params);
return kTfLiteOk;
}
if (NumDimensions(input) == 4) {
Softmax4DFloat(input, output, params);
return kTfLiteOk;
}
TF_LITE_KERNEL_LOG(
context, "Only 1D, 2D and 4D tensors supported currently, got %dD.",
NumDimensions(input));
return kTfLiteError;
SoftmaxFloat(input, output, params);
return kTfLiteOk;
}
case kTfLiteUInt8:
case kTfLiteInt8: {
if (NumDimensions(input) == 1) {
Softmax1DQuantized(input, output, params, data);
return kTfLiteOk;
}
if (NumDimensions(input) == 2) {
Softmax2DQuantized(input, output, params, data);
return kTfLiteOk;
}
if (NumDimensions(input) == 3 || NumDimensions(input) == 4) {
SoftmaxQuantized(input, output, params, data);
return kTfLiteOk;
}
TF_LITE_KERNEL_LOG(
context,
"Only 1D, 2D, 3D and 4D tensors supported currently, got %dD.",
NumDimensions(input));
return kTfLiteError;
SoftmaxQuantized(input, output, params, data);
return kTfLiteOk;
}
default:
TF_LITE_KERNEL_LOG(
context,
"Only float32, uint8_t and int8_t supported currently, got %d.",
"Only float32, uint8_t and int8_t input supported currently, got %d.",
input->type);
return kTfLiteError;
}

View File

@ -76,90 +76,17 @@ void* Init(TfLiteContext* context, const char* buffer, size_t length) {
void Free(TfLiteContext* context, void* buffer) {}
TfLiteStatus SoftmaxPrepare(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
const TfLiteTensor* input = GetInput(context, node, 0);
TF_LITE_ENSURE(context, NumDimensions(input) >= 1);
return kTfLiteOk;
}
// Takes a 1D tensor and performs softmax along it.
void Softmax1DFloat(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params) {
const int input_size = input->dims->data[0];
tflite::reference_ops::Softmax(input->data.f, input_size, 1, params->beta,
output->data.f);
}
// Takes a 2D tensor and perform softmax along the last dimension.
void Softmax2DFloat(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params) {
const int batch_size = input->dims->data[0];
const int input_size = input->dims->data[1];
tflite::reference_ops::Softmax(input->data.f, input_size, batch_size,
params->beta, output->data.f);
}
void Softmax1DQuantized(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params, OpData* data) {
// TODO(ahentz): this is arguably a dirty trick. Since the implementation
// always traverses the last dimension of a 4D tensor, we will pretend our 1D
// tensor is 4D in a special way. We will convert a (Y) shape into a (1,
// 1, 1, Y) shape.
const int input_size = input->dims->data[0];
const int32_t shape_data[4] = {1, 1, 1, input_size};
RuntimeShape shape(4, shape_data);
SoftmaxParams op_params;
op_params.input_multiplier = data->input_multiplier;
op_params.input_left_shift = data->input_left_shift;
op_params.diff_min = data->diff_min;
if (input->type == kTfLiteUInt8) {
tflite::reference_ops::Softmax(op_params, shape,
GetTensorData<uint8_t>(input), shape,
GetTensorData<uint8_t>(output));
} else {
if (output->type == kTfLiteInt16) {
tflite::reference_ops::Softmax(op_params, shape,
GetTensorData<int8_t>(input), shape,
GetTensorData<int16_t>(output));
} else {
tflite::reference_ops::Softmax(op_params, shape,
GetTensorData<int8_t>(input), shape,
GetTensorData<int8_t>(output));
}
}
}
void Softmax2DQuantized(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params, OpData* data) {
// TODO(ahentz): this is arguably a dirty trick. Since the implementation
// always traverses the last dimension of a 4D tensor, we will pretend our 2D
// tensor is 4D in a special way. We will convert a (X, Y) shape into a (X,
// 1, 1, Y) shape.
const int batch_size = input->dims->data[0];
const int input_size = input->dims->data[1];
const int32_t shape_data[4] = {batch_size, 1, 1, input_size};
RuntimeShape shape(4, shape_data);
SoftmaxParams op_params;
op_params.input_multiplier = data->input_multiplier;
op_params.input_left_shift = data->input_left_shift;
op_params.diff_min = data->diff_min;
if (input->type == kTfLiteUInt8) {
tflite::reference_ops::Softmax(op_params, shape,
GetTensorData<uint8_t>(input), shape,
GetTensorData<uint8_t>(output));
} else {
if (output->type == kTfLiteInt16) {
tflite::reference_ops::Softmax(op_params, shape,
GetTensorData<int8_t>(input), shape,
GetTensorData<int16_t>(output));
} else {
tflite::reference_ops::Softmax(op_params, shape,
GetTensorData<int8_t>(input), shape,
GetTensorData<int8_t>(output));
}
}
}
// Takes a 4D tensor and perform softmax along the forth dimension.
void Softmax4DFloat(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params) {
// Takes a tensor and performs softmax along the last dimension.
void SoftmaxFloat(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params) {
SoftmaxParams op_params;
op_params.beta = static_cast<double>(params->beta);
tflite::reference_ops::Softmax(
@ -201,51 +128,20 @@ TfLiteStatus SoftmaxEval(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_STATUS(
CalculateSoftmaxOpData(context, input, output, params, data));
// TODO(ahentz): consider an implementation that works for many (all?)
// dimensions.
switch (input->type) {
case kTfLiteFloat32: {
if (NumDimensions(input) == 1) {
Softmax1DFloat(input, output, params);
return kTfLiteOk;
}
if (NumDimensions(input) == 2) {
Softmax2DFloat(input, output, params);
return kTfLiteOk;
}
if (NumDimensions(input) == 4) {
Softmax4DFloat(input, output, params);
return kTfLiteOk;
}
TF_LITE_KERNEL_LOG(
context, "Only 1D, 2D and 4D tensors supported currently, got %dD.",
NumDimensions(input));
return kTfLiteError;
SoftmaxFloat(input, output, params);
return kTfLiteOk;
}
case kTfLiteInt8:
case kTfLiteUInt8: {
if (NumDimensions(input) == 1) {
Softmax1DQuantized(input, output, params, data);
return kTfLiteOk;
}
if (NumDimensions(input) == 2) {
Softmax2DQuantized(input, output, params, data);
return kTfLiteOk;
}
if (NumDimensions(input) == 3 || NumDimensions(input) == 4) {
SoftmaxQuantized(input, output, params, data);
return kTfLiteOk;
}
TF_LITE_KERNEL_LOG(
context,
"Only 1D, 2D, 3D and 4D tensors supported currently, got %dD.",
NumDimensions(input));
return kTfLiteError;
SoftmaxQuantized(input, output, params, data);
return kTfLiteOk;
}
default:
TF_LITE_KERNEL_LOG(
context,
"Only float32, uint8_t and int8_t supported currently, got %d.",
"Only float32, uint8_t and int8_t input supported currently, got %d.",
input->type);
return kTfLiteError;
}

View File

@ -160,24 +160,22 @@ TfLiteStatus CalculateSoftmaxOpData(TfLiteContext* context,
} // namespace
void Softmax2DQuantized(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params, OpData* data) {
const int batch_size = input->dims->data[0];
const int input_size = input->dims->data[1];
const int32_t shape_data[4] = {batch_size, 1, 1, input_size};
RuntimeShape shape(4, shape_data);
void SoftmaxQuantized(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params, OpData* data) {
SoftmaxParams op_params;
op_params.input_multiplier = data->input_multiplier;
op_params.input_left_shift = data->input_left_shift;
op_params.diff_min = data->diff_min;
if (output->type == kTfLiteInt16) {
xtensa::hifimini::Softmax(op_params, shape, GetTensorData<int8_t>(input),
shape, GetTensorData<int16_t>(output));
xtensa::hifimini::Softmax(
op_params, GetTensorShape(input), GetTensorData<int8_t>(input),
GetTensorShape(output), GetTensorData<int16_t>(output));
} else {
xtensa::hifimini::Softmax(op_params, shape, GetTensorData<int8_t>(input),
shape, GetTensorData<int8_t>(output));
xtensa::hifimini::Softmax(
op_params, GetTensorShape(input), GetTensorData<int8_t>(input),
GetTensorShape(output), GetTensorData<int8_t>(output));
}
}
@ -190,8 +188,11 @@ void Free(TfLiteContext* context, void* buffer) {}
TfLiteStatus SoftmaxPrepare(TfLiteContext* context, TfLiteNode* node) {
auto* params = reinterpret_cast<TfLiteSoftmaxParams*>(node->builtin_data);
TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
const TfLiteTensor* input = GetInput(context, node, 0);
TfLiteTensor* output = GetOutput(context, node, 0);
TF_LITE_ENSURE(context, NumDimensions(input) >= 1);
// TODO(b/132070898): Use statically slotted OpData structures until a
// scratch memory API is ready.
@ -213,17 +214,12 @@ TfLiteStatus SoftmaxEval(TfLiteContext* context, TfLiteNode* node) {
switch (input->type) {
case kTfLiteInt8: {
if (NumDimensions(input) == 2) {
Softmax2DQuantized(input, output, params, op_data);
return kTfLiteOk;
}
TF_LITE_KERNEL_LOG(context,
"Only 2D tensors supported currently, got %dD.",
NumDimensions(input));
return kTfLiteError;
SoftmaxQuantized(input, output, params, op_data);
return kTfLiteOk;
}
default:
TF_LITE_KERNEL_LOG(context, "Only int8_t supported currently, got %d.",
TF_LITE_KERNEL_LOG(context,
"Only int8_t input supported currently, got %d.",
input->type);
return kTfLiteError;
}