xref: /aosp_15_r20/external/tensorflow/tensorflow/lite/toco/tflite/export.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/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