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