| /* Copyright 2019 The TensorFlow Authors. All Rights Reserved. |
| |
| Licensed under the Apache License, Version 2.0 (the "License"); |
| you may not use this file except in compliance with the License. |
| You may obtain a copy of the License at |
| |
| http://www.apache.org/licenses/LICENSE-2.0 |
| |
| Unless required by applicable law or agreed to in writing, software |
| distributed under the License is distributed on an "AS IS" BASIS, |
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| See the License for the specific language governing permissions and |
| limitations under the License. |
| ==============================================================================*/ |
| #include "tensorflow/lite/toco/tflite/export.h" |
| |
| #include "flatbuffers/flexbuffers.h" |
| #include "absl/strings/str_join.h" |
| #include "tensorflow/core/lib/core/errors.h" |
| #include "tensorflow/lite/context.h" |
| #include "tensorflow/lite/schema/schema_generated.h" |
| #include "tensorflow/lite/toco/tflite/op_version.h" |
| #include "tensorflow/lite/toco/tflite/operator.h" |
| #include "tensorflow/lite/toco/tflite/types.h" |
| #include "tensorflow/lite/toco/tooling_util.h" |
| #include "tensorflow/lite/tools/optimize/quantize_weights.h" |
| #include "tensorflow/lite/version.h" |
| |
| namespace toco { |
| |
| namespace tflite { |
| |
| using flatbuffers::FlatBufferBuilder; |
| using flatbuffers::Offset; |
| using flatbuffers::Vector; |
| using ::tflite::Buffer; |
| using ::tflite::BuiltinOperator; |
| using ::tflite::BuiltinOperator_CUSTOM; |
| using ::tflite::BuiltinOperator_MAX; |
| using ::tflite::BuiltinOperator_MIN; |
| using ::tflite::CreateBuffer; |
| using ::tflite::CreateMetadata; |
| using ::tflite::CreateModel; |
| using ::tflite::CreateOperator; |
| using ::tflite::CreateTensor; |
| using ::tflite::Metadata; |
| using ::tflite::Operator; |
| using ::tflite::OperatorCode; |
| using ::tflite::SubGraph; |
| using ::tflite::Tensor; |
| |
| namespace { |
| |
| // Check if a TensorFlow Op is a control flow op by its name. |
| bool IsControlFlowOp(const string& tensorflow_op) { |
| // Technically this is equalivent to `::tensorflow::Node::IsControlFlow()`. |
| // It requires to construct a `::tensorflow::Graph` to use that helper |
| // function, so we simply hardcode the list of control flow ops here. |
| if (tensorflow_op == "Switch" || tensorflow_op == "RefSwitch" || |
| tensorflow_op == "Merge" || tensorflow_op == "RefMerge" || |
| tensorflow_op == "Enter" || tensorflow_op == "RefEnter" || |
| tensorflow_op == "Exit" || tensorflow_op == "RefExit" || |
| tensorflow_op == "NextIteration" || tensorflow_op == "RefNextIteration") { |
| return true; |
| } |
| // TODO(ycling): Also check how to handle Variable ops and Assign ops. |
| return false; |
| } |
| |
| // Check if a TensorFlow Op is unsupported by the Flex runtime. |
| bool IsUnsupportedFlexOp(const string& tensorflow_op) { |
| if (IsControlFlowOp(tensorflow_op)) { |
| return true; |
| } |
| // `HashTableV2` isn't supported for now since it requires an additional |
| // initialization step. |
| // TODO(b/117651199): Support `HashTableV2` with Flex runtime. |
| if (tensorflow_op == "HashTableV2") { |
| return true; |
| } |
| return false; |
| } |
| |
| // Map from operator name to TF Lite enum value, for all builtins. |
| const std::map<string, BuiltinOperator>& GetBuiltinOpsMap() { |
| static std::map<string, BuiltinOperator>* builtin_ops = nullptr; |
| if (builtin_ops == nullptr) { |
| builtin_ops = new std::map<string, BuiltinOperator>(); |
| |
| for (int i = BuiltinOperator_MIN; i <= BuiltinOperator_MAX; ++i) { |
| BuiltinOperator op = static_cast<BuiltinOperator>(i); |
| string name = EnumNameBuiltinOperator(op); |
| if (op != BuiltinOperator_CUSTOM && !name.empty()) { |
| (*builtin_ops)[name] = op; |
| } |
| } |
| } |
| return *builtin_ops; |
| } |
| |
| void WriteModelToString(const flatbuffers::FlatBufferBuilder& builder, |
| string* file_contents) { |
| const uint8_t* buffer = builder.GetBufferPointer(); |
| int size = builder.GetSize(); |
| *file_contents = string(reinterpret_cast<const char*>(buffer), size); |
| } |
| |
| } // Anonymous namespace. |
| |
| namespace details { |
| |
| OperatorKey::OperatorKey( |
| const ::toco::OperatorSignature& op_signature, |
| const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type, |
| bool enable_select_tf_ops) { |
| // Get the op name (by Toco definition). |
| const ::toco::Operator& op = *op_signature.op; |
| string name = HelpfulOperatorTypeName(op); |
| |
| bool is_builtin = false; |
| const auto& builtin_ops = GetBuiltinOpsMap(); |
| if (ops_by_type.count(op.type) != 0) { |
| version_ = ops_by_type.at(op.type)->GetVersion(op_signature); |
| name = ops_by_type.at(op.type)->name(); |
| is_builtin = (builtin_ops.count(name) > 0); |
| } |
| |
| if (is_builtin) { |
| // For TFLite supported builtin ops, find out its BuiltinOperator enum used |
| // in FlatBuffer. |
| type_ = builtin_ops.at(name); |
| return; |
| } |
| // The logic below is all for custom ops or Flex ops. |
| is_custom_op_ = true; |
| type_ = BuiltinOperator_CUSTOM; |
| |
| if (op.type == OperatorType::kUnsupported) { |
| const TensorFlowUnsupportedOperator& unsupported_op = |
| static_cast<const TensorFlowUnsupportedOperator&>(op); |
| const auto tensorflow_op = unsupported_op.tensorflow_op; |
| |
| if (ShouldExportAsFlexOp(enable_select_tf_ops, |
| unsupported_op.tensorflow_op)) { |
| is_custom_op_ = false; |
| is_flex_op_ = true; |
| flex_tensorflow_op_ = tensorflow_op; |
| custom_code_ = |
| string(::tflite::kFlexCustomCodePrefix) + flex_tensorflow_op_; |
| } else { |
| custom_code_ = tensorflow_op; |
| } |
| } else if (enable_select_tf_ops && !op.tensorflow_node_def.empty()) { |
| // For Toco-supported/TFLite-unsupported ops, if the TensorFlow NodeDef |
| // is retained in the Toco Operator, we produce a Flex op if Flex mode |
| // is enabled. |
| is_custom_op_ = false; |
| is_flex_op_ = true; |
| flex_tensorflow_op_ = name; |
| custom_code_ = |
| string(::tflite::kFlexCustomCodePrefix) + flex_tensorflow_op_; |
| } else { |
| // If Flex is disabled or the original TensorFlow NodeDef isn't available, |
| // we produce a custom op. This gives developers a chance to implement |
| // custom ops. |
| custom_code_ = name; |
| } |
| |
| if (is_flex_op_) { |
| if (IsUnsupportedFlexOp(flex_tensorflow_op_)) { |
| is_unsupported_flex_op_ = true; |
| } |
| } |
| } |
| |
| void LoadTensorsMap(const Model& model, TensorsMap* tensors_map) { |
| // First find a list of unique array names. |
| std::set<string> names; |
| for (const auto& array_pair : model.GetArrayMap()) { |
| names.insert(array_pair.first); |
| } |
| |
| // Now assign indices to them and fill in the map. |
| int index = 0; |
| for (const auto& name : names) { |
| (*tensors_map)[name] = index; |
| ++index; |
| } |
| } |
| |
| void LoadOperatorsMap( |
| const Model& model, OperatorsMap* operators_map, |
| const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type, |
| bool enable_select_tf_ops) { |
| // First find a list of unique operator types. |
| std::set<OperatorKey> keys; |
| for (const auto& op : model.operators) { |
| const toco::OperatorSignature op_signature = {op.get(), &model}; |
| keys.insert(OperatorKey(op_signature, ops_by_type, enable_select_tf_ops)); |
| } |
| // Now assign indices to them and fill in the map. |
| int index = 0; |
| for (const auto& key : keys) { |
| (*operators_map)[key] = index; |
| ++index; |
| } |
| } |
| |
| } // namespace details |
| |
| Offset<Vector<Offset<Tensor>>> 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) { |
| // In the end we will need to produce a vector sorted by the indices of the |
| // tensors in the tensors_map. |
| std::map<int, Offset<Tensor>> ordered_tensors; |
| |
| for (const auto& array_pair : model.GetArrayMap()) { |
| const string& tensor_name = array_pair.first; |
| const toco::Array& array = *array_pair.second; |
| |
| int buffer_index = buffers_to_write->size(); |
| auto type = DataType::Serialize(array.data_type); |
| buffers_to_write->push_back(DataBuffer::Serialize(array, builder)); |
| |
| std::vector<int> shape; |
| if (array.has_shape()) { |
| shape.reserve(array.shape().dims().size()); |
| for (const auto& d : array.shape().dims()) { |
| shape.push_back(d); |
| } |
| } |
| |
| Offset<Vector<float>> min; |
| Offset<Vector<float>> max; |
| Offset<Vector<float>> scale; |
| Offset<Vector<int64_t>> zero_point; |
| if (array.minmax) { |
| min = builder->CreateVector( |
| std::vector<float>{static_cast<float>(array.minmax->min)}); |
| max = builder->CreateVector( |
| std::vector<float>{static_cast<float>(array.minmax->max)}); |
| } |
| if (array.quantization_params) { |
| scale = builder->CreateVector(std::vector<float>{ |
| static_cast<float>(array.quantization_params->scale)}); |
| zero_point = builder->CreateVector( |
| std::vector<int64_t>{array.quantization_params->zero_point}); |
| } |
| auto q_param = ::tflite::CreateQuantizationParameters(*builder, min, max, |
| scale, zero_point); |
| |
| int index = tensors_map.at(tensor_name); |
| bool is_variable = |
| variable_tensor_indices.find(index) != variable_tensor_indices.end(); |
| ordered_tensors[index] = |
| CreateTensor(*builder, builder->CreateVector(shape), type, buffer_index, |
| builder->CreateString(tensor_name), q_param, is_variable); |
| } |
| |
| std::vector<Offset<Tensor>> tensor_vector; |
| tensor_vector.reserve(ordered_tensors.size()); |
| for (const auto& tensor : ordered_tensors) { |
| tensor_vector.push_back(tensor.second); |
| } |
| |
| return builder->CreateVector(tensor_vector); |
| } |
| |
| Offset<Vector<int32_t>> ExportInputTensors( |
| const Model& model, const details::TensorsMap& tensors_map, |
| FlatBufferBuilder* builder) { |
| std::vector<int32_t> inputs; |
| for (const auto& input : model.flags.input_arrays()) { |
| inputs.push_back(tensors_map.at(input.name())); |
| } |
| return builder->CreateVector<int32_t>(inputs); |
| } |
| |
| Offset<Vector<int32_t>> ExportOutputTensors( |
| const Model& model, const details::TensorsMap& tensors_map, |
| FlatBufferBuilder* builder) { |
| std::vector<int32_t> outputs; |
| for (const string& output : model.flags.output_arrays()) { |
| outputs.push_back(tensors_map.at(output)); |
| } |
| return builder->CreateVector<int32_t>(outputs); |
| } |
| |
| Offset<Vector<Offset<OperatorCode>>> 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) { |
| // Map from operator name to TF Lite enum value, for all builtins. |
| std::map<string, BuiltinOperator> builtin_ops; |
| for (int i = BuiltinOperator_MIN; i <= BuiltinOperator_MAX; ++i) { |
| BuiltinOperator op = static_cast<BuiltinOperator>(i); |
| string name = EnumNameBuiltinOperator(op); |
| if (op != BuiltinOperator_CUSTOM && !name.empty()) { |
| builtin_ops[name] = op; |
| } |
| } |
| |
| // We will need to produce a vector of codes in the same order as they |
| // appear in the operators_map. |
| std::map<int, Offset<OperatorCode>> ordered_opcodes; |
| |
| for (const auto& op : model.operators) { |
| const toco::OperatorSignature op_signature = {op.get(), &model}; |
| const details::OperatorKey operator_key = details::OperatorKey( |
| op_signature, ops_by_type, params.enable_select_tf_ops); |
| int op_index = operators_map.at(operator_key); |
| |
| flatbuffers::Offset<flatbuffers::String> custom_code = 0; |
| if (!operator_key.custom_code().empty()) { |
| custom_code = builder->CreateString(operator_key.custom_code()); |
| } |
| |
| ordered_opcodes[op_index] = CreateOperatorCode( |
| *builder, operator_key.type(), custom_code, operator_key.version()); |
| } |
| |
| std::vector<Offset<OperatorCode>> opcode_vector; |
| opcode_vector.reserve(ordered_opcodes.size()); |
| for (const auto& opcode : ordered_opcodes) { |
| opcode_vector.push_back(opcode.second); |
| } |
| |
| return builder->CreateVector(opcode_vector); |
| } |
| |
| Offset<Vector<Offset<Operator>>> 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) { |
| variable_tensor_indices->clear(); |
| |
| auto is_tflite_builtin = [](const BaseOperator* op) { |
| const auto& tflite_builtins = GetBuiltinOpsMap(); |
| return (op && tflite_builtins.find(op->name()) != tflite_builtins.end()); |
| }; |
| |
| // The operators are in execution order, so we just follow tf.mini order. |
| std::vector<Offset<Operator>> op_vector; |
| for (const auto& op : model.operators) { |
| std::vector<int32_t> inputs; |
| for (const string& input : op->inputs) { |
| // -1 is the ID for optional tensor in TFLite output |
| int id = model.IsOptionalArray(input) ? -1 : tensors_map.at(input); |
| inputs.push_back(id); |
| } |
| std::vector<int32_t> outputs; |
| for (const string& output : op->outputs) { |
| outputs.push_back(tensors_map.at(output)); |
| } |
| const toco::OperatorSignature op_signature = {op.get(), &model}; |
| const auto key = details::OperatorKey(op_signature, ops_by_type, |
| params.enable_select_tf_ops); |
| int op_index = operators_map.at(key); |
| |
| auto tflite_op_it = ops_by_type.find(op->type); |
| BaseOperator* tflite_op = tflite_op_it == ops_by_type.end() |
| ? nullptr |
| : tflite_op_it->second.get(); |
| |
| // This is a custom op unless we can find it in ops_by_type, and even then |
| // it could be a custom op (such as kUnsupported). |
| auto options = Options::Custom(0); |
| |
| std::vector<bool> mutating_input_variables; |
| |
| // It is conceivable that an op is exportable via Serialize() but does not |
| // have a corresponding TFLITE builtin. In that case, when flex mode is |
| // enabled we should export it as a flex op, not as a native. |
| bool export_as_flex_op = !is_tflite_builtin(tflite_op) && |
| key.is_flex_op() && |
| !op->tensorflow_node_def.empty(); |
| if (export_as_flex_op) { |
| auto fbb = WriteFlexOpOptions(op->tensorflow_node_def); |
| if (fbb) { |
| options = Options::Custom(builder->CreateVector(fbb->GetBuffer())); |
| } |
| } else if (tflite_op) { |
| options = tflite_op->Serialize(*op, builder); |
| mutating_input_variables = tflite_op->GetMutatingInputVariables(*op); |
| |
| if (!mutating_input_variables.empty()) { |
| for (size_t i = 0; i < op->inputs.size(); ++i) { |
| if (!mutating_input_variables[i]) { |
| continue; |
| } |
| int32_t variable_tensor_index = tensors_map.at(op->inputs[i]); |
| variable_tensor_indices->insert(variable_tensor_index); |
| } |
| } |
| } else { |
| // We don't know much about this op. It doesn't have a serializer and |
| // it is not supposed to be exported as a flex op. We will treat it as |
| // a regular custom op: we will still create an operator for it, but it |
| // will not have any 'options'. |
| } |
| |
| // The only supported CustomOptionFormat is FLEXBUFFERS now. |
| op_vector.push_back(CreateOperator( |
| *builder, op_index, builder->CreateVector(inputs), |
| builder->CreateVector(outputs), options.type, options.builtin, |
| options.custom, ::tflite::CustomOptionsFormat_FLEXBUFFERS, |
| builder->CreateVector(mutating_input_variables))); |
| } |
| |
| return builder->CreateVector(op_vector); |
| } |
| |
| Offset<Vector<Offset<Buffer>>> ExportBuffers( |
| const Model& model, |
| const std::vector<Offset<Vector<uint8_t>>>& buffers_to_write, |
| FlatBufferBuilder* builder) { |
| std::vector<Offset<Buffer>> buffer_vector; |
| buffer_vector.reserve(buffers_to_write.size()); |
| for (const auto buffer : buffers_to_write) { |
| buffer_vector.push_back(CreateBuffer(*builder, buffer)); |
| } |
| return builder->CreateVector(buffer_vector); |
| } |
| |
| tensorflow::Status Export(const Model& model, string* output_file_contents, |
| const ExportParams& params) { |
| const auto ops_by_type = BuildOperatorByTypeMap(params.enable_select_tf_ops); |
| return Export(model, output_file_contents, params, ops_by_type); |
| } |
| |
| void ParseControlFlowErrors(std::set<string>* custom_ops, |
| std::vector<string>* error_msgs) { |
| std::set<string> unsupported_control_flow_ops; |
| // Check if unsupported ops contains control flow ops. It's impossible |
| // to implement these ops as custom ops at the moment. |
| for (const auto& op : *custom_ops) { |
| if (IsControlFlowOp(op)) { |
| unsupported_control_flow_ops.insert(op); |
| } |
| } |
| if (!unsupported_control_flow_ops.empty()) { |
| error_msgs->push_back(absl::StrCat( |
| "TensorFlow Lite currently doesn't support control flow ops: ", |
| absl::StrJoin(unsupported_control_flow_ops, ", "), ".", |
| " We are working on supporting control flow ops, please see github " |
| "issue at " |
| "https://github.com/tensorflow/tensorflow/issues/28485.")); |
| } |
| // Remove control flow ops from `custom_ops` set so that they won't be |
| // reported again in later messages. |
| for (const auto& op : unsupported_control_flow_ops) { |
| custom_ops->erase(op); |
| } |
| } |
| |
| // Exports a string buffer that contains the model's minimum required runtime |
| // version. |
| void ExportModelVersionBuffer( |
| const Model& model, std::vector<Offset<Vector<uint8_t>>>* buffers_to_write, |
| FlatBufferBuilder* builder) { |
| const std::string min_runtime = GetMinimumRuntimeVersionForModel(model); |
| buffers_to_write->push_back(builder->CreateVector( |
| reinterpret_cast<const uint8_t*>(min_runtime.data()), |
| min_runtime.size())); |
| } |
| |
| tensorflow::Status Export( |
| const Model& model, string* output_file_contents, |
| const ExportParams& params, |
| const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type) { |
| for (const string& input_array : model.GetInvalidInputArrays()) { |
| if (model.HasArray(input_array)) { |
| return tensorflow::errors::InvalidArgument(absl::StrCat( |
| "Placeholder ", input_array, " should be specied by input_arrays.")); |
| } |
| } |
| |
| flatbuffers::FlatBufferBuilder builder(/*initial_size=*/10240); |
| |
| details::TensorsMap tensors_map; |
| details::LoadTensorsMap(model, &tensors_map); |
| |
| details::OperatorsMap operators_map; |
| details::LoadOperatorsMap(model, &operators_map, ops_by_type, |
| params.enable_select_tf_ops); |
| |
| std::vector<Offset<Vector<uint8_t>>> buffers_to_write; |
| // Insert an empty buffer to the beginning of the list. |
| buffers_to_write.push_back(0); |
| |
| auto op_codes = |
| ExportOperatorCodes(model, ops_by_type, operators_map, &builder, params); |
| |
| for (const auto& op : model.operators) { |
| if (op->type == OperatorType::kFakeQuant) { |
| LOG(WARNING) << "FAKE_QUANT operation " << LogName(*op) |
| << " was not converted. If running quantized make sure you " |
| "are passing --inference_type=QUANTIZED_UINT8 and values " |
| "for --std_values and --mean_values."; |
| } |
| } |
| |
| // The set of used builtin ops. |
| std::set<string> builtin_ops; |
| // The set of custom ops (not including Flex ops). |
| std::set<string> custom_ops; |
| // The set of Flex ops which are not supported. |
| std::set<string> unsupported_flex_ops; |
| |
| for (const auto& it : operators_map) { |
| const details::OperatorKey& key = it.first; |
| if (key.is_custom_op()) { |
| custom_ops.insert(key.custom_code()); |
| } |
| if (key.is_unsupported_flex_op()) { |
| unsupported_flex_ops.insert(key.flex_tensorflow_op()); |
| } |
| if (!key.is_custom_op() && !key.is_flex_op() && |
| !key.is_unsupported_flex_op()) { |
| builtin_ops.insert(EnumNameBuiltinOperator(key.type())); |
| } |
| } |
| |
| if (!custom_ops.empty()) { |
| if (!params.allow_custom_ops) { |
| auto please_report_bug_message = []() { |
| return "We are continually in the process of adding support to " |
| "TensorFlow Lite for more ops. It would be helpful if you could " |
| "inform us of how this conversion went by opening a github " |
| "issue at " |
| "https://github.com/tensorflow/tensorflow/issues/new?template=" |
| "40-tflite-op-request.md\n and pasting the following:\n\n"; |
| }; |
| |
| std::vector<string> error_msgs; |
| ParseControlFlowErrors(&custom_ops, &error_msgs); |
| |
| // Remove ExpandDims and ReorderAxes from unimplemented list unless they |
| // compose the list. Both ops are removed during graph transformations. |
| // However, if an op is unimplemented earlier in the model, the graph |
| // transformation is unable to run because the output shape is not |
| // defined. This causes unnecessary confusion during model conversion |
| // time. |
| std::set<string> custom_ops_final; |
| for (const auto& op_type : custom_ops) { |
| if (op_type != "ReorderAxes" && op_type != "ExpandDims") { |
| custom_ops_final.insert(op_type); |
| } |
| } |
| if (custom_ops_final.empty()) { |
| custom_ops_final = custom_ops; |
| } |
| |
| if (!custom_ops_final.empty()) { |
| if (params.enable_select_tf_ops) { |
| error_msgs.push_back(absl::StrCat( |
| "Some of the operators in the model are not supported by " |
| "the standard TensorFlow Lite runtime and are not recognized " |
| "by " |
| "TensorFlow. If you have a custom " |
| "implementation for them you can disable this error with " |
| "--allow_custom_ops, or by setting allow_custom_ops=True " |
| "when calling tf.lite.TFLiteConverter(). Here is a list " |
| "of builtin operators you are using: ", |
| absl::StrJoin(builtin_ops, ", "), |
| ". Here is a list " |
| "of operators for which you will need custom implementations: ", |
| absl::StrJoin(custom_ops_final, ", "), ".")); |
| } else { |
| error_msgs.push_back(absl::StrCat( |
| "Some of the operators in the model are not supported by " |
| "the standard TensorFlow Lite runtime. If those are native " |
| "TensorFlow operators, you might be able to use the extended " |
| "runtime by passing --enable_select_tf_ops, or by setting " |
| "target_ops=TFLITE_BUILTINS,SELECT_TF_OPS when calling " |
| "tf.lite.TFLiteConverter(). Otherwise, if you have a " |
| "custom implementation for them you can disable this error " |
| "with " |
| "--allow_custom_ops, or by setting allow_custom_ops=True " |
| "when calling tf.lite.TFLiteConverter(). Here is a list " |
| "of builtin operators you are using: ", |
| absl::StrJoin(builtin_ops, ", "), |
| ". Here is a list " |
| "of operators for which you will need custom implementations: ", |
| absl::StrJoin(custom_ops_final, ", "), ".")); |
| } |
| } |
| if (!error_msgs.empty()) { |
| return tensorflow::errors::InvalidArgument(absl::StrCat( |
| please_report_bug_message(), absl::StrJoin(error_msgs, " "))); |
| } |
| } |
| } |
| |
| if (!unsupported_flex_ops.empty()) { |
| return tensorflow::errors::InvalidArgument( |
| absl::StrCat("Some of the operators in the model are not supported by " |
| "TensorFlow Flex runtime: ", |
| absl::StrJoin(unsupported_flex_ops, ", "), ".")); |
| } |
| |
| std::set<int32_t> variable_tensor_indices; |
| auto ops = ExportOperators(model, ops_by_type, operators_map, tensors_map, |
| &builder, &variable_tensor_indices, params); |
| |
| auto tensors = ExportTensors(model, tensors_map, &builder, &buffers_to_write, |
| variable_tensor_indices); |
| auto inputs = ExportInputTensors(model, tensors_map, &builder); |
| auto outputs = ExportOutputTensors(model, tensors_map, &builder); |
| |
| // TODO(aselle): add support to toco for multiple subgraphs. |
| auto subgraph = CreateSubGraph(builder, tensors, inputs, outputs, ops, |
| /* name */ 0); |
| std::vector<flatbuffers::Offset<SubGraph>> subgraphs = {subgraph}; |
| |
| // TODO(wangtz): offline memory planning for activation Tensors. |
| if (!params.allow_dynamic_tensors) { |
| return tensorflow::errors::Unimplemented( |
| "Unsupported flag: allow_dynamic_tensors. Offline memory planning is " |
| "not implemented yet."); |
| } |
| |
| // Write the minimum required runtime version into metadata. |
| auto metadata = |
| CreateMetadata(builder, builder.CreateString("min_runtime_version"), |
| buffers_to_write.size()); |
| ExportModelVersionBuffer(model, &buffers_to_write, &builder); |
| std::vector<flatbuffers::Offset<Metadata>> metadatas = {metadata}; |
| |
| auto buffers = ExportBuffers(model, buffers_to_write, &builder); |
| auto description = builder.CreateString("TOCO Converted."); |
| |
| auto new_model_location = |
| CreateModel(builder, TFLITE_SCHEMA_VERSION, op_codes, |
| builder.CreateVector(subgraphs), description, buffers, |
| /* metadata_buffer */ 0, builder.CreateVector(metadatas)); |
| ::tflite::FinishModelBuffer(builder, new_model_location); |
| |
| if (params.quantize_weights == QuantizedBufferType::NONE) { |
| WriteModelToString(builder, output_file_contents); |
| } else { |
| // Call the quantize_weights tool. |
| LOG(INFO) << "Quantizing TFLite model after conversion to flatbuffer. " |
| "dump_graphviz will only output the model before this " |
| "transformation. To visualize the output graph use " |
| "lite/tools/optimize.py."; |
| flatbuffers::FlatBufferBuilder q_builder(/*initial_size=*/10240); |
| const uint8_t* buffer = builder.GetBufferPointer(); |
| const ::tflite::Model* input_model = ::tflite::GetModel(buffer); |
| ::tflite::optimize::BufferType quantized_type; |
| if (params.quantize_weights == QuantizedBufferType::INT8) { |
| quantized_type = ::tflite::optimize::BufferType::QUANTIZED_INT8; |
| } else if (params.quantize_weights == QuantizedBufferType::FLOAT16) { |
| quantized_type = ::tflite::optimize::BufferType::QUANTIZED_FLOAT16; |
| } else { |
| return tensorflow::errors::InvalidArgument( |
| "Quantized type not recognized"); |
| } |
| if (::tflite::optimize::QuantizeWeights(&q_builder, input_model, |
| quantized_type) != kTfLiteOk) { |
| return tensorflow::errors::InvalidArgument( |
| "Quantize weights transformation failed."); |
| } |
| WriteModelToString(q_builder, output_file_contents); |
| } |
| |
| return tensorflow::Status(); |
| } |
| |
| } // namespace tflite |
| |
| } // namespace toco |