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/toco/tflite/export.h"
16
17 #include <string>
18
19 #include "flatbuffers/flexbuffers.h"
20 #include "absl/strings/str_join.h"
21 #include "tensorflow/core/lib/core/errors.h"
22 #include "tensorflow/lite/context.h"
23 #include "tensorflow/lite/schema/schema_conversion_utils.h"
24 #include "tensorflow/lite/schema/schema_generated.h"
25 #include "tensorflow/lite/toco/tflite/operator.h"
26 #include "tensorflow/lite/toco/tflite/types.h"
27 #include "tensorflow/lite/toco/tooling_util.h"
28 #include "tensorflow/lite/tools/optimize/quantize_weights.h"
29 #include "tensorflow/lite/tools/versioning/runtime_version.h"
30 #include "tensorflow/lite/version.h"
31
32 namespace toco {
33
34 namespace tflite {
35
36 using flatbuffers::FlatBufferBuilder;
37 using flatbuffers::Offset;
38 using flatbuffers::Vector;
39 using ::tflite::Buffer;
40 using ::tflite::BuiltinOperator;
41 using ::tflite::BuiltinOperator_CUSTOM;
42 using ::tflite::BuiltinOperator_MAX;
43 using ::tflite::BuiltinOperator_MIN;
44 using ::tflite::CreateBuffer;
45 using ::tflite::CreateMetadata;
46 using ::tflite::CreateModel;
47 using ::tflite::CreateOperator;
48 using ::tflite::CreateTensor;
49 using ::tflite::Metadata;
50 using ::tflite::Operator;
51 using ::tflite::OperatorCode;
52 using ::tflite::SubGraph;
53 using ::tflite::Tensor;
54
55 namespace {
56
57 // Check if a TensorFlow Op is a control flow op by its name.
IsControlFlowOp(const std::string & tensorflow_op)58 bool IsControlFlowOp(const std::string& tensorflow_op) {
59 // Technically this is equivalent to `::tensorflow::Node::IsControlFlow()`.
60 // It requires to construct a `::tensorflow::Graph` to use that helper
61 // function, so we simply hardcode the list of control flow ops here.
62 if (tensorflow_op == "Switch" || tensorflow_op == "RefSwitch" ||
63 tensorflow_op == "Merge" || tensorflow_op == "RefMerge" ||
64 tensorflow_op == "Enter" || tensorflow_op == "RefEnter" ||
65 tensorflow_op == "Exit" || tensorflow_op == "RefExit" ||
66 tensorflow_op == "NextIteration" || tensorflow_op == "RefNextIteration") {
67 return true;
68 }
69 // TODO(ycling): Also check how to handle Variable ops and Assign ops.
70 return false;
71 }
72
73 // Check if a TensorFlow Op is unsupported by the Flex runtime.
IsUnsupportedFlexOp(const std::string & tensorflow_op)74 bool IsUnsupportedFlexOp(const std::string& tensorflow_op) {
75 if (IsControlFlowOp(tensorflow_op)) {
76 return true;
77 }
78 // `HashTableV2` isn't supported for now since it requires an additional
79 // initialization step.
80 // TODO(b/117651199): Support `HashTableV2` with Flex runtime.
81 if (tensorflow_op == "HashTableV2") {
82 return true;
83 }
84 return false;
85 }
86
87 // Map from operator name to TF Lite enum value, for all builtins.
GetBuiltinOpsMap()88 const std::map<std::string, BuiltinOperator>& GetBuiltinOpsMap() {
89 static std::map<std::string, BuiltinOperator>* builtin_ops = nullptr;
90 if (builtin_ops == nullptr) {
91 builtin_ops = new std::map<std::string, BuiltinOperator>();
92
93 for (int i = BuiltinOperator_MIN; i <= BuiltinOperator_MAX; ++i) {
94 BuiltinOperator op = static_cast<BuiltinOperator>(i);
95 std::string name = EnumNameBuiltinOperator(op);
96 if (op != BuiltinOperator_CUSTOM && !name.empty()) {
97 (*builtin_ops)[name] = op;
98 }
99 }
100 }
101 return *builtin_ops;
102 }
103
WriteModelToString(const flatbuffers::FlatBufferBuilder & builder,std::string * file_contents)104 void WriteModelToString(const flatbuffers::FlatBufferBuilder& builder,
105 std::string* file_contents) {
106 const uint8_t* buffer = builder.GetBufferPointer();
107 int size = builder.GetSize();
108 *file_contents = std::string(reinterpret_cast<const char*>(buffer), size);
109 }
110
111 } // Anonymous namespace.
112
113 namespace details {
114
OperatorKey(const::toco::OperatorSignature & op_signature,const std::map<OperatorType,std::unique_ptr<BaseOperator>> & ops_by_type,bool enable_select_tf_ops)115 OperatorKey::OperatorKey(
116 const ::toco::OperatorSignature& op_signature,
117 const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type,
118 bool enable_select_tf_ops) {
119 // Get the op name (by Toco definition).
120 const ::toco::Operator& op = *op_signature.op;
121 std::string name = HelpfulOperatorTypeName(op);
122
123 bool is_builtin = false;
124 const auto& builtin_ops = GetBuiltinOpsMap();
125 if (ops_by_type.count(op.type) != 0) {
126 version_ = ops_by_type.at(op.type)->GetVersion(op_signature);
127 name = ops_by_type.at(op.type)->name();
128 is_builtin = (builtin_ops.count(name) > 0);
129 }
130
131 if (is_builtin) {
132 // For TFLite supported builtin ops, find out its BuiltinOperator enum used
133 // in FlatBuffer.
134 type_ = builtin_ops.at(name);
135 return;
136 }
137 // The logic below is all for custom ops or Flex ops.
138 is_custom_op_ = true;
139 type_ = BuiltinOperator_CUSTOM;
140
141 if (op.type == OperatorType::kUnsupported) {
142 const TensorFlowUnsupportedOperator& unsupported_op =
143 static_cast<const TensorFlowUnsupportedOperator&>(op);
144 const auto tensorflow_op = unsupported_op.tensorflow_op;
145
146 if (ShouldExportAsFlexOp(enable_select_tf_ops,
147 unsupported_op.tensorflow_op)) {
148 is_custom_op_ = false;
149 is_flex_op_ = true;
150 flex_tensorflow_op_ = tensorflow_op;
151 custom_code_ =
152 std::string(::tflite::kFlexCustomCodePrefix) + flex_tensorflow_op_;
153 } else {
154 custom_code_ = tensorflow_op;
155 }
156 } else if (enable_select_tf_ops && !op.tensorflow_node_def.empty()) {
157 // For Toco-supported/TFLite-unsupported ops, if the TensorFlow NodeDef
158 // is retained in the Toco Operator, we produce a Flex op if Flex mode
159 // is enabled.
160 is_custom_op_ = false;
161 is_flex_op_ = true;
162 flex_tensorflow_op_ = name;
163 custom_code_ =
164 std::string(::tflite::kFlexCustomCodePrefix) + flex_tensorflow_op_;
165 } else {
166 // If Flex is disabled or the original TensorFlow NodeDef isn't available,
167 // we produce a custom op. This gives developers a chance to implement
168 // custom ops.
169 custom_code_ = name;
170 }
171
172 if (is_flex_op_) {
173 if (IsUnsupportedFlexOp(flex_tensorflow_op_)) {
174 is_unsupported_flex_op_ = true;
175 }
176 }
177 }
178
LoadTensorsMap(const Model & model,TensorsMap * tensors_map)179 void LoadTensorsMap(const Model& model, TensorsMap* tensors_map) {
180 // First find a list of unique array names.
181 std::set<std::string> names;
182 for (const auto& array_pair : model.GetArrayMap()) {
183 names.insert(array_pair.first);
184 }
185
186 // Now assign indices to them and fill in the map.
187 int index = 0;
188 for (const auto& name : names) {
189 (*tensors_map)[name] = index;
190 ++index;
191 }
192 }
193
LoadOperatorsMap(const Model & model,OperatorsMap * operators_map,const std::map<OperatorType,std::unique_ptr<BaseOperator>> & ops_by_type,bool enable_select_tf_ops)194 void LoadOperatorsMap(
195 const Model& model, OperatorsMap* operators_map,
196 const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type,
197 bool enable_select_tf_ops) {
198 // First find a list of unique operator types.
199 std::set<OperatorKey> keys;
200 for (const auto& op : model.operators) {
201 const toco::OperatorSignature op_signature = {op.get(), &model};
202 keys.insert(OperatorKey(op_signature, ops_by_type, enable_select_tf_ops));
203 }
204 // Now assign indices to them and fill in the map.
205 int index = 0;
206 for (const auto& key : keys) {
207 (*operators_map)[key] = index;
208 ++index;
209 }
210 }
211
212 } // namespace details
213
ExportTensors(const Model & model,const details::TensorsMap & tensors_map,FlatBufferBuilder * builder,std::vector<Offset<Vector<uint8_t>>> * buffers_to_write,const std::set<int32_t> & variable_tensor_indices)214 Offset<Vector<Offset<Tensor>>> ExportTensors(
215 const Model& model, const details::TensorsMap& tensors_map,
216 FlatBufferBuilder* builder,
217 std::vector<Offset<Vector<uint8_t>>>* buffers_to_write,
218 const std::set<int32_t>& variable_tensor_indices) {
219 // In the end we will need to produce a vector sorted by the indices of the
220 // tensors in the tensors_map.
221 std::map<int, Offset<Tensor>> ordered_tensors;
222
223 for (const auto& array_pair : model.GetArrayMap()) {
224 const std::string& tensor_name = array_pair.first;
225 const toco::Array& array = *array_pair.second;
226
227 int buffer_index = buffers_to_write->size();
228 auto type = DataType::Serialize(array.data_type);
229 buffers_to_write->push_back(DataBuffer::Serialize(array, builder));
230
231 std::vector<int> shape;
232 if (array.has_shape()) {
233 shape.reserve(array.shape().dims().size());
234 for (const auto& d : array.shape().dims()) {
235 shape.push_back(d);
236 }
237 }
238
239 Offset<Vector<float>> min;
240 Offset<Vector<float>> max;
241 Offset<Vector<float>> scale;
242 Offset<Vector<int64_t>> zero_point;
243 if (array.minmax) {
244 min = builder->CreateVector(
245 std::vector<float>{static_cast<float>(array.minmax->min)});
246 max = builder->CreateVector(
247 std::vector<float>{static_cast<float>(array.minmax->max)});
248 }
249 if (array.quantization_params) {
250 scale = builder->CreateVector(std::vector<float>{
251 static_cast<float>(array.quantization_params->scale)});
252 zero_point = builder->CreateVector(
253 std::vector<int64_t>{array.quantization_params->zero_point});
254 }
255 auto q_param = ::tflite::CreateQuantizationParameters(*builder, min, max,
256 scale, zero_point);
257
258 int index = tensors_map.at(tensor_name);
259 bool is_variable =
260 variable_tensor_indices.find(index) != variable_tensor_indices.end();
261 ordered_tensors[index] =
262 CreateTensor(*builder, builder->CreateVector(shape), type, buffer_index,
263 builder->CreateString(tensor_name), q_param, is_variable);
264 }
265
266 std::vector<Offset<Tensor>> tensor_vector;
267 tensor_vector.reserve(ordered_tensors.size());
268 for (const auto& tensor : ordered_tensors) {
269 tensor_vector.push_back(tensor.second);
270 }
271
272 return builder->CreateVector(tensor_vector);
273 }
274
ExportInputTensors(const Model & model,const details::TensorsMap & tensors_map,FlatBufferBuilder * builder)275 Offset<Vector<int32_t>> ExportInputTensors(
276 const Model& model, const details::TensorsMap& tensors_map,
277 FlatBufferBuilder* builder) {
278 std::vector<int32_t> inputs;
279 for (const auto& input : model.flags.input_arrays()) {
280 inputs.push_back(tensors_map.at(input.name()));
281 }
282 return builder->CreateVector<int32_t>(inputs);
283 }
284
ExportOutputTensors(const Model & model,const details::TensorsMap & tensors_map,FlatBufferBuilder * builder)285 Offset<Vector<int32_t>> ExportOutputTensors(
286 const Model& model, const details::TensorsMap& tensors_map,
287 FlatBufferBuilder* builder) {
288 std::vector<int32_t> outputs;
289 for (const std::string& output : model.flags.output_arrays()) {
290 outputs.push_back(tensors_map.at(output));
291 }
292 return builder->CreateVector<int32_t>(outputs);
293 }
294
ExportOperatorCodes(const Model & model,const std::map<OperatorType,std::unique_ptr<BaseOperator>> & ops_by_type,const details::OperatorsMap & operators_map,FlatBufferBuilder * builder,const ExportParams & params)295 Offset<Vector<Offset<OperatorCode>>> ExportOperatorCodes(
296 const Model& model,
297 const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type,
298 const details::OperatorsMap& operators_map, FlatBufferBuilder* builder,
299 const ExportParams& params) {
300 // Map from operator name to TF Lite enum value, for all builtins.
301 std::map<std::string, BuiltinOperator> builtin_ops;
302 for (int i = BuiltinOperator_MIN; i <= BuiltinOperator_MAX; ++i) {
303 BuiltinOperator op = static_cast<BuiltinOperator>(i);
304 std::string name = EnumNameBuiltinOperator(op);
305 if (op != BuiltinOperator_CUSTOM && !name.empty()) {
306 builtin_ops[name] = op;
307 }
308 }
309
310 // We will need to produce a vector of codes in the same order as they
311 // appear in the operators_map.
312 std::map<int, Offset<OperatorCode>> ordered_opcodes;
313
314 for (const auto& op : model.operators) {
315 const toco::OperatorSignature op_signature = {op.get(), &model};
316 const details::OperatorKey operator_key = details::OperatorKey(
317 op_signature, ops_by_type, params.enable_select_tf_ops);
318 int op_index = operators_map.at(operator_key);
319
320 flatbuffers::Offset<flatbuffers::String> custom_code = 0;
321 if (!operator_key.custom_code().empty()) {
322 custom_code = builder->CreateString(operator_key.custom_code());
323 }
324
325 ordered_opcodes[op_index] = CreateOperatorCode(
326 *builder, operator_key.type(), custom_code, operator_key.version());
327 }
328
329 std::vector<Offset<OperatorCode>> opcode_vector;
330 opcode_vector.reserve(ordered_opcodes.size());
331 for (const auto& opcode : ordered_opcodes) {
332 opcode_vector.push_back(opcode.second);
333 }
334
335 return builder->CreateVector(opcode_vector);
336 }
337
ExportOperators(const Model & model,const std::map<OperatorType,std::unique_ptr<BaseOperator>> & ops_by_type,const details::OperatorsMap & operators_map,const details::TensorsMap & tensors_map,FlatBufferBuilder * builder,std::set<int32_t> * variable_tensor_indices,const ExportParams & params)338 Offset<Vector<Offset<Operator>>> ExportOperators(
339 const Model& model,
340 const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type,
341 const details::OperatorsMap& operators_map,
342 const details::TensorsMap& tensors_map, FlatBufferBuilder* builder,
343 std::set<int32_t>* variable_tensor_indices, const ExportParams& params) {
344 variable_tensor_indices->clear();
345
346 auto is_tflite_builtin = [](const BaseOperator* op) {
347 const auto& tflite_builtins = GetBuiltinOpsMap();
348 return (op && tflite_builtins.find(op->name()) != tflite_builtins.end());
349 };
350
351 // The operators are in execution order, so we just follow tf.mini order.
352 std::vector<Offset<Operator>> op_vector;
353 for (const auto& op : model.operators) {
354 std::vector<int32_t> inputs;
355 for (const std::string& input : op->inputs) {
356 // -1 is the ID for optional tensor in TFLite output
357 int id = model.IsOptionalArray(input) ? -1 : tensors_map.at(input);
358 inputs.push_back(id);
359 }
360 std::vector<int32_t> outputs;
361 for (const std::string& output : op->outputs) {
362 outputs.push_back(tensors_map.at(output));
363 }
364 const toco::OperatorSignature op_signature = {op.get(), &model};
365 const auto key = details::OperatorKey(op_signature, ops_by_type,
366 params.enable_select_tf_ops);
367 int op_index = operators_map.at(key);
368
369 auto tflite_op_it = ops_by_type.find(op->type);
370 BaseOperator* tflite_op = tflite_op_it == ops_by_type.end()
371 ? nullptr
372 : tflite_op_it->second.get();
373
374 // This is a custom op unless we can find it in ops_by_type, and even then
375 // it could be a custom op (such as kUnsupported).
376 auto options = Options::Custom(0);
377
378 std::vector<bool> mutating_input_variables;
379
380 // It is conceivable that an op is exportable via Serialize() but does not
381 // have a corresponding TFLITE builtin. In that case, when flex mode is
382 // enabled we should export it as a flex op, not as a native.
383 bool export_as_flex_op = !is_tflite_builtin(tflite_op) &&
384 key.is_flex_op() &&
385 !op->tensorflow_node_def.empty();
386 if (export_as_flex_op) {
387 auto fbb = WriteFlexOpOptions(op->tensorflow_node_def);
388 if (fbb) {
389 options = Options::Custom(builder->CreateVector(fbb->GetBuffer()));
390 }
391 } else if (tflite_op) {
392 options = tflite_op->Serialize(*op, builder);
393 mutating_input_variables = tflite_op->GetMutatingInputVariables(*op);
394
395 if (!mutating_input_variables.empty()) {
396 for (uint32_t i = 0; i < op->inputs.size(); ++i) {
397 if (!mutating_input_variables[i]) {
398 continue;
399 }
400 int32_t variable_tensor_index = tensors_map.at(op->inputs[i]);
401 variable_tensor_indices->insert(variable_tensor_index);
402 }
403 }
404 } else {
405 // We don't know much about this op. It doesn't have a serializer and
406 // it is not supposed to be exported as a flex op. We will treat it as
407 // a regular custom op: we will still create an operator for it, but it
408 // will not have any 'options'.
409 }
410
411 // The only supported CustomOptionFormat is FLEXBUFFERS now.
412 op_vector.push_back(CreateOperator(
413 *builder, op_index, builder->CreateVector(inputs),
414 builder->CreateVector(outputs), options.type, options.builtin,
415 options.custom, ::tflite::CustomOptionsFormat_FLEXBUFFERS,
416 builder->CreateVector(mutating_input_variables)));
417 }
418
419 return builder->CreateVector(op_vector);
420 }
421
ExportBuffers(const Model & model,const std::vector<Offset<Vector<uint8_t>>> & buffers_to_write,FlatBufferBuilder * builder)422 Offset<Vector<Offset<Buffer>>> ExportBuffers(
423 const Model& model,
424 const std::vector<Offset<Vector<uint8_t>>>& buffers_to_write,
425 FlatBufferBuilder* builder) {
426 std::vector<Offset<Buffer>> buffer_vector;
427 buffer_vector.reserve(buffers_to_write.size());
428 for (const auto buffer : buffers_to_write) {
429 buffer_vector.push_back(CreateBuffer(*builder, buffer));
430 }
431 return builder->CreateVector(buffer_vector);
432 }
433
Export(const Model & model,std::string * output_file_contents,const ExportParams & params)434 tensorflow::Status Export(const Model& model, std::string* output_file_contents,
435 const ExportParams& params) {
436 const auto ops_by_type = BuildOperatorByTypeMap(params.enable_select_tf_ops);
437 return Export(model, output_file_contents, params, ops_by_type);
438 }
439
ParseControlFlowErrors(std::set<std::string> * custom_ops,std::vector<std::string> * error_msgs)440 void ParseControlFlowErrors(std::set<std::string>* custom_ops,
441 std::vector<std::string>* error_msgs) {
442 std::set<std::string> unsupported_control_flow_ops;
443 // Check if unsupported ops contains control flow ops. It's impossible
444 // to implement these ops as custom ops at the moment.
445 for (const auto& op : *custom_ops) {
446 if (IsControlFlowOp(op)) {
447 unsupported_control_flow_ops.insert(op);
448 }
449 }
450 if (!unsupported_control_flow_ops.empty()) {
451 error_msgs->push_back(absl::StrCat(
452 "TensorFlow Lite currently doesn't support control flow ops: ",
453 absl::StrJoin(unsupported_control_flow_ops, ", "), ".",
454 " We are working on supporting control flow ops, please see github "
455 "issue at "
456 "https://github.com/tensorflow/tensorflow/issues/28485."));
457 }
458 // Remove control flow ops from `custom_ops` set so that they won't be
459 // reported again in later messages.
460 for (const auto& op : unsupported_control_flow_ops) {
461 custom_ops->erase(op);
462 }
463 }
464
465 // Exports appropriate model metadata.
466 // This consists of a 16-byte placeholder string buffer that will be populated
467 // with the model's minimum required runtime version after the model is
468 // finalized. We choose 16 bytes as it's the alignment of flatbuffer buffers
469 // and prevents wasted space if the final string is shorter than 16 bytes.
470 // Note: This logic matches that used in the MLIR flatbuffer export path.
ExportMetadata(const Model & model,std::vector<Offset<Vector<uint8_t>>> * buffers_to_write,FlatBufferBuilder * builder)471 flatbuffers::Offset<tflite::Metadata> ExportMetadata(
472 const Model& model, std::vector<Offset<Vector<uint8_t>>>* buffers_to_write,
473 FlatBufferBuilder* builder) {
474 auto metadata =
475 CreateMetadata(*builder, builder->CreateString("min_runtime_version"),
476 buffers_to_write->size());
477 const std::string min_runtime_placeholder = std::string(16, '\0');
478 buffers_to_write->push_back(builder->CreateVector(
479 reinterpret_cast<const uint8_t*>(min_runtime_placeholder.data()),
480 min_runtime_placeholder.size()));
481 return metadata;
482 }
483
Export(const Model & model,std::string * output_file_contents,const ExportParams & params,const std::map<OperatorType,std::unique_ptr<BaseOperator>> & ops_by_type)484 tensorflow::Status Export(
485 const Model& model, std::string* output_file_contents,
486 const ExportParams& params,
487 const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type) {
488 for (const std::string& input_array : model.GetInvalidInputArrays()) {
489 if (model.HasArray(input_array)) {
490 return tensorflow::errors::InvalidArgument(
491 absl::StrCat("Placeholder ", input_array,
492 " should be specified by "
493 "input_arrays."));
494 }
495 }
496
497 flatbuffers::FlatBufferBuilder builder(/*initial_size=*/10240);
498
499 details::TensorsMap tensors_map;
500 details::LoadTensorsMap(model, &tensors_map);
501
502 details::OperatorsMap operators_map;
503 details::LoadOperatorsMap(model, &operators_map, ops_by_type,
504 params.enable_select_tf_ops);
505
506 std::vector<Offset<Vector<uint8_t>>> buffers_to_write;
507 // Insert an empty buffer to the beginning of the list.
508 buffers_to_write.push_back(0);
509
510 auto op_codes =
511 ExportOperatorCodes(model, ops_by_type, operators_map, &builder, params);
512
513 for (const auto& op : model.operators) {
514 if (op->type == OperatorType::kFakeQuant) {
515 LOG(WARNING) << "FAKE_QUANT operation " << LogName(*op)
516 << " was not converted. If running quantized make sure you "
517 "are passing --inference_type=QUANTIZED_UINT8 and values "
518 "for --std_values and --mean_values.";
519 }
520 }
521
522 // The set of used builtin ops.
523 std::set<std::string> builtin_ops;
524 // The set of custom ops (not including Flex ops).
525 std::set<std::string> custom_ops;
526 // The set of Flex ops which are not supported.
527 std::set<std::string> unsupported_flex_ops;
528
529 for (const auto& it : operators_map) {
530 const details::OperatorKey& key = it.first;
531 if (key.is_custom_op()) {
532 custom_ops.insert(key.custom_code());
533 }
534 if (key.is_unsupported_flex_op()) {
535 unsupported_flex_ops.insert(key.flex_tensorflow_op());
536 }
537 if (!key.is_custom_op() && !key.is_flex_op() &&
538 !key.is_unsupported_flex_op()) {
539 builtin_ops.insert(EnumNameBuiltinOperator(key.type()));
540 }
541 }
542
543 if (!custom_ops.empty()) {
544 if (!params.allow_custom_ops) {
545 auto please_report_bug_message = []() {
546 return "We are continually in the process of adding support to "
547 "TensorFlow Lite for more ops. It would be helpful if you could "
548 "inform us of how this conversion went by opening a github "
549 "issue at "
550 "https://github.com/tensorflow/tensorflow/issues/new?template="
551 "40-tflite-op-request.md\n and pasting the following:\n\n";
552 };
553
554 std::vector<std::string> error_msgs;
555 ParseControlFlowErrors(&custom_ops, &error_msgs);
556
557 // Remove ExpandDims and ReorderAxes from unimplemented list unless they
558 // compose the list. Both ops are removed during graph transformations.
559 // However, if an op is unimplemented earlier in the model, the graph
560 // transformation is unable to run because the output shape is not
561 // defined. This causes unnecessary confusion during model conversion
562 // time.
563 std::set<std::string> custom_ops_final;
564 for (const auto& op_type : custom_ops) {
565 if (op_type != "ReorderAxes" && op_type != "ExpandDims") {
566 custom_ops_final.insert(op_type);
567 }
568 }
569 if (custom_ops_final.empty()) {
570 custom_ops_final = custom_ops;
571 }
572
573 if (!custom_ops_final.empty()) {
574 if (params.enable_select_tf_ops) {
575 error_msgs.push_back(absl::StrCat(
576 "Some of the operators in the model are not supported by "
577 "the standard TensorFlow Lite runtime and are not recognized "
578 "by "
579 "TensorFlow. If you have a custom "
580 "implementation for them you can disable this error with "
581 "--allow_custom_ops, or by setting allow_custom_ops=True "
582 "when calling tf.lite.TFLiteConverter(). Here is a list "
583 "of builtin operators you are using: ",
584 absl::StrJoin(builtin_ops, ", "),
585 ". Here is a list "
586 "of operators for which you will need custom implementations: ",
587 absl::StrJoin(custom_ops_final, ", "), "."));
588 } else {
589 error_msgs.push_back(absl::StrCat(
590 "Some of the operators in the model are not supported by "
591 "the standard TensorFlow Lite runtime. If those are native "
592 "TensorFlow operators, you might be able to use the extended "
593 "runtime by passing --enable_select_tf_ops, or by setting "
594 "target_ops=TFLITE_BUILTINS,SELECT_TF_OPS when calling "
595 "tf.lite.TFLiteConverter(). Otherwise, if you have a "
596 "custom implementation for them you can disable this error "
597 "with "
598 "--allow_custom_ops, or by setting allow_custom_ops=True "
599 "when calling tf.lite.TFLiteConverter(). Here is a list "
600 "of builtin operators you are using: ",
601 absl::StrJoin(builtin_ops, ", "),
602 ". Here is a list "
603 "of operators for which you will need custom implementations: ",
604 absl::StrJoin(custom_ops_final, ", "), "."));
605 }
606 }
607 if (!error_msgs.empty()) {
608 return tensorflow::errors::InvalidArgument(absl::StrCat(
609 please_report_bug_message(), absl::StrJoin(error_msgs, " ")));
610 }
611 }
612 }
613
614 if (!unsupported_flex_ops.empty()) {
615 return tensorflow::errors::InvalidArgument(
616 absl::StrCat("Some of the operators in the model are not supported by "
617 "TensorFlow Flex runtime: ",
618 absl::StrJoin(unsupported_flex_ops, ", "), "."));
619 }
620
621 std::set<int32_t> variable_tensor_indices;
622 auto ops = ExportOperators(model, ops_by_type, operators_map, tensors_map,
623 &builder, &variable_tensor_indices, params);
624
625 auto tensors = ExportTensors(model, tensors_map, &builder, &buffers_to_write,
626 variable_tensor_indices);
627 auto inputs = ExportInputTensors(model, tensors_map, &builder);
628 auto outputs = ExportOutputTensors(model, tensors_map, &builder);
629
630 // TODO(aselle): add support to toco for multiple subgraphs.
631 auto subgraph = CreateSubGraph(builder, tensors, inputs, outputs, ops,
632 /* name */ 0);
633 std::vector<flatbuffers::Offset<SubGraph>> subgraphs = {subgraph};
634
635 // TODO(wangtz): offline memory planning for activation Tensors.
636 if (!params.allow_dynamic_tensors) {
637 return tensorflow::errors::Unimplemented(
638 "Unsupported flag: allow_dynamic_tensors. Offline memory planning is "
639 "not implemented yet.");
640 }
641
642 auto metadata = ExportMetadata(model, &buffers_to_write, &builder);
643 std::vector<flatbuffers::Offset<Metadata>> metadatas = {metadata};
644
645 auto buffers = ExportBuffers(model, buffers_to_write, &builder);
646 auto description = builder.CreateString("TOCO Converted.");
647
648 auto new_model_location =
649 CreateModel(builder, TFLITE_SCHEMA_VERSION, op_codes,
650 builder.CreateVector(subgraphs), description, buffers,
651 /* metadata_buffer */ 0, builder.CreateVector(metadatas));
652 ::tflite::FinishModelBuffer(builder, new_model_location);
653 ::tflite::UpdateMinimumRuntimeVersionForModel(builder.GetBufferPointer());
654
655 if (params.quantize_weights == QuantizedBufferType::NONE) {
656 WriteModelToString(builder, output_file_contents);
657 } else {
658 // Call the quantize_weights tool.
659 LOG(INFO) << "Quantizing TFLite model after conversion to flatbuffer. "
660 "dump_graphviz will only output the model before this "
661 "transformation. To visualize the output graph use "
662 "lite/tools/optimize.py.";
663 flatbuffers::FlatBufferBuilder q_builder(/*initial_size=*/10240);
664 const uint8_t* buffer = builder.GetBufferPointer();
665 const ::tflite::Model* input_model = ::tflite::GetModel(buffer);
666 ::tflite::optimize::BufferType quantized_type;
667 if (params.quantize_weights == QuantizedBufferType::INT8) {
668 quantized_type = ::tflite::optimize::BufferType::QUANTIZED_INT8;
669 } else if (params.quantize_weights == QuantizedBufferType::FLOAT16) {
670 quantized_type = ::tflite::optimize::BufferType::QUANTIZED_FLOAT16;
671 } else {
672 return tensorflow::errors::InvalidArgument(
673 "Quantized type not recognized");
674 }
675 if (::tflite::optimize::QuantizeWeights(
676 &q_builder, input_model, quantized_type,
677 !params.disable_per_channel,
678 ::tflite::optimize::QuantizerType::OLD_QUANTIZER) != kTfLiteOk) {
679 return tensorflow::errors::InvalidArgument(
680 "Quantize weights transformation failed.");
681 }
682 WriteModelToString(q_builder, output_file_contents);
683 }
684
685 return tensorflow::Status();
686 }
687
688 } // namespace tflite
689
690 } // namespace toco
691