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