Add transpose conv support to the NNAPI delegate
PiperOrigin-RevId: 260909964
This commit is contained in:
Lev Proleev
TensorFlower Gardener
parent
b32786df61
commit
4902e6bf68
@ -557,7 +557,14 @@ class NNAPIOpBuilder {
|
||||
TfLiteStatus AddVectorInt32Operand(const int32_t* values,
|
||||
uint32_t num_values) {
|
||||
return AddVectorOperand<int32_t>(values, num_values,
|
||||
ANEURALNETWORKS_TENSOR_INT32);
|
||||
ANEURALNETWORKS_TENSOR_INT32,
|
||||
/*scale=*/0.f, /*zero_point=*/0);
|
||||
}
|
||||
|
||||
TfLiteStatus AddVectorInt32Operand(const int32_t* values, uint32_t num_values,
|
||||
float scale, int32_t zero_point) {
|
||||
return AddVectorOperand<int32_t>(
|
||||
values, num_values, ANEURALNETWORKS_TENSOR_INT32, scale, zero_point);
|
||||
}
|
||||
|
||||
TfLiteStatus AddVectorFloat32Operand(const float* values,
|
||||
@ -781,9 +788,13 @@ class NNAPIOpBuilder {
|
||||
|
||||
template <typename T>
|
||||
TfLiteStatus AddVectorOperand(const T* values, uint32_t num_values,
|
||||
int32_t nn_type) {
|
||||
ANeuralNetworksOperandType operand_type{
|
||||
.type = nn_type, .dimensionCount = 1, .dimensions = &num_values};
|
||||
int32_t nn_type, float scale,
|
||||
int32_t zero_point) {
|
||||
ANeuralNetworksOperandType operand_type{.type = nn_type,
|
||||
.dimensionCount = 1,
|
||||
.dimensions = &num_values,
|
||||
.scale = scale,
|
||||
.zeroPoint = zero_point};
|
||||
|
||||
RETURN_TFLITE_ERROR_IF_NN_ERROR(
|
||||
context_,
|
||||
@ -797,6 +808,13 @@ class NNAPIOpBuilder {
|
||||
return kTfLiteOk;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
TfLiteStatus AddVectorOperand(const T* values, uint32_t num_values,
|
||||
int32_t nn_type) {
|
||||
return AddVectorOperand(values, num_values, nn_type, /*scale=*/0.f,
|
||||
/*zero_point=*/0);
|
||||
}
|
||||
|
||||
TfLiteStatus AddFloat32OutputTensor(uint32_t dimension_count,
|
||||
const uint32_t* dimension_data,
|
||||
int* ann_index_out) {
|
||||
@ -1812,6 +1830,91 @@ class NNAPIDelegateKernel {
|
||||
return BasicMappingFn<ANEURALNETWORKS_SIN>;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinTransposeConv:
|
||||
if (version == 1 && android_sdk_version >= kMinSdkVersionForNNAPI12) {
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
const bool hybrid_op = IsHybridOperator(mapping_args.context,
|
||||
kTfLiteBuiltinTransposeConv,
|
||||
mapping_args.node);
|
||||
mapping_args.builder->AddTensorInput(/*kDataInputTensor*/ 2,
|
||||
hybrid_op);
|
||||
mapping_args.builder->AddTensorInput(/*kWeightsTensor*/ 1,
|
||||
hybrid_op);
|
||||
|
||||
// NNAPI requires a bias tensor, so we allocate a new tensor to fill
|
||||
// it with zeroes. It is deleted with other tensors in the context
|
||||
// during subgraph destructor call.
|
||||
int bias_index = -1;
|
||||
mapping_args.context->AddTensors(mapping_args.context, 1,
|
||||
&bias_index);
|
||||
TfLiteTensor* bias_tensor =
|
||||
&mapping_args.context->tensors[bias_index];
|
||||
const auto input_type =
|
||||
mapping_args.context
|
||||
->tensors[mapping_args.node->inputs
|
||||
->data[/*kDataInputTensor*/ 2]]
|
||||
.type;
|
||||
if (input_type == kTfLiteFloat32) {
|
||||
bias_tensor->type = kTfLiteFloat32;
|
||||
} else {
|
||||
bias_tensor->type = kTfLiteInt32;
|
||||
}
|
||||
|
||||
// Create an array with a required bias shape and resize the bias
|
||||
// tensor.
|
||||
TfLiteIntArray* bias_shape = TfLiteIntArrayCreate(1);
|
||||
const TfLiteTensor& output_shape =
|
||||
mapping_args.context->tensors
|
||||
[mapping_args.node->inputs->data[/*kOutputShapeTensor*/ 0]];
|
||||
const int output_depth = output_shape.data.i32[3];
|
||||
bias_shape->data[0] = output_depth;
|
||||
bias_tensor->allocation_type = kTfLiteDynamic;
|
||||
mapping_args.context->ResizeTensor(mapping_args.context,
|
||||
bias_tensor, bias_shape);
|
||||
|
||||
// Set tensor's values to zeroes and add it using AddVector*, so
|
||||
// that the values are copied to NNAPI. We don't use the AddTensor
|
||||
// function because it doesn't copy values and the tensor we just
|
||||
// created is not in the node->inputs.
|
||||
if (input_type == kTfLiteFloat32) {
|
||||
memset(bias_tensor->data.f, 0, output_depth * sizeof(float));
|
||||
mapping_args.builder->AddVectorFloat32Operand(bias_tensor->data.f,
|
||||
output_depth);
|
||||
} else {
|
||||
memset(bias_tensor->data.i32, 0, output_depth * sizeof(int));
|
||||
const TfLiteTensor& input_tensor =
|
||||
mapping_args.context->tensors
|
||||
[mapping_args.node->inputs->data[/*kDataInputTensor*/ 2]];
|
||||
const TfLiteTensor& filter_tensor =
|
||||
mapping_args.context->tensors
|
||||
[mapping_args.node->inputs->data[/*kWeightsTensor*/ 1]];
|
||||
// NNAPI requires bias scale to be a product of an input scale and
|
||||
// a filter scale.
|
||||
bias_tensor->params.scale =
|
||||
input_tensor.params.scale * filter_tensor.params.scale;
|
||||
mapping_args.builder->AddVectorInt32Operand(
|
||||
bias_tensor->data.i32, output_depth,
|
||||
input_tensor.params.scale * filter_tensor.params.scale,
|
||||
/*zero_point=*/0);
|
||||
}
|
||||
|
||||
mapping_args.builder->AddTensorInput(/*kOutputShapeTensor*/ 0,
|
||||
hybrid_op);
|
||||
|
||||
auto builtin = reinterpret_cast<TfLiteTransposeConvParams*>(
|
||||
mapping_args.node->builtin_data);
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->padding);
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->stride_width);
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->stride_height);
|
||||
mapping_args.builder->AddScalarInt32Operand(
|
||||
/*ANEURALNETWORKS_FUSED_NONE*/ 0);
|
||||
// Use NHWC layout for input and output
|
||||
mapping_args.builder->AddScalarBoolOperand(false);
|
||||
return ANEURALNETWORKS_TRANSPOSE_CONV;
|
||||
};
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinSqrt:
|
||||
if (version == 1 && android_sdk_version >= kMinSdkVersionForNNAPI12 &&
|
||||
IsFloat(context->tensors[node->inputs->data[0]].type)) {
|
||||
@ -3087,13 +3190,17 @@ class NNAPIDelegateKernel {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
if ((reg->builtin_code == kTfLiteBuiltinSplit) &&
|
||||
(input_index == node->inputs->data[0])) {
|
||||
// Skip the axis input tensor; it will be added as a scalar operand
|
||||
// by the Map() mapping.
|
||||
continue;
|
||||
}
|
||||
if (reg->builtin_code == kTfLiteBuiltinTransposeConv) {
|
||||
// Everything is added during Map since input tensors
|
||||
// have different order.
|
||||
continue;
|
||||
}
|
||||
|
||||
// Pad and Padv2 have an optional parameter for a pad value which has
|
||||
// to be converted to a scalar type in NN API.
|
||||
|
||||
@ -1460,6 +1460,7 @@ cc_test(
|
||||
name = "transpose_conv_test",
|
||||
size = "small",
|
||||
srcs = ["transpose_conv_test.cc"],
|
||||
tags = ["tflite_nnapi"],
|
||||
deps = [
|
||||
":builtin_ops",
|
||||
":test_main",
|
||||
|
||||
@ -130,6 +130,7 @@ enum {
|
||||
ANEURALNETWORKS_SQRT = 88,
|
||||
ANEURALNETWORKS_TILE = 89,
|
||||
ANEURALNETWORKS_TOPK_V2 = 90,
|
||||
ANEURALNETWORKS_TRANSPOSE_CONV = 91,
|
||||
ANEURALNETWORKS_UNIDIRECTIONAL_SEQUENCE_LSTM = 92,
|
||||
ANEURALNETWORKS_UNIDIRECTIONAL_SEQUENCE_RNN = 93,
|
||||
ANEURALNETWORKS_RESIZE_NEAREST_NEIGHBOR = 94,
|
||||
|
||||
Loading…
Reference in New Issue
Block a user