xref: /aosp_15_r20/external/tensorflow/tensorflow/lite/delegates/coreml/builders/convolution_op_builder.cc (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #include "tensorflow/lite/delegates/coreml/builders/convolution_op_builder.h"
16 
17 #include <algorithm>
18 #include <memory>
19 #include <string>
20 
21 #include "google/protobuf/repeated_field.h"
22 #include "tensorflow/lite/c/common.h"
23 #include "tensorflow/lite/delegates/coreml/builders/activation_layer_builder.h"
24 #include "tensorflow/lite/delegates/coreml/builders/op_factory.h"
25 #include "tensorflow/lite/delegates/coreml/builders/op_validator.h"
26 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
27 #include "tensorflow/lite/kernels/kernel_util.h"
28 
29 namespace tflite {
30 namespace delegates {
31 namespace coreml {
DebugName()32 const std::string& ConvolutionOpBuilder::DebugName() {
33   if (debug_name_.empty()) SetDebugName("ConvolutionOpBuilder", node_id_);
34   return debug_name_;
35 }
36 
SetWeights(TfLiteTensor * weights)37 void ConvolutionOpBuilder::SetWeights(TfLiteTensor* weights) {
38   weights_ = weights;
39 }
40 
SetBias(TfLiteTensor * bias)41 void ConvolutionOpBuilder::SetBias(TfLiteTensor* bias) { bias_ = bias; }
42 
SetOutputShape(TfLiteTensor * output_shape)43 void ConvolutionOpBuilder::SetOutputShape(TfLiteTensor* output_shape) {
44   output_shape_ = output_shape;
45 }
46 
Build()47 CoreML::Specification::NeuralNetworkLayer* ConvolutionOpBuilder::Build() {
48   if (layer_ == nullptr) {
49     layer_ = std::make_unique<CoreML::Specification::NeuralNetworkLayer>();
50   }
51   layer_->set_name(DebugName());
52 
53   int stride_height = 0;
54   int stride_width = 0;
55   int dilation_height = 0;
56   int dilation_width = 0;
57   TfLitePadding padding;
58 
59   switch (conv_type_) {
60     case ConvolutionType::kConv: {
61       const auto* conv_params =
62           reinterpret_cast<const TfLiteConvParams*>(builtin_data_);
63       stride_height = conv_params->stride_height;
64       stride_width = conv_params->stride_width;
65       dilation_height = conv_params->dilation_height_factor;
66       dilation_width = conv_params->dilation_width_factor;
67       padding = conv_params->padding;
68 
69       layer_->mutable_convolution()->set_ngroups(1);
70       break;
71     }
72     case ConvolutionType::kDepthwiseConv: {
73       const auto* depthwise_conv_params =
74           reinterpret_cast<const TfLiteDepthwiseConvParams*>(builtin_data_);
75       stride_height = depthwise_conv_params->stride_height;
76       stride_width = depthwise_conv_params->stride_width;
77       dilation_height = depthwise_conv_params->dilation_height_factor;
78       dilation_width = depthwise_conv_params->dilation_width_factor;
79       padding = depthwise_conv_params->padding;
80 
81       // n_groups = kernel_channel / depth_multiplier
82       layer_->mutable_convolution()->set_ngroups(
83           weights_->dims->data[3] / depthwise_conv_params->depth_multiplier);
84       break;
85     }
86     case ConvolutionType::kTransposeConv: {
87       const auto* transpose_conv_params =
88           reinterpret_cast<const TfLiteTransposeConvParams*>(builtin_data_);
89       const int height_index = 1;
90       const int width_index = 2;
91 
92       stride_height = transpose_conv_params->stride_height;
93       stride_width = transpose_conv_params->stride_width;
94       padding = transpose_conv_params->padding;
95 
96       layer_->mutable_convolution()->mutable_outputshape()->Add(
97           GetTensorData<int32_t>(output_shape_)[height_index]);
98       layer_->mutable_convolution()->mutable_outputshape()->Add(
99           GetTensorData<int32_t>(output_shape_)[width_index]);
100       break;
101     }
102   }
103 
104   // If not set, it will default to (1,1)
105   if (stride_height) {
106     layer_->mutable_convolution()->add_stride(stride_height);
107     layer_->mutable_convolution()->add_stride(stride_width);
108   }
109 
110   if (dilation_height) {
111     layer_->mutable_convolution()->add_dilationfactor(dilation_height);
112     layer_->mutable_convolution()->add_dilationfactor(dilation_width);
113   }
114 
115   switch (padding) {
116     case kTfLitePaddingSame:
117       layer_->mutable_convolution()->mutable_same();
118       break;
119     case kTfLitePaddingValid:
120       layer_->mutable_convolution()->mutable_valid();
121       break;
122     case kTfLitePaddingUnknown:
123       fprintf(stderr, "Padding is unknown.\n");
124       break;
125   }
126 
127   FillCoreMLWeights();
128   FillCoreMLBias();
129 
130   return layer_.release();
131 }
132 
FillCoreMLWeights()133 void ConvolutionOpBuilder::FillCoreMLWeights() {
134   if (conv_type_ == ConvolutionType::kDepthwiseConv) {
135     layer_->mutable_convolution()->set_kernelchannels(1);
136     layer_->mutable_convolution()->set_outputchannels(weights_->dims->data[3]);
137   } else {
138     layer_->mutable_convolution()->set_kernelchannels(weights_->dims->data[3]);
139     layer_->mutable_convolution()->set_outputchannels(weights_->dims->data[0]);
140   }
141   layer_->mutable_convolution()->add_kernelsize(weights_->dims->data[1]);
142   layer_->mutable_convolution()->add_kernelsize(weights_->dims->data[2]);
143 
144   TransposeKernelWeights();  // Should be called after CoreML shape is set.
145 }
146 
TransposeKernelWeights()147 void ConvolutionOpBuilder::TransposeKernelWeights() {
148   RuntimeShape tfl_shape(4, weights_->dims->data);
149   // CoreML kernel has shape of (C_out, C_in, H, W)
150   RuntimeShape coreml_shape(
151       {static_cast<int>(layer_->convolution().outputchannels()),
152        static_cast<int>(layer_->convolution().kernelchannels()),
153        static_cast<int>(layer_->convolution().kernelsize()[0]),
154        static_cast<int>(layer_->convolution().kernelsize()[1])});
155 
156   TransposeParams params;
157 
158   if (conv_type_ == ConvolutionType::kDepthwiseConv) {
159     // DepthwiseConv2D: TFL kernel has shape of (1, H, W, C_out),
160     // and CoreML kernel has shape of (C_out, 1, H, W)
161     params = {/*perm_count=*/4, /*perm=*/{3, 0, 1, 2}};
162   } else {
163     // Conv2D and TransposeConv: TFL kernel has shape of (C_out, H, W, C_in),
164     // and CoreML kernel has shape of (C_out, C_in, H, W)
165     params = {/*perm_count=*/4, /*perm=*/{0, 3, 1, 2}};
166   }
167 
168   if (conv_type_ == ConvolutionType::kTransposeConv) {
169     layer_->mutable_convolution()->set_isdeconvolution(true);
170   }
171 
172   if (weights_->type == kTfLiteFloat32) {
173     auto* coreml_weights =
174         layer_->mutable_convolution()->mutable_weights()->mutable_floatvalue();
175     coreml_weights->Resize(NumElements(weights_), 0);
176 
177     optimized_ops::Transpose<float>(params, tfl_shape, weights_->data.f,
178                                     coreml_shape,
179                                     coreml_weights->mutable_data());
180   } else if (weights_->type == kTfLiteFloat16) {
181     auto* coreml_weights = layer_->mutable_convolution()
182                                ->mutable_weights()
183                                ->mutable_float16value();
184     // float16value has type of bytes (std::string)
185     coreml_weights->resize(weights_->bytes, 0);
186 
187     optimized_ops::Transpose<uint16_t>(
188         params, tfl_shape, reinterpret_cast<uint16_t*>(weights_->data.raw),
189         coreml_shape, reinterpret_cast<uint16_t*>(&coreml_weights->front()));
190   }
191 }
192 
FillCoreMLBias()193 void ConvolutionOpBuilder::FillCoreMLBias() {
194   if (bias_ != nullptr) {
195     layer_->mutable_convolution()->set_hasbias(true);
196     if (bias_->type == kTfLiteFloat32) {
197       std::copy(bias_->data.f, bias_->data.f + NumElements(bias_->dims),
198                 google::protobuf::RepeatedFieldBackInserter(layer_->mutable_convolution()
199                                                       ->mutable_bias()
200                                                       ->mutable_floatvalue()));
201     } else if (bias_->type == kTfLiteFloat16) {
202       // float16value has type of bytes (std::string)
203       layer_->mutable_convolution()
204           ->mutable_bias()
205           ->mutable_float16value()
206           ->assign(bias_->data.raw, bias_->bytes);
207     }
208   }
209 }
210 
PopulateSubgraph(TfLiteContext * context)211 TfLiteStatus ConvolutionOpBuilder::PopulateSubgraph(TfLiteContext* context) {
212   TfLiteFusedActivation activation;
213   switch (conv_type_) {
214     case ConvolutionType::kConv: {
215       const auto* conv_params =
216           reinterpret_cast<const TfLiteConvParams*>(builtin_data_);
217       activation = conv_params->activation;
218       break;
219     }
220     case ConvolutionType::kDepthwiseConv: {
221       const auto* depthwise_conv_params =
222           reinterpret_cast<const TfLiteDepthwiseConvParams*>(builtin_data_);
223       activation = depthwise_conv_params->activation;
224       break;
225     }
226     case ConvolutionType::kTransposeConv: {
227       activation = kTfLiteActNone;
228       break;
229     }
230   }
231 
232   if (activation == kTfLiteActNone) {
233     builder_output_ = AddOutput();
234   } else {
235     ActivationLayerBuilder* activation_builder =
236         reinterpret_cast<ActivationLayerBuilder*>(
237             graph_builder_->AddBuilder(CreateActivationLayerBuilder, nullptr));
238     activation_builder->SetActivation(activation);
239     activation_builder->AddInput(AddOutput());
240     activation_builder->PopulateSubgraph(context);
241     builder_output_ = activation_builder->GetOutput(context);
242   }
243   return kTfLiteOk;
244 }
245 
RegisterInputs(const TfLiteIntArray * inputs,TfLiteContext * context)246 TfLiteStatus ConvolutionOpBuilder::RegisterInputs(const TfLiteIntArray* inputs,
247                                                   TfLiteContext* context) {
248   if (conv_type_ == ConvolutionType::kTransposeConv) {
249     if (inputs->size != 3) {
250       TF_LITE_KERNEL_LOG(context,
251                          "Transpose Conv should have 3 inputs, %d given.",
252                          inputs->size);
253       return kTfLiteError;
254     }
255     AddInput(inputs->data[2]);
256     SetOutputShape(&context->tensors[inputs->data[0]]);
257   } else {
258     if (inputs->size != 2 && inputs->size != 3) {
259       TF_LITE_KERNEL_LOG(context,
260                          "Convolution and depthwise convolution should have 2 "
261                          "or 3 inputs, %d given.",
262                          inputs->size);
263       return kTfLiteError;
264     }
265     AddInput(inputs->data[0]);
266     if (inputs->size > 2) {
267       SetBias(&context->tensors[inputs->data[2]]);
268     }
269   }
270   SetWeights(&context->tensors[inputs->data[1]]);
271   return kTfLiteOk;
272 }
273 
RegisterOutputs(const TfLiteIntArray * outputs,TfLiteContext * context)274 TfLiteStatus ConvolutionOpBuilder::RegisterOutputs(
275     const TfLiteIntArray* outputs, TfLiteContext* context) {
276   if (outputs->size != 1) {
277     TF_LITE_KERNEL_LOG(context, "Wrong # of outputs!.");
278     return kTfLiteError;
279   }
280   TensorID output_tensor = GetOutput(context);
281   if (output_tensor.NodeID() == -1) {
282     TF_LITE_KERNEL_LOG(context, "Failed to build output tensor.");
283     return kTfLiteError;
284   }
285   graph_builder_->AddTensorWithID(outputs->data[0], output_tensor);
286   return kTfLiteOk;
287 }
288 
CreateConvolutionOpBuilder(GraphBuilder * graph_builder)289 OpBuilder* CreateConvolutionOpBuilder(GraphBuilder* graph_builder) {
290   return new ConvolutionOpBuilder(graph_builder, ConvolutionType::kConv);
291 }
292 
CreateDepthwiseConvolutionOpBuilder(GraphBuilder * graph_builder)293 OpBuilder* CreateDepthwiseConvolutionOpBuilder(GraphBuilder* graph_builder) {
294   return new ConvolutionOpBuilder(graph_builder,
295                                   ConvolutionType::kDepthwiseConv);
296 }
297 
CreateTransposeConvolutionOpBuilder(GraphBuilder * graph_builder)298 OpBuilder* CreateTransposeConvolutionOpBuilder(GraphBuilder* graph_builder) {
299   return new ConvolutionOpBuilder(graph_builder,
300                                   ConvolutionType::kTransposeConv);
301 }
302 
IsConvolutionOpSupported(const TfLiteRegistration * registration,const TfLiteNode * node,TfLiteContext * context)303 bool IsConvolutionOpSupported(const TfLiteRegistration* registration,
304                               const TfLiteNode* node, TfLiteContext* context) {
305   if (node->builtin_data == nullptr) return false;
306 
307   TfLiteFusedActivation activation;
308 
309   if (registration->builtin_code == kTfLiteBuiltinConv2d) {
310     const auto* conv_params =
311         reinterpret_cast<const TfLiteConvParams*>(node->builtin_data);
312     activation = conv_params->activation;
313   } else if (registration->builtin_code == kTfLiteBuiltinDepthwiseConv2d) {
314     const auto* depthwise_conv_params =
315         reinterpret_cast<const TfLiteDepthwiseConvParams*>(node->builtin_data);
316     activation = depthwise_conv_params->activation;
317   } else if (registration->builtin_code == kTfLiteBuiltinTransposeConv) {
318     activation = kTfLiteActNone;
319   } else {
320     TF_LITE_KERNEL_LOG(
321         context,
322         "Invalid op: op must be Conv2D, DepthwiseConv2D or TransposeConv.");
323     return false;
324   }
325 
326   if (activation == kTfLiteActSignBit) {
327     return false;
328   }
329 
330   const int kOutputShapeTensor = 0;  // Only used for TransposeConv
331   const int kWeightTensor = 1;
332   const int kBiasTensor = 2;  // Only used for non-TransposeConv
333   const TfLiteTensor* weights;
334   TF_LITE_ENSURE_OK(context,
335                     GetInputSafe(context, node, kWeightTensor, &weights));
336   const int max_kernel_size = 16384;
337   if (!IsConstantTensor(weights)) {
338     return false;
339   }
340   if (weights->dims->data[1] > max_kernel_size ||
341       weights->dims->data[2] > max_kernel_size) {
342     return false;
343   }
344   if (registration->builtin_code == kTfLiteBuiltinTransposeConv) {
345     if (!IsConstantTensor(GetInput(context, node, kOutputShapeTensor))) {
346       return false;
347     }
348   } else {
349     if (node->inputs->size >= kBiasTensor &&
350         !IsConstantTensor(GetInput(context, node, kBiasTensor))) {
351       return false;
352     }
353   }
354 
355   return true;
356 }
357 
IsDepthwiseConvolutionOpSupported(const TfLiteRegistration * registration,const TfLiteNode * node,TfLiteContext * context)358 bool IsDepthwiseConvolutionOpSupported(const TfLiteRegistration* registration,
359                                        const TfLiteNode* node,
360                                        TfLiteContext* context) {
361   return IsConvolutionOpSupported(registration, node, context);
362 }
363 
IsTransposeConvolutionOpSupported(const TfLiteRegistration * registration,const TfLiteNode * node,TfLiteContext * context)364 bool IsTransposeConvolutionOpSupported(const TfLiteRegistration* registration,
365                                        const TfLiteNode* node,
366                                        TfLiteContext* context) {
367   return IsConvolutionOpSupported(registration, node, context);
368 }
369 
370 }  // namespace coreml
371 }  // namespace delegates
372 }  // namespace tflite
373