wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/lite/toco/tflite/BUILD
+204
View File
@@ -0,0 +1,204 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
)
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
licenses = ["notice"],
)
cc_library(
name = "operator",
srcs = [
"operator.cc",
],
hdrs = [
"builtin_operator.h",
"custom_operator.h",
"operator.h",
"simple_operator.h",
],
compatible_with = get_compatible_with_portable(),
visibility = [
"//tensorflow/lite/toco:__subpackages__",
],
deps = [
":types",
"//tensorflow/compiler/mlir/lite/delegates/flex:allowlisted_flex_ops_lib",
"//tensorflow/compiler/mlir/lite/tools/versioning",
"//tensorflow/compiler/mlir/lite/tools/versioning:op_signature",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/lite/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/schema:schema_fbs",
"//tensorflow/lite/toco:graph_transformations",
"//tensorflow/lite/toco:model",
"//tensorflow/lite/toco:runtime",
"//tensorflow/lite/toco:toco_port",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/memory",
"@flatbuffers",
],
)
tf_cc_test(
name = "operator_test",
srcs = [
"operator_test.cc",
],
deps = [
":operator",
"//tensorflow/core:ops",
"//tensorflow/core:protos_all_cc",
"//tensorflow/lite/schema:schema_fbs",
"//tensorflow/lite/toco:model",
"//tensorflow/lite/toco:runtime",
"//tensorflow/lite/toco:tooling_util",
"@com_google_absl//absl/log:check",
"@com_google_googletest//:gtest_main",
"@flatbuffers",
],
)
cc_library(
name = "types",
srcs = [
"types.cc",
],
hdrs = [
"types.h",
],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/lite:string_util",
"//tensorflow/lite/schema:schema_fbs",
"//tensorflow/lite/toco:model",
"//tensorflow/lite/toco:runtime",
"@com_google_absl//absl/log",
"@flatbuffers//:runtime_cc",
],
)
tf_cc_test(
name = "types_test",
srcs = [
"types_test.cc",
],
deps = [
":types",
"//tensorflow/core:ops",
"//tensorflow/lite/schema:schema_fbs",
"//tensorflow/lite/toco:model",
"//tensorflow/lite/toco:runtime",
"@com_google_googletest//:gtest_main",
"@flatbuffers//:runtime_cc",
],
)
cc_library(
name = "export",
srcs = [
"export.cc",
],
hdrs = [
"export.h",
],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:public"],
deps = [
":operator",
":types",
"//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy:quantize_weights",
"//tensorflow/compiler/mlir/lite/schema:schema_conversion_utils",
"//tensorflow/compiler/mlir/lite/tools/versioning",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core/platform:status",
"//tensorflow/lite:schema_fbs_version",
"//tensorflow/lite:util",
"//tensorflow/lite/c:c_api_types",
"//tensorflow/lite/schema:schema_fbs",
"//tensorflow/lite/toco:model",
"//tensorflow/lite/toco:toco_port",
"//tensorflow/lite/toco:tooling_util",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@flatbuffers",
],
)
tf_cc_test(
name = "export_test",
srcs = [
"export_test.cc",
],
deps = [
":export",
":operator",
":types",
"//tensorflow/compiler/mlir/lite/schema:schema_utils",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:ops",
"//tensorflow/core:protos_all_cc",
"//tensorflow/lite/schema:schema_fbs",
"//tensorflow/lite/toco:model",
"@com_google_absl//absl/log",
"@com_google_googletest//:gtest_main",
"@flatbuffers",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
cc_library(
name = "import",
srcs = [
"import.cc",
],
hdrs = [
"import.h",
],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:public"],
deps = [
":operator",
":types",
"//tensorflow/compiler/mlir/lite/schema:schema_utils",
"//tensorflow/lite:framework",
"//tensorflow/lite/core:framework",
"//tensorflow/lite/core/tools:verifier",
"//tensorflow/lite/schema:schema_fbs",
"//tensorflow/lite/toco:model",
"//tensorflow/lite/toco:model_flags_proto_cc",
"//tensorflow/lite/toco:tooling_util",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/strings:string_view",
"@flatbuffers",
],
)
tf_cc_test(
name = "import_test",
srcs = [
"import_test.cc",
],
tags = [
"no_oss", # TODO(b/273558651): Enable after updating flatbuffer version.
],
deps = [
":import",
"//tensorflow/compiler/mlir/lite/schema:schema_conversion_utils",
"//tensorflow/core:ops",
"//tensorflow/lite:schema_fbs_version",
"//tensorflow/lite/schema:schema_fbs",
"//tensorflow/lite/toco:model",
"//tensorflow/lite/toco:model_flags_proto_cc",
"//tensorflow/lite/toco:toco_port",
"@com_google_googletest//:gtest_main",
"@flatbuffers",
],
)
tensorflow/lite/toco/tflite/builtin_operator.h
+88
View File
@@ -0,0 +1,88 @@
/* Copyright 2017 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_TOCO_TFLITE_BUILTIN_OPERATOR_H_
#define TENSORFLOW_LITE_TOCO_TFLITE_BUILTIN_OPERATOR_H_
#include <memory>
#include "absl/memory/memory.h"
#include "tensorflow/compiler/mlir/lite/tools/versioning/op_version.h"
#include "tensorflow/lite/toco/tflite/operator.h"
namespace toco {
namespace tflite {
// Builtin operators have special TF Lite objects describing their options.
// This class has the boilerplate code for creating those.
//
// Template arguments:
// - T1 must derive from ::toco::Operator.
// - T2 must be one of TF Lite's objects defining Builtin Options, such as
// ::tflite::Conv2DOptions.
template <typename T1, typename T2, ::tflite::BuiltinOptions TfLiteEnum>
class BuiltinOperator : public BaseOperator {
public:
using TocoOperator = T1;
using TfLiteOptions = T2;
BuiltinOperator(::tflite::BuiltinOperator op, OperatorType type)
: BaseOperator(::tflite::EnumNameBuiltinOperator(op), type),
builtin_op_(op) {}
// Build the configuration object in the given flatbuffer builder. Return
// its offset.
virtual flatbuffers::Offset<TfLiteOptions> WriteOptions(
const TocoOperator& op,
flatbuffers::FlatBufferBuilder* builder) const = 0;
// Read options from the TF Lite object and set the corresponding values in
// the tf.mini operator.
virtual void ReadOptions(const TfLiteOptions& opt,
TocoOperator* op) const = 0;
Options Serialize(const Operator& op,
flatbuffers::FlatBufferBuilder* builder) const override {
auto options = WriteOptions(static_cast<const TocoOperator&>(op), builder);
return Options::Builtin(TfLiteEnum, options.Union());
}
std::unique_ptr<Operator> Deserialize(
const BuiltinOptions* builtin_options,
const CustomOptions* custom_options) const override {
auto op = std::make_unique<TocoOperator>();
auto* options = static_cast<const TfLiteOptions*>(builtin_options);
if (options) {
ReadOptions(*options, op.get());
}
return std::unique_ptr<Operator>(op.release());
}
int GetVersion(const OperatorSignature& op_signature) const override {
return ::tflite::GetBuiltinOperatorVersion(
GetVersioningOpSig(builtin_op_, op_signature));
}
::tflite::BuiltinOperator builtin_op() const { return builtin_op_; }
private:
const ::tflite::BuiltinOperator builtin_op_;
};
} // namespace tflite
} // namespace toco
#endif // TENSORFLOW_LITE_TOCO_TFLITE_BUILTIN_OPERATOR_H_
tensorflow/lite/toco/tflite/custom_operator.h
+76
View File
@@ -0,0 +1,76 @@
/* Copyright 2017 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_TOCO_TFLITE_CUSTOM_OPERATOR_H_
#define TENSORFLOW_LITE_TOCO_TFLITE_CUSTOM_OPERATOR_H_
#include <memory>
#include "flatbuffers/flexbuffers.h"
#include "absl/memory/memory.h"
#include "tensorflow/lite/toco/tflite/operator.h"
namespace toco {
namespace tflite {
// Custom operators have a generic byte buffer describing their options. This
// class provides the boilerplate code for populating those options using
// flexbuffers. Note that most of toco's operators will likely be supported
// as builtin operators in TF Lite.
//
// Template argument T must derive from ::toco::Operator.
template <typename T>
class CustomOperator : public BaseOperator {
public:
using TocoOperator = T;
using BaseOperator::BaseOperator;
// Populate the given flexbuffer with options obtained from the tf.mini
// operator.
virtual void WriteOptions(const TocoOperator& op,
flexbuffers::Builder* fbb) const {}
// Set options in the given tf.mini operator using values from the flexbuffer
// map.
virtual void ReadOptions(const flexbuffers::Map& m, TocoOperator* op) const {}
Options Serialize(const Operator& op,
flatbuffers::FlatBufferBuilder* builder) const override {
flexbuffers::Builder fbb;
fbb.Map(
[&]() { WriteOptions(static_cast<const TocoOperator&>(op), &fbb); });
fbb.Finish();
return Options::Custom(builder->CreateVector(fbb.GetBuffer()));
}
std::unique_ptr<Operator> Deserialize(
const BuiltinOptions* builtin_options,
const CustomOptions* custom_options) const override {
auto op = std::make_unique<TocoOperator>();
if (custom_options) {
auto flexbuffer_map =
flexbuffers::GetRoot(custom_options->data(), custom_options->size())
.AsMap();
ReadOptions(flexbuffer_map, op.get());
}
return std::unique_ptr<Operator>(op.release());
}
};
} // namespace tflite
} // namespace toco
#endif // TENSORFLOW_LITE_TOCO_TFLITE_CUSTOM_OPERATOR_H_
tensorflow/lite/toco/tflite/export.cc
+703
View File
@@ -0,0 +1,703 @@
/* 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 <cstdint>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <vector>
#include "flatbuffers/flexbuffers.h"
#include "absl/log/log.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "flatbuffers/string.h" // from @flatbuffers
#include "flatbuffers/vector.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantize_weights.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_conversion_utils.h"
#include "tensorflow/compiler/mlir/lite/tools/versioning/runtime_version.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/toco/model.h"
#include "tensorflow/lite/toco/tflite/operator.h"
#include "tensorflow/lite/toco/tflite/types.h"
#include "tensorflow/lite/toco/toco_types.h"
#include "tensorflow/lite/toco/tooling_util.h"
#include "tensorflow/lite/util.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 std::string& tensorflow_op) {
// Technically this is equivalent 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;
}
return false;
}
// Check if a TensorFlow Op is unsupported by the Flex runtime.
bool IsUnsupportedFlexOp(const std::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<std::string, BuiltinOperator>& GetBuiltinOpsMap() {
static std::map<std::string, BuiltinOperator>* builtin_ops = nullptr;
if (builtin_ops == nullptr) {
builtin_ops = new std::map<std::string, BuiltinOperator>();
for (int i = BuiltinOperator_MIN; i <= BuiltinOperator_MAX; ++i) {
BuiltinOperator op = static_cast<BuiltinOperator>(i);
std::string name = EnumNameBuiltinOperator(op);
if (op != BuiltinOperator_CUSTOM && !name.empty()) {
(*builtin_ops)[name] = op;
}
}
}
return *builtin_ops;
}
void WriteModelToString(const flatbuffers::FlatBufferBuilder& builder,
std::string* file_contents) {
const uint8_t* buffer = builder.GetBufferPointer();
int size = builder.GetSize();
*file_contents = std::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;
std::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_ =
std::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_ =
std::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<std::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 std::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 std::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<std::string, BuiltinOperator> builtin_ops;
for (int i = BuiltinOperator_MIN; i <= BuiltinOperator_MAX; ++i) {
BuiltinOperator op = static_cast<BuiltinOperator>(i);
std::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 std::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 std::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 (uint32_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);
}
absl::Status Export(const Model& model, std::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<std::string>* custom_ops,
std::vector<std::string>* error_msgs) {
std::set<std::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 appropriate model metadata.
// This consists of a 16-byte placeholder string buffer that will be populated
// with the model's minimum required runtime version after the model is
// finalized. We choose 16 bytes as it's the alignment of flatbuffer buffers
// and prevents wasted space if the final string is shorter than 16 bytes.
// Note: This logic matches that used in the MLIR flatbuffer export path.
flatbuffers::Offset<tflite::Metadata> ExportMetadata(
const Model& model, std::vector<Offset<Vector<uint8_t>>>* buffers_to_write,
FlatBufferBuilder* builder) {
auto metadata =
CreateMetadata(*builder, builder->CreateString("min_runtime_version"),
buffers_to_write->size());
const std::string min_runtime_placeholder = std::string(16, '\0');
buffers_to_write->push_back(builder->CreateVector(
reinterpret_cast<const uint8_t*>(min_runtime_placeholder.data()),
min_runtime_placeholder.size()));
return metadata;
}
absl::Status Export(
const Model& model, std::string* output_file_contents,
const ExportParams& params,
const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type) {
for (const std::string& input_array : model.GetInvalidInputArrays()) {
if (model.HasArray(input_array)) {
return absl::InvalidArgumentError(absl::StrCat("Placeholder ",
input_array,
" should be specified 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<std::string> builtin_ops;
// The set of custom ops (not including Flex ops).
std::set<std::string> custom_ops;
// The set of Flex ops which are not supported.
std::set<std::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<std::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<std::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 absl::InvalidArgumentError(absl::StrCat(
please_report_bug_message(), absl::StrJoin(error_msgs, " ")));
}
}
}
if (!unsupported_flex_ops.empty()) {
return absl::InvalidArgumentError(
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 absl::UnimplementedError(
"Unsupported flag: allow_dynamic_tensors. Offline memory planning is "
"not implemented yet.");
}
auto metadata = ExportMetadata(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);
::tflite::UpdateMinimumRuntimeVersionForModel(builder.GetBufferPointer());
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);
::mlir::lite::toco_legacy::BufferType quantized_type;
if (params.quantize_weights == QuantizedBufferType::INT8) {
quantized_type = ::mlir::lite::toco_legacy::BufferType::QUANTIZED_INT8;
} else if (params.quantize_weights == QuantizedBufferType::FLOAT16) {
quantized_type = ::mlir::lite::toco_legacy::BufferType::QUANTIZED_FLOAT16;
} else {
return absl::InvalidArgumentError("Quantized type not recognized");
}
if (!::mlir::lite::toco_legacy::QuantizeWeights(
&q_builder, input_model, quantized_type,
!params.disable_per_channel,
::mlir::lite::toco_legacy::QuantizerType::OLD_QUANTIZER)
.ok()) {
return absl::InvalidArgumentError(
"Quantize weights transformation failed.");
}
WriteModelToString(q_builder, output_file_contents);
}
return absl::Status();
}
} // namespace tflite
} // namespace toco
tensorflow/lite/toco/tflite/export.h
+183
View File
@@ -0,0 +1,183 @@
/* Copyright 2017 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_TOCO_TFLITE_EXPORT_H_
#define TENSORFLOW_LITE_TOCO_TFLITE_EXPORT_H_
#include <string>
#include "absl/log/log.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/toco/model.h"
#include "tensorflow/lite/toco/tflite/operator.h"
#include "tensorflow/lite/util.h"
namespace toco {
namespace tflite {
enum class QuantizedBufferType { NONE, INT8, FLOAT16 };
// The parameters for exporting a TFLite model.
struct ExportParams {
bool allow_custom_ops = false;
bool allow_dynamic_tensors = true;
bool enable_select_tf_ops = false;
QuantizedBufferType quantize_weights = QuantizedBufferType::NONE;
// Whether to use per-tensor (false) or per-channel (true) for hybrid quant.
bool disable_per_channel = false;
};
// Transform the given tf.mini model into a TF Lite flatbuffer and deposit the
// result in the given string.
absl::Status Export(const Model& model, std::string* output_file_contents,
const ExportParams& params);
// Export API with custom TFLite operator mapping.
absl::Status Export(
const Model& model, std::string* output_file_contents,
const ExportParams& params,
const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type);
// This is for backward-compatibility.
inline void Export(const Model& model, bool allow_custom_ops,
bool quantize_weights, std::string* output_file_contents) {
ExportParams params;
params.allow_custom_ops = allow_custom_ops;
params.quantize_weights =
quantize_weights ? QuantizedBufferType::INT8 : QuantizedBufferType::NONE;
auto status = Export(model, output_file_contents, params);
if (!status.ok()) LOG(QFATAL) << status.message();
}
// This is for backward-compatibility.
inline void Export(
const Model& model, bool allow_custom_ops, bool quantize_weights,
std::string* output_file_contents,
const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type) {
ExportParams params;
params.allow_custom_ops = allow_custom_ops;
params.quantize_weights =
quantize_weights ? QuantizedBufferType::INT8 : QuantizedBufferType::NONE;
auto status = Export(model, output_file_contents, params, ops_by_type);
if (!status.ok()) LOG(QFATAL) << status.message();
}
// This is for backward-compatibility.
inline void Export(const Model& model, std::string* output_file_contents) {
ExportParams params;
params.allow_custom_ops = true;
auto status = Export(model, output_file_contents, params);
if (!status.ok()) LOG(QFATAL) << status.message();
}
namespace details {
// A map from tensor name to its final position in the TF Lite buffer.
using TensorsMap = std::unordered_map<std::string, int>;
// A key to identify an operator.
// Only when `type` is `kUnsupported`, `custom_code` is filled to
// identify which operation is used.
class OperatorKey {
public:
OperatorKey() {}
// Construct OperatorKey by Toco op.
OperatorKey(
const ::toco::OperatorSignature& op_signature,
const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type,
bool enable_select_tf_ops);
// Construct OperatorKey by type, custom code and version.
// Note that this construct doesn't set the additional information including
// `is_custom_op`, `is_flex_op`, `is_unsupported_flex_op`.
OperatorKey(::tflite::BuiltinOperator type, const std::string& custom_code,
int version)
: type_(type), custom_code_(custom_code), version_(version) {}
// Only `type`, `custom_code` and `version` is used to compute hash and
// identity.
::tflite::BuiltinOperator type() const { return type_; }
const std::string& custom_code() const { return custom_code_; }
int version() const { return version_; }
// The attributes below are not used to compute hash and identity.
//
// Return true if the op is a custom op. Note it will return false for Flex
// ops.
bool is_custom_op() const { return is_custom_op_; }
// Return true if the op is a Flex op.
bool is_flex_op() const { return is_flex_op_; }
// Return true if the op is a Flex op but it's knwon that the op is not
// supported by Flex runtime.
bool is_unsupported_flex_op() const { return is_unsupported_flex_op_; }
// Return the original TensorFlow op name for a Flex op.
const std::string& flex_tensorflow_op() const { return flex_tensorflow_op_; }
bool operator<(const OperatorKey& other) const {
if (type_ < other.type_)
return true;
else if (type_ > other.type_)
return false;
else if (custom_code_ < other.custom_code_)
return true;
else if (custom_code_ > other.custom_code_)
return false;
else
return version_ < other.version_;
}
bool operator==(const OperatorKey& other) const {
return type_ == other.type_ && custom_code_ == other.custom_code_ &&
version_ == other.version_;
}
struct Hash {
size_t operator()(const OperatorKey& key) const {
return ::tflite::CombineHashes(
{std::hash<size_t>()(static_cast<size_t>(key.type())),
std::hash<std::string>()(key.custom_code()),
std::hash<int>()(key.version())});
}
};
private:
::tflite::BuiltinOperator type_ = ::tflite::BuiltinOperator_CUSTOM;
std::string custom_code_;
int version_ = 1;
bool is_custom_op_ = false;
bool is_flex_op_ = false;
bool is_unsupported_flex_op_ = false;
// The original TensorFlow op name for the flex op. Filled only when
// `is_flex_op` is true.
std::string flex_tensorflow_op_;
};
// A map from OperatorKey to its final position in the TF Lite buffer.
using OperatorsMap = std::unordered_map<OperatorKey, int, OperatorKey::Hash>;
void LoadTensorsMap(const Model& model, TensorsMap* tensors_map);
void LoadOperatorsMap(
const Model& model, OperatorsMap* operators_map,
const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type,
bool enable_select_tf_ops);
} // namespace details
} // namespace tflite
} // namespace toco
#endif // TENSORFLOW_LITE_TOCO_TFLITE_EXPORT_H_
tensorflow/lite/toco/tflite/export_test.cc
+850
View File
@@ -0,0 +1,850 @@
/* Copyright 2017 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 <algorithm>
#include <cstdint>
#include <initializer_list>
#include <map>
#include <memory>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/log/log.h"
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/schema/schema_utils.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/toco/model.h"
#include "tensorflow/lite/toco/tflite/builtin_operator.h"
#include "tensorflow/lite/toco/tflite/operator.h"
#include "tensorflow/lite/toco/tflite/types.h"
namespace toco {
namespace tflite {
namespace {
using ::testing::ElementsAre;
using ::testing::HasSubstr;
class ExportTest : public ::testing::Test {
protected:
void ResetOperators() { input_model_.operators.clear(); }
void AddTensorsByName(std::initializer_list<std::string> names) {
for (const std::string& name : names) {
input_model_.GetOrCreateArray(name);
}
}
void AddOperatorsByName(std::initializer_list<std::string> names) {
for (const std::string& name : names) {
if (name == "Conv") {
auto* op = new ConvOperator;
op->padding.type = PaddingType::kSame;
op->inputs = {"input", "filter"};
op->outputs = {"output"};
Array& input_array = input_model_.GetOrCreateArray(op->inputs[0]);
Array& filter_array = input_model_.GetOrCreateArray(op->inputs[1]);
Array& output_array = input_model_.GetOrCreateArray(op->outputs[0]);
input_array.data_type = ArrayDataType::kFloat;
filter_array.data_type = ArrayDataType::kFloat;
output_array.data_type = ArrayDataType::kFloat;
input_model_.operators.emplace_back(op);
} else if (name == "Add") {
auto* op = new AddOperator;
op->inputs = {"input1", "input2"};
op->outputs = {"output"};
Array& input1_array = input_model_.GetOrCreateArray(op->inputs[0]);
Array& input2_array = input_model_.GetOrCreateArray(op->inputs[1]);
Array& output_array = input_model_.GetOrCreateArray(op->outputs[0]);
input1_array.data_type = ArrayDataType::kFloat;
input2_array.data_type = ArrayDataType::kFloat;
output_array.data_type = ArrayDataType::kFloat;
input_model_.operators.emplace_back(op);
} else if (name == "Sub") {
auto* op = new SubOperator;
op->inputs = {"input1", "input2"};
op->outputs = {"output"};
Array& input1_array = input_model_.GetOrCreateArray(op->inputs[0]);
Array& input2_array = input_model_.GetOrCreateArray(op->inputs[1]);
Array& output_array = input_model_.GetOrCreateArray(op->outputs[0]);
input1_array.data_type = ArrayDataType::kFloat;
input2_array.data_type = ArrayDataType::kFloat;
output_array.data_type = ArrayDataType::kFloat;
input1_array.copy_shape({1, 2, 2, 2});
input2_array.copy_shape({1, 2, 2, 2});
output_array.copy_shape({1, 2, 2, 2});
input_model_.operators.emplace_back(op);
} else if (name == "Assert") {
auto* op = new TensorFlowAssertOperator;
// Even though assert is known to TOCO, it doesn't have a tflite
// serializer, so it has to be exported as a custom op. If we attach a
// NodeDef to it, however, it will be exported as a flex op instead.
::tensorflow::NodeDef node_def;
node_def.set_name("Assert");
node_def.set_op("Assert");
node_def.SerializeToString(&op->tensorflow_node_def);
input_model_.operators.emplace_back(op);
} else {
auto* op = new TensorFlowUnsupportedOperator;
op->tensorflow_op = name;
input_model_.operators.emplace_back(op);
}
}
}
void BuildQuantizableTestModel() {
input_model_.GetOrCreateArray("inputs");
Array& weight_array = input_model_.GetOrCreateArray("weights");
// Make the buffer large enough for QuantizeWeights transformation to take
// effect.
int buf_size = 1296;
auto weight_buf = std::make_unique<float[]>(buf_size);
for (int i = 0; i < buf_size; i++) {
// Fill the array with some garbage values.
weight_buf[i] = static_cast<float>(i % 128);
}
weight_array.data_type = ArrayDataType::kFloat;
// Initialize shape for the input array.
Shape* weight_array_shape = weight_array.mutable_shape();
std::vector<int>* weight_array_shape_dim =
weight_array_shape->mutable_dims();
weight_array_shape_dim->resize(4, 6);
auto& weight_array_buffer =
weight_array.GetMutableBuffer<ArrayDataType::kFloat>();
weight_array_buffer.data.resize(buf_size);
float* buf_ptr =
weight_array.GetMutableBuffer<ArrayDataType::kFloat>().data.data();
std::copy(weight_buf.get(), weight_buf.get() + buf_size, buf_ptr);
{
auto* op = new ConvOperator;
op->padding.type = PaddingType::kSame;
op->inputs = {"inputs", "weights"};
op->outputs = {"output"};
Array& input_array = input_model_.GetArray(op->inputs[0]);
Array& filter_array = input_model_.GetArray(op->inputs[1]);
Array& output_array = input_model_.GetOrCreateArray(op->outputs[0]);
input_array.data_type = ArrayDataType::kFloat;
filter_array.data_type = ArrayDataType::kFloat;
output_array.data_type = ArrayDataType::kFloat;
input_model_.operators.emplace_back(op);
}
{
auto* op = new AddOperator;
op->inputs = {"input1", "input2"};
op->outputs = {"output"};
Array& input1_array = input_model_.GetOrCreateArray(op->inputs[0]);
Array& input2_array = input_model_.GetOrCreateArray(op->inputs[1]);
Array& output_array = input_model_.GetOrCreateArray(op->outputs[0]);
input1_array.data_type = ArrayDataType::kFloat;
input2_array.data_type = ArrayDataType::kFloat;
output_array.data_type = ArrayDataType::kFloat;
input_model_.operators.emplace_back(op);
}
}
absl::Status ExportAndReturnStatus(const ExportParams& params) {
std::string result;
return Export(input_model_, &result, params);
}
std::vector<std::string> ExportAndSummarizeOperators(
const ExportParams& params) {
std::vector<std::string> names;
std::string result;
auto status = Export(input_model_, &result, params);
if (!status.ok()) {
LOG(INFO) << status.message();
return names;
}
auto* model = ::tflite::GetModel(result.data());
for (const ::tflite::OperatorCode* opcode : *model->operator_codes()) {
auto builtin_code = GetBuiltinCode(opcode);
if (builtin_code != ::tflite::BuiltinOperator_CUSTOM) {
names.push_back(std::string("builtin:") +
::tflite::EnumNameBuiltinOperator(builtin_code));
} else {
names.push_back(std::string("custom:") +
opcode->custom_code()->c_str());
}
}
return names;
}
std::vector<uint32_t> ExportAndGetOperatorIndices(
const ExportParams& params) {
std::vector<uint32_t> indices;
std::string result;
if (!Export(input_model_, &result, params).ok()) return indices;
auto* model = ::tflite::GetModel(result.data());
auto operators = (*model->subgraphs())[0]->operators();
for (const auto* op : *operators) {
indices.push_back(op->opcode_index());
}
return indices;
}
Model input_model_;
};
TEST_F(ExportTest, LoadTensorsMap) {
AddTensorsByName({"tensor_one", "tensor_two"});
details::TensorsMap tensors;
details::LoadTensorsMap(input_model_, &tensors);
EXPECT_EQ(0, tensors["tensor_one"]);
EXPECT_EQ(1, tensors["tensor_two"]);
}
TEST_F(ExportTest, LoadOperatorsMap) {
AddOperatorsByName({"Conv", "Add", "MyCrazyOp", "Sub"});
details::OperatorsMap operators;
const auto ops_by_type = BuildOperatorByTypeMap();
details::LoadOperatorsMap(input_model_, &operators, ops_by_type, false);
EXPECT_EQ(
0, operators[details::OperatorKey(::tflite::BuiltinOperator_ADD, "", 1)]);
EXPECT_EQ(1, operators[details::OperatorKey(::tflite::BuiltinOperator_CONV_2D,
"", 1)]);
EXPECT_EQ(2, operators[details::OperatorKey(::tflite::BuiltinOperator_CUSTOM,
"MyCrazyOp", 1)]);
EXPECT_EQ(
3, operators[details::OperatorKey(::tflite::BuiltinOperator_SUB, "", 1)]);
}
TEST_F(ExportTest, UnsupportedFunctionality) {
AddOperatorsByName({"Conv"});
ExportParams params;
params.allow_dynamic_tensors = false;
auto status = ExportAndReturnStatus(params);
EXPECT_EQ(status.code(), ::tensorflow::error::UNIMPLEMENTED);
EXPECT_THAT(status.message(),
HasSubstr("Unsupported flag: allow_dynamic_tensors."));
}
TEST_F(ExportTest, Export) {
AddOperatorsByName({"Conv", "Add", "MyCrazyOp", "Sub"});
ExportParams params;
params.allow_custom_ops = true;
params.enable_select_tf_ops = false;
params.quantize_weights = QuantizedBufferType::NONE;
EXPECT_THAT(ExportAndSummarizeOperators(params),
ElementsAre("builtin:ADD", "builtin:CONV_2D", "custom:MyCrazyOp",
"builtin:SUB"));
EXPECT_THAT(ExportAndGetOperatorIndices(params), ElementsAre(1, 0, 2, 3));
}
TEST_F(ExportTest, ExportMinRuntime) {
AddOperatorsByName({"Conv", "Add", "Sub"});
ExportParams params;
params.allow_custom_ops = true;
params.enable_select_tf_ops = false;
params.quantize_weights = QuantizedBufferType::NONE;
std::string output;
auto status = Export(input_model_, &output, params);
auto* model = ::tflite::GetModel(output.data());
EXPECT_EQ(model->metadata()->size(), 1);
EXPECT_EQ(model->metadata()->Get(0)->name()->str(), "min_runtime_version");
auto buf = model->metadata()->Get(0)->buffer();
auto* buffer = (*model->buffers())[buf];
auto* array = buffer->data();
EXPECT_EQ(reinterpret_cast<const char*>(array->data()), std::string("1.6.0"));
}
TEST_F(ExportTest, ExportEmptyMinRuntime) {
AddOperatorsByName({"Switch", "MyCustomOp", "Assert"});
ExportParams params;
params.allow_custom_ops = true;
std::string output;
auto status = Export(input_model_, &output, params);
auto* model = ::tflite::GetModel(output.data());
EXPECT_EQ(model->metadata()->size(), 1);
EXPECT_EQ(model->metadata()->Get(0)->name()->str(), "min_runtime_version");
auto buf = model->metadata()->Get(0)->buffer();
auto* buffer = (*model->buffers())[buf];
auto* array = buffer->data();
EXPECT_EQ(reinterpret_cast<const char*>(array->data()), std::string(""));
}
TEST_F(ExportTest, UnsupportedControlFlowErrors) {
AddOperatorsByName({"Conv", "Add", "Switch", "Merge"});
ExportParams params;
params.allow_custom_ops = false;
// The model contains control flow ops which are not convertible, so we should
// check the returned error message.
std::string output;
const auto ops_by_type = BuildOperatorByTypeMap();
auto status = Export(input_model_, &output, params, ops_by_type);
EXPECT_EQ(status.message(),
"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\nTensorFlow Lite currently doesn't support control "
"flow ops: Merge, Switch. We are working on supporting control "
"flow ops, please see github issue at "
"https://github.com/tensorflow/tensorflow/issues/28485.");
}
TEST_F(ExportTest, UnsupportedOpsAndNeedEnableFlex) {
AddOperatorsByName({"Conv", "Add", "BatchNormWithGlobalNormalization"});
ExportParams params;
params.allow_custom_ops = false;
params.enable_select_tf_ops = false;
std::string output;
const auto ops_by_type = BuildOperatorByTypeMap();
auto status = Export(input_model_, &output, params, ops_by_type);
EXPECT_EQ(
status.message(),
"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\nSome 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: ADD, CONV_2D. Here is a list of operators for which you will "
"need custom implementations: BatchNormWithGlobalNormalization.");
}
TEST_F(ExportTest, UnsupportedOpsNeedCustomImplementation) {
AddOperatorsByName({"Conv", "Add", "MyCustomOp1", "MyCustomOp2"});
ExportParams params;
params.allow_custom_ops = false;
params.enable_select_tf_ops = true;
std::string output;
const auto ops_by_type = BuildOperatorByTypeMap();
auto status = Export(input_model_, &output, params, ops_by_type);
EXPECT_EQ(
status.message(),
"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\nSome 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: ADD, CONV_2D. Here is a list "
"of operators for which you will need custom implementations: "
"MyCustomOp1, MyCustomOp2.");
}
TEST_F(ExportTest, UnsupportedControlFlowAndCustomOpsErrors) {
AddOperatorsByName(
{"Conv", "Add", "Switch", "Merge", "MyCustomOp1", "MyCustomOp2"});
ExportParams params;
params.allow_custom_ops = false;
// The model contains control flow ops which are not convertible, so we should
// check the returned error message.
std::string output;
const auto ops_by_type = BuildOperatorByTypeMap();
auto status = Export(input_model_, &output, params, ops_by_type);
EXPECT_EQ(
status.message(),
"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\nTensorFlow Lite currently doesn't support control flow "
"ops: Merge, Switch. We are working on supporting control flow ops, "
"please see github issue at "
"https://github.com/tensorflow/tensorflow/issues/28485. 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: ADD, CONV_2D. Here is a list of operators for which you will "
"need custom implementations: MyCustomOp1, MyCustomOp2.");
}
TEST_F(ExportTest, QuantizeWeights) {
// Sanity check for quantize_weights parameter.
BuildQuantizableTestModel();
std::string unquantized_result;
Export(input_model_, true, /*quantize_weights*/ false, &unquantized_result);
BuildQuantizableTestModel();
std::string quantized_result;
Export(input_model_, true, /*quantize_weights*/ true, &quantized_result);
// The quantized models should be smaller.
EXPECT_LT(quantized_result.size(), unquantized_result.size());
}
class OpSetsTest : public ExportTest {
public:
enum OpSet { kTfLiteBuiltins, kSelectTfOps, kCustomOps };
void SetAllowedOpSets(std::initializer_list<OpSet> sets) {
import_all_ops_as_unsupported_ = true;
params_.allow_custom_ops = false;
params_.enable_select_tf_ops = false;
params_.quantize_weights = QuantizedBufferType::NONE;
for (const OpSet& i : sets) {
switch (i) {
case kTfLiteBuiltins:
import_all_ops_as_unsupported_ = false;
break;
case kSelectTfOps:
params_.enable_select_tf_ops = true;
break;
case kCustomOps:
params_.allow_custom_ops = true;
break;
}
}
}
std::vector<std::string> ImportExport(
std::initializer_list<std::string> op_names) {
ResetOperators();
if (!import_all_ops_as_unsupported_) {
AddOperatorsByName(op_names);
} else {
for (const std::string& name : op_names) {
auto* op = new TensorFlowUnsupportedOperator;
op->tensorflow_op = name;
input_model_.operators.emplace_back(op);
}
}
return ExportAndSummarizeOperators(params_);
}
private:
bool import_all_ops_as_unsupported_;
ExportParams params_;
};
TEST_F(OpSetsTest, BuiltinsOnly) {
// --target_op_set=TFLITE_BUILTINS
SetAllowedOpSets({kTfLiteBuiltins});
EXPECT_THAT(ImportExport({"Add", "AdjustHue", "UnrollAndFold", "Assert"}),
ElementsAre());
EXPECT_THAT(ImportExport({"Add"}), ElementsAre("builtin:ADD"));
// --target_op_set=TFLITE_BUILTINS --allow_custom_ops
SetAllowedOpSets({kTfLiteBuiltins, kCustomOps});
EXPECT_THAT(ImportExport({"Add", "AdjustHue", "UnrollAndFold", "Assert"}),
ElementsAre("builtin:ADD", "custom:AdjustHue", "custom:Assert",
"custom:UnrollAndFold"));
}
TEST_F(OpSetsTest, TfSelectOnly) {
// --target_op_set=SELECT_TF_OPS
SetAllowedOpSets({kSelectTfOps});
EXPECT_THAT(ImportExport({"Add", "AdjustHue", "RandomUniform",
"UnrollAndFold", "Assert"}),
ElementsAre());
EXPECT_THAT(ImportExport({"Add"}), ElementsAre("custom:FlexAdd"));
// --target_op_set=SELECT_TF_OPS --allow_custom_ops
SetAllowedOpSets({kSelectTfOps, kCustomOps});
EXPECT_THAT(
ImportExport(
{"Add", "AdjustHue", "RandomUniform", "UnrollAndFold", "Assert"}),
ElementsAre("custom:FlexAdd", "custom:FlexAdjustHue", "custom:FlexAssert",
"custom:FlexRandomUniform", "custom:UnrollAndFold"));
}
TEST_F(OpSetsTest, BuiltinsAndTfSelect) {
// --target_op_set=TFLITE_BUILTINS,SELECT_TF_OPS
SetAllowedOpSets({kTfLiteBuiltins, kSelectTfOps});
EXPECT_THAT(ImportExport({"Add", "AdjustHue", "UnrollAndFold", "Assert"}),
ElementsAre());
EXPECT_THAT(ImportExport({"Add", "RandomUniform"}),
ElementsAre("builtin:ADD", "custom:FlexRandomUniform"));
// --target_op_set=TFLITE_BUILTINS,SELECT_TF_OPS --allow_custom_ops
SetAllowedOpSets({kTfLiteBuiltins, kSelectTfOps, kCustomOps});
EXPECT_THAT(
ImportExport(
{"Add", "AdjustHue", "RandomUniform", "UnrollAndFold", "Assert"}),
ElementsAre("builtin:ADD", "custom:FlexAdjustHue", "custom:FlexAssert",
"custom:FlexRandomUniform", "custom:UnrollAndFold"));
}
// This test is based on a hypothetical scenario that dilation is supported
// only in Conv version 2. So Toco populates version=1 when dilation parameters
// are all 1, and version=2 otherwise.
class FakeConvolutionOperator
: public BuiltinOperator<ConvOperator, ::tflite::Conv2DOptions,
::tflite::BuiltinOptions_Conv2DOptions> {
public:
FakeConvolutionOperator()
: BuiltinOperator(::tflite::BuiltinOperator_CONV_2D,
OperatorType::kConv) {}
// Returning the op version according to the op parameters.
int GetVersion(const OperatorSignature& op_signature) const override {
const TocoOperator& conv_op =
static_cast<const TocoOperator&>(*op_signature.op);
if (conv_op.dilation_width_factor != 1 ||
conv_op.dilation_height_factor != 1) {
// Version 2 if dilation is used.
return 2;
}
return 1;
}
// Note: The read / write code doesn't need to be changed if we stick with
// the restrictions:
// * Only adding parameters at the bottom of the Flatbuffer tables.
// * When the default value of parameters are used, the op works consistently
// with the previous version.
flatbuffers::Offset<TfLiteOptions> WriteOptions(
const TocoOperator& op,
flatbuffers::FlatBufferBuilder* builder) const override {
auto padding = Padding::Serialize(op.padding.type);
auto activation_function =
ActivationFunction::Serialize(op.fused_activation_function);
return ::tflite::CreateConv2DOptions(*builder, padding, op.stride_width,
op.stride_height, activation_function,
op.dilation_width_factor,
op.dilation_height_factor);
}
void ReadOptions(const TfLiteOptions& options,
TocoOperator* op) const override {
op->padding.type = Padding::Deserialize(options.padding());
op->stride_width = options.stride_w();
op->stride_height = options.stride_h();
op->dilation_width_factor = options.dilation_w_factor();
op->dilation_height_factor = options.dilation_h_factor();
op->fused_activation_function =
ActivationFunction::Deserialize(options.fused_activation_function());
}
};
class VersionedOpExportTest : public ::testing::Test {
protected:
void SetUp() override {
input_model_.GetOrCreateArray("input");
input_model_.GetOrCreateArray("filter");
input_model_.GetOrCreateArray("output");
}
void AddConvOp(bool use_dilation) {
{
auto* op = new ConvOperator;
op->inputs.push_back("input");
op->inputs.push_back("filter");
op->outputs.push_back("output");
op->padding.type = PaddingType::kSame;
op->stride_width = 1;
op->stride_height = 1;
if (use_dilation) {
op->dilation_width_factor = 2;
op->dilation_height_factor = 2;
} else {
op->dilation_width_factor = 1;
op->dilation_height_factor = 1;
}
input_model_.operators.emplace_back(op);
}
}
std::map<OperatorType, std::unique_ptr<BaseOperator>>
BuildFakeOperatorByTypeMap() {
std::map<OperatorType, std::unique_ptr<BaseOperator>> result;
result[OperatorType::kConv] =
std::unique_ptr<BaseOperator>(new FakeConvolutionOperator);
return result;
}
Model input_model_;
};
TEST_F(VersionedOpExportTest, LoadOperatorsMapWithOpV1) {
AddConvOp(false);
details::OperatorsMap operators;
const auto ops_by_type = BuildFakeOperatorByTypeMap();
details::LoadOperatorsMap(input_model_, &operators, ops_by_type, false);
EXPECT_EQ(1, operators.size());
EXPECT_EQ(0, operators.at(details::OperatorKey(
::tflite::BuiltinOperator_CONV_2D, "", 1)));
}
TEST_F(VersionedOpExportTest, LoadOperatorsMapWithOpV2) {
AddConvOp(true);
details::OperatorsMap operators;
const auto ops_by_type = BuildFakeOperatorByTypeMap();
details::LoadOperatorsMap(input_model_, &operators, ops_by_type, false);
EXPECT_EQ(1, operators.size());
EXPECT_EQ(0, operators.at(details::OperatorKey(
::tflite::BuiltinOperator_CONV_2D, "", 2)));
}
TEST_F(VersionedOpExportTest, LoadOperatorsMapWithBothVersions) {
AddConvOp(false);
AddConvOp(true);
details::OperatorsMap operators;
const auto ops_by_type = BuildFakeOperatorByTypeMap();
details::LoadOperatorsMap(input_model_, &operators, ops_by_type, false);
EXPECT_EQ(2, operators.size());
EXPECT_EQ(0, operators.at(details::OperatorKey(
::tflite::BuiltinOperator_CONV_2D, "", 1)));
EXPECT_EQ(1, operators.at(details::OperatorKey(
::tflite::BuiltinOperator_CONV_2D, "", 2)));
}
TEST_F(VersionedOpExportTest, Export) {
AddConvOp(false);
AddConvOp(true);
std::string result;
const auto ops_by_type = BuildFakeOperatorByTypeMap();
Export(input_model_, true, false, &result, ops_by_type);
auto* model = ::tflite::GetModel(result.data());
auto operator_codes = model->operator_codes();
// Verify that 2 operator codes are populated. Both are CONV_2D but with
// different versions.
EXPECT_EQ(2, operator_codes->size());
EXPECT_EQ(::tflite::BuiltinOperator_CONV_2D,
GetBuiltinCode((*operator_codes)[0]));
EXPECT_EQ(1, (*operator_codes)[0]->version());
EXPECT_EQ(::tflite::BuiltinOperator_CONV_2D,
GetBuiltinCode((*operator_codes)[1]));
EXPECT_EQ(2, (*operator_codes)[1]->version());
// Verify that the 2 operators points to the correct indices of the operation
// codes.
auto operators = (*model->subgraphs())[0]->operators();
EXPECT_EQ(2, operators->size());
EXPECT_EQ(0, (*operators)[0]->opcode_index());
EXPECT_EQ(1, (*operators)[1]->opcode_index());
}
TEST(OperatorKeyTest, TestBuiltinOp) {
Model model;
auto op = std::make_unique<ConvOperator>();
// Test a normal float operation.
op->inputs = {"input", "filter"};
op->outputs = {"output"};
Array& input_array = model.GetOrCreateArray(op->inputs[0]);
Array& filter_array = model.GetOrCreateArray(op->inputs[1]);
Array& output_array = model.GetOrCreateArray(op->outputs[0]);
input_array.data_type = ArrayDataType::kFloat;
filter_array.data_type = ArrayDataType::kFloat;
output_array.data_type = ArrayDataType::kFloat;
const auto ops_by_type = BuildOperatorByTypeMap();
const toco::OperatorSignature op_signature = {op.get(), &model};
const auto key = details::OperatorKey(op_signature, ops_by_type, false);
EXPECT_EQ(key.type(), ::tflite::BuiltinOperator_CONV_2D);
EXPECT_EQ(key.custom_code(), "");
EXPECT_EQ(key.version(), 1);
}
TEST(OperatorKeyTest, TestBuiltinOpWithVersionedInputTypes) {
Model model;
auto op = std::make_unique<DequantizeOperator>();
op->inputs = {"input"};
op->outputs = {"output"};
Array& input_array = model.GetOrCreateArray(op->inputs[0]);
Array& output_array = model.GetOrCreateArray(op->outputs[0]);
input_array.data_type = ArrayDataType::kInt8;
output_array.data_type = ArrayDataType::kFloat;
const auto ops_by_type = BuildOperatorByTypeMap();
// Test a signed int8 dequantize operation.
const toco::OperatorSignature op_signature = {op.get(), &model};
const auto key = details::OperatorKey(op_signature, ops_by_type, false);
EXPECT_EQ(key.type(), ::tflite::BuiltinOperator_DEQUANTIZE);
EXPECT_EQ(key.custom_code(), "");
EXPECT_EQ(key.version(), 2);
}
TEST(OperatorKeyTest, TestCustomOp) {
Model model;
auto op = std::make_unique<TensorFlowUnsupportedOperator>();
op->tensorflow_op = "MyCrazyCustomOp";
const auto ops_by_type = BuildOperatorByTypeMap();
const toco::OperatorSignature op_signature = {op.get(), &model};
const auto key = details::OperatorKey(op_signature, ops_by_type, false);
EXPECT_EQ(key.type(), ::tflite::BuiltinOperator_CUSTOM);
EXPECT_EQ(key.custom_code(), "MyCrazyCustomOp");
EXPECT_EQ(key.version(), 1);
}
TEST(OperatorKeyTest, TestFlexOp) {
Model model;
auto op = std::make_unique<TensorFlowUnsupportedOperator>();
op->tensorflow_op = "BatchMatMul";
const auto ops_by_type = BuildOperatorByTypeMap();
{
const toco::OperatorSignature op_signature = {op.get(), &model};
const auto key = details::OperatorKey(op_signature, ops_by_type, false);
// It shouldn't be converted to Flex op if `allow_flex_op` is false.
EXPECT_EQ(key.type(), ::tflite::BuiltinOperator_CUSTOM);
EXPECT_EQ(key.custom_code(), "BatchMatMul");
EXPECT_EQ(key.version(), 1);
EXPECT_TRUE(key.is_custom_op());
EXPECT_FALSE(key.is_flex_op());
}
{
// Verify that the custom op name is prefixed by "Flex" and `is_flex_op`
// is true.
const toco::OperatorSignature op_signature = {op.get(), &model};
const auto key = details::OperatorKey(op_signature, ops_by_type, true);
EXPECT_EQ(key.type(), ::tflite::BuiltinOperator_CUSTOM);
EXPECT_EQ(key.custom_code(), "FlexBatchMatMul");
EXPECT_EQ(key.version(), 1);
EXPECT_FALSE(key.is_custom_op());
EXPECT_TRUE(key.is_flex_op());
}
}
TEST(OperatorKeyTest, TestFlexWithControlFlowOp) {
Model model;
auto op = std::make_unique<TensorFlowUnsupportedOperator>();
op->tensorflow_op = "Merge";
const auto ops_by_type = BuildOperatorByTypeMap();
const toco::OperatorSignature op_signature = {op.get(), &model};
const auto key = details::OperatorKey(op_signature, ops_by_type, true);
EXPECT_EQ(key.type(), ::tflite::BuiltinOperator_CUSTOM);
EXPECT_EQ(key.custom_code(), "FlexMerge");
EXPECT_EQ(key.version(), 1);
EXPECT_FALSE(key.is_custom_op());
EXPECT_TRUE(key.is_flex_op());
// The control flow ops should be marked as unsupported.
EXPECT_TRUE(key.is_unsupported_flex_op());
}
TEST(OperatorKeyTest, TestFlexWithUnsupportedOp) {
Model model;
auto op = std::make_unique<TensorFlowUnsupportedOperator>();
op->tensorflow_op = "UnsupportedOp";
const auto ops_by_type = BuildOperatorByTypeMap();
const toco::OperatorSignature op_signature = {op.get(), &model};
const auto key = details::OperatorKey(op_signature, ops_by_type, true);
EXPECT_EQ(key.type(), ::tflite::BuiltinOperator_CUSTOM);
EXPECT_EQ(key.custom_code(), "UnsupportedOp");
EXPECT_EQ(key.version(), 1);
EXPECT_FALSE(key.is_flex_op());
EXPECT_FALSE(key.is_unsupported_flex_op());
}
TEST(OperatorKeyTest, TestFlexWithPartiallySupportedOps) {
// Test Toco-supported/TFLite-unsupported operators.
Model model;
auto op = std::make_unique<TensorFlowAssertOperator>();
const auto ops_by_type = BuildOperatorByTypeMap();
{
// If NodeDef isn't retained in the Toco op, a regular custom op
// will be exported.
const toco::OperatorSignature op_signature = {op.get(), &model};
const auto key = details::OperatorKey(op_signature, ops_by_type, true);
EXPECT_EQ(key.type(), ::tflite::BuiltinOperator_CUSTOM);
EXPECT_EQ(key.custom_code(), "Assert");
EXPECT_EQ(key.version(), 1);
EXPECT_TRUE(key.is_custom_op());
EXPECT_FALSE(key.is_flex_op());
}
::tensorflow::NodeDef node_def;
node_def.set_name("TensorFlowAssert");
node_def.set_op("TensorFlowAssert");
node_def.SerializeToString(&op->tensorflow_node_def);
{
// If NodeDef is retained in the Toco op, a Flex op will be exported.
const toco::OperatorSignature op_signature = {op.get(), &model};
const auto key = details::OperatorKey(op_signature, ops_by_type, true);
EXPECT_EQ(key.type(), ::tflite::BuiltinOperator_CUSTOM);
EXPECT_EQ(key.custom_code(), "FlexAssert");
EXPECT_EQ(key.version(), 1);
EXPECT_FALSE(key.is_custom_op());
EXPECT_TRUE(key.is_flex_op());
}
}
// TODO(ahentz): tests for tensors, inputs, outputs, opcodes and operators.
} // namespace
} // namespace tflite
} // namespace toco
tensorflow/lite/toco/tflite/import.cc
+251
View File
@@ -0,0 +1,251 @@
/* Copyright 2017 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/import.h"
#include <cstddef>
#include <cstdint>
#include <map>
#include <memory>
#include <string>
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/strings/string_view.h"
#include "flatbuffers/verifier.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/schema/schema_utils.h"
#include "tensorflow/lite/core/tools/verifier.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/stderr_reporter.h"
#include "tensorflow/lite/toco/model.h"
#include "tensorflow/lite/toco/model_flags.pb.h"
#include "tensorflow/lite/toco/tflite/operator.h"
#include "tensorflow/lite/toco/tflite/types.h"
#include "tensorflow/lite/toco/tooling_util.h"
namespace toco {
namespace tflite {
namespace details {
void LoadTensorsTable(const ::tflite::Model& input_model,
TensorsTable* tensors_table) {
// TODO(aselle): add support to toco for multiple subgraphs.
auto tensors = (*input_model.subgraphs())[0]->tensors();
if (!tensors) return;
for (const auto* tensor : *tensors) {
tensors_table->push_back(tensor->name()->c_str());
}
}
void LoadOperatorsTable(const ::tflite::Model& input_model,
OperatorsTable* operators_table) {
auto opcodes = input_model.operator_codes();
if (!opcodes) return;
for (const auto* opcode : *opcodes) {
auto builtin_code = GetBuiltinCode(opcode);
if (builtin_code != ::tflite::BuiltinOperator_CUSTOM) {
operators_table->push_back(EnumNameBuiltinOperator(builtin_code));
} else {
operators_table->push_back(opcode->custom_code()->c_str());
}
}
}
} // namespace details
void ImportTensors(const ::tflite::Model& input_model, Model* model) {
auto tensors = (*input_model.subgraphs())[0]->tensors();
auto* buffers = input_model.buffers();
// auto tensors = input_model.tensors();
if (!tensors) return;
for (const auto* input_tensor : *tensors) {
Array& array = model->GetOrCreateArray(input_tensor->name()->c_str());
array.data_type = DataType::Deserialize(input_tensor->type());
int buffer_index = input_tensor->buffer();
auto* buffer = buffers->Get(buffer_index);
DataBuffer::Deserialize(*input_tensor, *buffer, &array);
auto shape = input_tensor->shape();
if (shape) {
// If the shape is 0-dimensional, make sure to record it as such,
// as oppose to leaving the array without a shape.
array.mutable_shape()->mutable_dims()->clear();
for (uint32_t i = 0; i < shape->Length(); ++i) {
auto d = shape->Get(i);
array.mutable_shape()->mutable_dims()->push_back(d);
}
}
auto quantization = input_tensor->quantization();
if (quantization) {
// Note that tf.mini only supports a single quantization parameters for
// the whole array.
if (quantization->min() && quantization->max()) {
CHECK_EQ(1, quantization->min()->Length());
CHECK_EQ(1, quantization->max()->Length());
MinMax& minmax = array.GetOrCreateMinMax();
minmax.min = quantization->min()->Get(0);
minmax.max = quantization->max()->Get(0);
}
if (quantization->scale() && quantization->zero_point()) {
CHECK_EQ(1, quantization->scale()->Length());
CHECK_EQ(1, quantization->zero_point()->Length());
QuantizationParams& q = array.GetOrCreateQuantizationParams();
q.scale = quantization->scale()->Get(0);
q.zero_point = quantization->zero_point()->Get(0);
}
}
}
}
void ImportOperators(
const ::tflite::Model& input_model,
const std::map<std::string, std::unique_ptr<BaseOperator>>& ops_by_name,
const details::TensorsTable& tensors_table,
const details::OperatorsTable& operators_table, Model* model) {
// TODO(aselle): add support for multiple subgraphs.
auto ops = (*input_model.subgraphs())[0]->operators();
if (!ops) return;
for (const auto* input_op : *ops) {
uint32_t index = input_op->opcode_index();
if (index > operators_table.size()) {
LOG(FATAL) << "Index " << index << " must be between zero and "
<< operators_table.size();
}
std::string opname = operators_table.at(index);
// Find and use the appropriate operator deserialization factory.
std::unique_ptr<Operator> new_op = nullptr;
if (ops_by_name.count(opname) == 0) {
std::string effective_opname = "TENSORFLOW_UNSUPPORTED";
if (ops_by_name.count(effective_opname) == 0) {
LOG(FATAL) << "Internal logic error: TENSORFLOW_UNSUPPORTED not found.";
}
new_op = ops_by_name.at(effective_opname)
->Deserialize(input_op->builtin_options(),
input_op->custom_options());
if (new_op->type == OperatorType::kUnsupported) {
auto* unsupported_op =
static_cast<TensorFlowUnsupportedOperator*>(new_op.get());
unsupported_op->tensorflow_op = opname;
// TODO(b/109932940): Remove this when quantized is removed.
// For now, we assume all ops are quantized.
unsupported_op->quantized = true;
} else {
LOG(FATAL) << "Expected a TensorFlowUnsupportedOperator";
}
} else {
new_op = ops_by_name.at(opname)->Deserialize(input_op->builtin_options(),
input_op->custom_options());
}
model->operators.emplace_back(new_op.release());
auto* op = model->operators.back().get();
// Make sure all the inputs and outputs are hooked up.
auto inputs = input_op->inputs();
for (uint32_t i = 0; i < inputs->Length(); i++) {
auto input_index = inputs->Get(i);
// input_index == -1 indicates optional tensor.
if (input_index != -1) {
const std::string& input_name = tensors_table.at(input_index);
op->inputs.push_back(input_name);
} else {
const std::string& tensor_name =
toco::AvailableArrayName(*model, "OptionalTensor");
model->CreateOptionalArray(tensor_name);
op->inputs.push_back(tensor_name);
}
}
auto outputs = input_op->outputs();
for (int i = 0, end = outputs->Length(); i < end; i++) {
auto output_index = outputs->Get(i);
const std::string& output_name = tensors_table.at(output_index);
op->outputs.push_back(output_name);
}
}
}
void ImportIOTensors(const ModelFlags& model_flags,
const ::tflite::Model& input_model,
const details::TensorsTable& tensors_table, Model* model) {
// Import from the first subgraph if input arrays have not been specified.
if (model_flags.input_arrays().empty()) {
auto inputs = (*input_model.subgraphs())[0]->inputs();
if (inputs) {
for (int input : *inputs) {
const std::string& input_name = tensors_table.at(input);
model->flags.add_input_arrays()->set_name(input_name);
}
}
}
// Import from the first subgraph if output arrays have not been specified.
if (model_flags.output_arrays().empty()) {
auto outputs = (*input_model.subgraphs())[0]->outputs();
if (outputs) {
for (int output : *outputs) {
const std::string& output_name = tensors_table.at(output);
model->flags.add_output_arrays(output_name);
}
}
}
}
namespace {
bool Verify(const void* buf, size_t len) {
::flatbuffers::Verifier verifier(static_cast<const uint8_t*>(buf), len);
return ::tflite::VerifyModelBuffer(verifier);
}
} // namespace
std::unique_ptr<Model> Import(const ModelFlags& model_flags,
absl::string_view input_file_contents) {
::tflite::AlwaysTrueResolver r;
if (!::tflite::Verify(input_file_contents.data(), input_file_contents.size(),
r, ::tflite::DefaultErrorReporter())) {
LOG(FATAL) << "Invalid flatbuffer.";
}
const ::tflite::Model* input_model =
::tflite::GetModel(input_file_contents.data());
// Full list of all known operators.
const auto ops_by_name = BuildOperatorByNameMap();
if (!input_model->subgraphs() || input_model->subgraphs()->size() != 1) {
LOG(FATAL) << "Number of subgraphs in tflite should be exactly 1.";
}
std::unique_ptr<Model> model;
model = std::make_unique<Model>();
details::TensorsTable tensors_table;
details::LoadTensorsTable(*input_model, &tensors_table);
details::OperatorsTable operators_table;
details::LoadOperatorsTable(*input_model, &operators_table);
ImportTensors(*input_model, model.get());
ImportOperators(*input_model, ops_by_name, tensors_table, operators_table,
model.get());
ImportIOTensors(model_flags, *input_model, tensors_table, model.get());
UndoWeightsShuffling(model.get());
return model;
}
} // namespace tflite
} // namespace toco
tensorflow/lite/toco/tflite/import.h
+55
View File
@@ -0,0 +1,55 @@
/* Copyright 2017 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_TOCO_TFLITE_IMPORT_H_
#define TENSORFLOW_LITE_TOCO_TFLITE_IMPORT_H_
#include <memory>
#include <string>
#include <vector>
#include "absl/strings/string_view.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/toco/model.h"
#include "tensorflow/lite/toco/model_flags.pb.h"
namespace toco {
namespace tflite {
// Parse the given string as TF Lite flatbuffer and return a new tf.mini model.
std::unique_ptr<Model> Import(const ModelFlags& model_flags,
absl::string_view input_file_contents);
namespace details {
// The names of all tensors found in a TF Lite model.
using TensorsTable = std::vector<std::string>;
// The names of all operators found in TF Lite model. If the operator is
// builtin, the string representation of the corresponding enum value is used
// as name.
using OperatorsTable = std::vector<std::string>;
void LoadTensorsTable(const ::tflite::Model &input_model,
TensorsTable *tensors_table);
void LoadOperatorsTable(const ::tflite::Model &input_model,
OperatorsTable *operators_table);
} // namespace details
} // namespace tflite
} // namespace toco
#endif // TENSORFLOW_LITE_TOCO_TFLITE_IMPORT_H_
tensorflow/lite/toco/tflite/import_test.cc
+275
View File
@@ -0,0 +1,275 @@
/* Copyright 2017 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/import.h"
#include <cstdint>
#include <initializer_list>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "flatbuffers/vector.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/schema/schema_conversion_utils.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/toco/model.h"
#include "tensorflow/lite/toco/model_flags.pb.h"
#include "tensorflow/lite/toco/toco_types.h"
#include "tensorflow/lite/version.h"
namespace toco {
namespace tflite {
namespace {
using ::testing::ElementsAre;
using flatbuffers::Offset;
using flatbuffers::Vector;
class ImportTest : public ::testing::Test {
protected:
template <typename T>
Offset<Vector<unsigned char>> CreateDataVector(const std::vector<T>& data) {
return builder_.CreateVector(reinterpret_cast<const uint8_t*>(data.data()),
sizeof(T) * data.size());
}
Offset<Vector<Offset<::tflite::Buffer>>> BuildBuffers() {
auto buf0 = ::tflite::CreateBuffer(builder_, CreateDataVector<float>({}));
auto buf1 = ::tflite::CreateBuffer(
builder_, CreateDataVector<float>({1.0f, 2.0f, 3.0f, 4.0f}));
auto buf2 =
::tflite::CreateBuffer(builder_, CreateDataVector<float>({3.0f, 4.0f}));
return builder_.CreateVector(
std::vector<Offset<::tflite::Buffer>>({buf0, buf1, buf2}));
}
Offset<Vector<Offset<::tflite::Tensor>>> BuildTensors() {
auto q = ::tflite::CreateQuantizationParameters(
builder_,
/*min=*/builder_.CreateVector<float>({0.1f}),
/*max=*/builder_.CreateVector<float>({0.2f}),
/*scale=*/builder_.CreateVector<float>({0.3f}),
/*zero_point=*/builder_.CreateVector<int64_t>({100LL}));
auto t1 =
::tflite::CreateTensor(builder_, builder_.CreateVector<int>({1, 2, 2}),
::tflite::TensorType_FLOAT32, 1,
builder_.CreateString("tensor_one"), q);
auto t2 =
::tflite::CreateTensor(builder_, builder_.CreateVector<int>({2, 1}),
::tflite::TensorType_FLOAT32, 0,
builder_.CreateString("tensor_two"), q);
return builder_.CreateVector(
std::vector<Offset<::tflite::Tensor>>({t1, t2}));
}
Offset<Vector<Offset<::tflite::OperatorCode>>> BuildOpCodes(
std::initializer_list<::tflite::BuiltinOperator> op_codes) {
std::vector<Offset<::tflite::OperatorCode>> op_codes_vector;
for (auto op : op_codes) {
op_codes_vector.push_back(::tflite::CreateOperatorCode(builder_, op, 0));
}
return builder_.CreateVector(op_codes_vector);
}
Offset<Vector<Offset<::tflite::OperatorCode>>> BuildOpCodes() {
return BuildOpCodes({::tflite::BuiltinOperator_MAX_POOL_2D,
::tflite::BuiltinOperator_CONV_2D});
}
Offset<Vector<Offset<::tflite::Operator>>> BuildOperators(
std::initializer_list<int> inputs, std::initializer_list<int> outputs) {
auto is = builder_.CreateVector<int>(inputs);
if (inputs.size() == 0) is = 0;
auto os = builder_.CreateVector<int>(outputs);
if (outputs.size() == 0) os = 0;
auto op = ::tflite::CreateOperator(
builder_, 0, is, os, ::tflite::BuiltinOptions_Conv2DOptions,
::tflite::CreateConv2DOptions(builder_, ::tflite::Padding_VALID, 1, 1,
::tflite::ActivationFunctionType_NONE)
.Union(),
/*custom_options=*/0, ::tflite::CustomOptionsFormat_FLEXBUFFERS);
return builder_.CreateVector(std::vector<Offset<::tflite::Operator>>({op}));
}
Offset<Vector<Offset<::tflite::Operator>>> BuildOperators() {
return BuildOperators({0}, {1});
}
Offset<Vector<Offset<::tflite::SubGraph>>> BuildSubGraphs(
Offset<Vector<Offset<::tflite::Tensor>>> tensors,
Offset<Vector<Offset<::tflite::Operator>>> operators,
int num_sub_graphs = 1) {
std::vector<int32_t> inputs = {0};
std::vector<int32_t> outputs = {1};
std::vector<Offset<::tflite::SubGraph>> v;
for (int i = 0; i < num_sub_graphs; ++i) {
v.push_back(::tflite::CreateSubGraph(
builder_, tensors, builder_.CreateVector(inputs),
builder_.CreateVector(outputs), operators,
builder_.CreateString("subgraph")));
}
return builder_.CreateVector(v);
}
// This is a very simplistic model. We are not interested in testing all the
// details here, since tf.mini's testing framework will be exercising all the
// conversions multiple times, and the conversion of operators is tested by
// separate unittests.
void BuildTestModel() {
auto buffers = BuildBuffers();
auto tensors = BuildTensors();
auto opcodes = BuildOpCodes();
auto operators = BuildOperators();
auto subgraphs = BuildSubGraphs(tensors, operators);
auto s = builder_.CreateString("");
::tflite::FinishModelBuffer(
builder_, ::tflite::CreateModel(builder_, TFLITE_SCHEMA_VERSION,
opcodes, subgraphs, s, buffers));
input_model_ = ::tflite::GetModel(builder_.GetBufferPointer());
}
std::string InputModelAsString() {
return std::string(reinterpret_cast<char*>(builder_.GetBufferPointer()),
builder_.GetSize());
}
flatbuffers::FlatBufferBuilder builder_;
const ::tflite::Model* input_model_ = nullptr;
};
TEST_F(ImportTest, LoadTensorsTable) {
BuildTestModel();
details::TensorsTable tensors;
details::LoadTensorsTable(*input_model_, &tensors);
EXPECT_THAT(tensors, ElementsAre("tensor_one", "tensor_two"));
}
TEST_F(ImportTest, LoadOperatorsTable) {
BuildTestModel();
details::OperatorsTable operators;
details::LoadOperatorsTable(*input_model_, &operators);
EXPECT_THAT(operators, ElementsAre("MAX_POOL_2D", "CONV_2D"));
}
TEST_F(ImportTest, Tensors) {
BuildTestModel();
auto model = Import(ModelFlags(), InputModelAsString());
ASSERT_GT(model->HasArray("tensor_one"), 0);
Array& a1 = model->GetArray("tensor_one");
EXPECT_EQ(ArrayDataType::kFloat, a1.data_type);
EXPECT_THAT(a1.GetBuffer<ArrayDataType::kFloat>().data,
ElementsAre(1.0f, 2.0f, 3.0f, 4.0f));
ASSERT_TRUE(a1.has_shape());
EXPECT_THAT(a1.shape().dims(), ElementsAre(1, 2, 2));
const auto& mm = a1.minmax;
ASSERT_TRUE(mm.get());
EXPECT_FLOAT_EQ(0.1, mm->min);
EXPECT_FLOAT_EQ(0.2, mm->max);
const auto& q = a1.quantization_params;
ASSERT_TRUE(q.get());
EXPECT_FLOAT_EQ(0.3, q->scale);
EXPECT_EQ(100, q->zero_point);
}
TEST_F(ImportTest, NoBuffers) {
auto buffers = 0;
auto tensors = BuildTensors();
auto opcodes = BuildOpCodes();
auto operators = BuildOperators();
auto subgraphs = BuildSubGraphs(tensors, operators);
auto comment = builder_.CreateString("");
::tflite::FinishModelBuffer(
builder_, ::tflite::CreateModel(builder_, TFLITE_SCHEMA_VERSION, opcodes,
subgraphs, comment, buffers));
EXPECT_DEATH(Import(ModelFlags(), InputModelAsString()),
"Missing 'buffers' section.");
}
TEST_F(ImportTest, NoInputs) {
auto buffers = BuildBuffers();
auto tensors = BuildTensors();
auto opcodes = BuildOpCodes();
auto operators = BuildOperators({}, {1});
auto subgraphs = BuildSubGraphs(tensors, operators);
auto comment = builder_.CreateString("");
::tflite::FinishModelBuffer(
builder_, ::tflite::CreateModel(builder_, TFLITE_SCHEMA_VERSION, opcodes,
subgraphs, comment, buffers));
EXPECT_DEATH(Import(ModelFlags(), InputModelAsString()),
"Missing 'inputs' for operator.");
}
TEST_F(ImportTest, NoOutputs) {
auto buffers = BuildBuffers();
auto tensors = BuildTensors();
auto opcodes = BuildOpCodes();
auto operators = BuildOperators({0}, {});
auto subgraphs = BuildSubGraphs(tensors, operators);
auto comment = builder_.CreateString("");
::tflite::FinishModelBuffer(
builder_, ::tflite::CreateModel(builder_, TFLITE_SCHEMA_VERSION, opcodes,
subgraphs, comment, buffers));
EXPECT_DEATH(Import(ModelFlags(), InputModelAsString()),
"Missing 'outputs' for operator.");
}
TEST_F(ImportTest, InvalidOpCode) {
auto buffers = BuildBuffers();
auto tensors = BuildTensors();
auto opcodes = BuildOpCodes({static_cast<::tflite::BuiltinOperator>(-1),
::tflite::BuiltinOperator_CONV_2D});
auto operators = BuildOperators();
auto subgraphs = BuildSubGraphs(tensors, operators);
auto comment = builder_.CreateString("");
::tflite::FinishModelBuffer(
builder_, ::tflite::CreateModel(builder_, TFLITE_SCHEMA_VERSION, opcodes,
subgraphs, comment, buffers));
EXPECT_DEATH(Import(ModelFlags(), InputModelAsString()),
"Operator id '-1' is out of range.");
}
TEST_F(ImportTest, MultipleSubGraphs) {
auto buffers = BuildBuffers();
auto tensors = BuildTensors();
auto opcodes = BuildOpCodes();
auto operators = BuildOperators();
auto subgraphs = BuildSubGraphs(tensors, operators, 2);
auto comment = builder_.CreateString("");
::tflite::FinishModelBuffer(
builder_, ::tflite::CreateModel(builder_, TFLITE_SCHEMA_VERSION, opcodes,
subgraphs, comment, buffers));
input_model_ = ::tflite::GetModel(builder_.GetBufferPointer());
EXPECT_DEATH(Import(ModelFlags(), InputModelAsString()),
"Number of subgraphs in tflite should be exactly 1.");
}
// TODO(ahentz): still need tests for Operators and IOTensors.
} // namespace
} // namespace tflite
} // namespace toco
tensorflow/lite/toco/tflite/operator.cc
+2150
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/toco/tflite/operator.h
+139
View File
@@ -0,0 +1,139 @@
/* Copyright 2017 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_TOCO_TFLITE_OPERATOR_H_
#define TENSORFLOW_LITE_TOCO_TFLITE_OPERATOR_H_
#include <cstdint>
#include <map>
#include <memory>
#include <string>
#include <vector>
#include "flatbuffers/flatbuffers.h"
#include "flatbuffers/flexbuffers.h"
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "flatbuffers/vector.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/tools/versioning/op_signature.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/toco/model.h"
namespace toco {
namespace tflite {
class BaseOperator;
// Return a map contained all know TF Lite Operators, keyed by their names.
std::map<std::string, std::unique_ptr<BaseOperator>> BuildOperatorByNameMap(
bool enable_select_tf_ops = false);
// Return a map contained all know TF Lite Operators, keyed by the type of
// their tf.mini counterparts.
std::map<OperatorType, std::unique_ptr<BaseOperator>> BuildOperatorByTypeMap(
bool enable_select_tf_ops = false);
// Write the custom option FlexBuffer with a serialized TensorFlow NodeDef
// for a Flex op.
std::unique_ptr<flexbuffers::Builder> WriteFlexOpOptions(
const std::string& tensorflow_node_def);
// These are the flatbuffer types for custom and builtin options.
using CustomOptions = flatbuffers::Vector<uint8_t>;
using BuiltinOptions = void;
// A simple wrapper around the flatbuffer objects used to describe options that
// configure operators.
struct Options {
// Build custom options.
static Options Custom(flatbuffers::Offset<CustomOptions> offset) {
return {::tflite::BuiltinOptions_NONE, 0, offset};
}
// Build builtin options of the given type.
static Options Builtin(::tflite::BuiltinOptions type,
flatbuffers::Offset<BuiltinOptions> offset) {
return {type, offset, 0};
}
::tflite::BuiltinOptions type;
flatbuffers::Offset<BuiltinOptions> builtin;
flatbuffers::Offset<CustomOptions> custom;
};
// A BaseOperator encapsulates the relationship between operators in tf.mini
// and TF lite, and provides methods for converting between those two formats.
class BaseOperator {
public:
// Build an operator with the given TF Lite name and tf.mini type.
BaseOperator(const std::string& name, OperatorType type)
: name_(name), type_(type) {}
virtual ~BaseOperator() = default;
std::string name() const { return name_; }
OperatorType type() const { return type_; }
// Given a tf.mini operator, create the corresponding flatbuffer options and
// return their offsets.
virtual Options Serialize(const Operator& op,
flatbuffers::FlatBufferBuilder* builder) const = 0;
// Read TF Lite options and create the appropriate tf.mini operator.
virtual std::unique_ptr<Operator> Deserialize(
const BuiltinOptions* builtin_options,
const CustomOptions* custom_options) const = 0;
// Get the op version using the OperatorSignature.
// The function needs to be overridden to return the op version based on the
// parameters. Note:
// * The first version for each op should be 1 (to be consistent with the
// default value in Flatbuffer. `return 1;` is okay for newly implemented
// ops.
// * When multiple versions are defined for an op, this function could be
// overridden. (See example in `operator_test.cc` and
// 'tools/versioning/op_version.cc`)
virtual int GetVersion(const OperatorSignature& op_signature) const = 0;
// Given a Toco `Operator`, return a list of booleans indicating the op
// mutates which input variables.
// * If the op mutates any input variables, it should return a list of bool
// with the same length as inputs.
// * Otherwise, it will return an empty list.
virtual std::vector<bool> GetMutatingInputVariables(
const Operator& op) const {
// Most ops don't have variable tensors. This function can be overridden.
return std::vector<bool>();
}
private:
std::string name_;
OperatorType type_;
};
// Helper function to create ::tflite::OpSignature from the given
// ::tflite::BuiltinOperator and OperatorSignature.
::tflite::OpSignature GetVersioningOpSig(const ::tflite::BuiltinOperator op,
const OperatorSignature& op_signature);
// Helper function to determine if a unsupported TensorFlow op should be
// exported as an Flex op or a regular custom op.
bool ShouldExportAsFlexOp(bool enable_select_tf_ops,
const std::string& tensorflow_op_name);
} // namespace tflite
} // namespace toco
#endif // TENSORFLOW_LITE_TOCO_TFLITE_OPERATOR_H_
tensorflow/lite/toco/tflite/operator_test.cc
+1314
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/toco/tflite/simple_operator.h
+66
View File
@@ -0,0 +1,66 @@
/* Copyright 2017 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_TOCO_TFLITE_SIMPLE_OPERATOR_H_
#define TENSORFLOW_LITE_TOCO_TFLITE_SIMPLE_OPERATOR_H_
#include "tensorflow/compiler/mlir/lite/tools/versioning/op_version.h"
#include "tensorflow/lite/toco/tflite/operator.h"
namespace toco {
namespace tflite {
// Simple operators don't have any configuration options and can be trivially
// serialized and deserialized. Note that most of toco's operators will
// likely be supported as builtin operators in TF Lite. Simple (and custom)
// operators are mostly a convenience for the times when tf.mini supports more
// operators than TF Lite.
//
// Template argument T must derive from ::toco::Operator.
template <typename T>
class SimpleOperator : public BaseOperator {
public:
using BaseOperator::BaseOperator;
SimpleOperator(::tflite::BuiltinOperator op, OperatorType type)
: BaseOperator(::tflite::EnumNameBuiltinOperator(op), type),
builtin_op_(op) {}
Options Serialize(const Operator& op,
flatbuffers::FlatBufferBuilder* builder) const override {
return Options();
}
std::unique_ptr<Operator> Deserialize(
const BuiltinOptions* builtin_options,
const CustomOptions* custom_options) const override {
return std::unique_ptr<Operator>(new T);
}
int GetVersion(const OperatorSignature& op_signature) const override {
return ::tflite::GetBuiltinOperatorVersion(
GetVersioningOpSig(builtin_op_, op_signature));
}
::tflite::BuiltinOperator builtin_op() const { return builtin_op_; }
private:
const ::tflite::BuiltinOperator builtin_op_;
};
} // namespace tflite
} // namespace toco
#endif // TENSORFLOW_LITE_TOCO_TFLITE_SIMPLE_OPERATOR_H_
tensorflow/lite/toco/tflite/types.cc
+269
View File
@@ -0,0 +1,269 @@
/* Copyright 2017 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/types.h"
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>
#include "absl/log/log.h"
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "flatbuffers/vector.h" // from @flatbuffers
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/string_util.h"
#include "tensorflow/lite/toco/model.h"
#include "tensorflow/lite/toco/runtime/types.h"
namespace toco {
namespace tflite {
namespace {
DataBuffer::FlatBufferOffset CopyStringToBuffer(
const Array& array, flatbuffers::FlatBufferBuilder* builder) {
const auto& src_data = array.GetBuffer<ArrayDataType::kString>().data;
::tflite::DynamicBuffer dyn_buffer;
for (const std::string& str : src_data) {
dyn_buffer.AddString(str.c_str(), str.length());
}
char* tensor_buffer;
int bytes = dyn_buffer.WriteToBuffer(&tensor_buffer);
std::vector<uint8_t> dst_data(bytes);
memcpy(dst_data.data(), tensor_buffer, bytes);
free(tensor_buffer);
return builder->CreateVector(dst_data.data(), bytes);
}
// vector<bool> may be implemented using a bit-set, so we can't just
// reinterpret_cast, accessing its data as vector<bool> and let flatbuffer
// CreateVector handle it.
// Background: https://isocpp.org/blog/2012/11/on-vectorbool
DataBuffer::FlatBufferOffset CopyBoolToBuffer(
const Array& array, flatbuffers::FlatBufferBuilder* builder) {
const auto& src_data = array.GetBuffer<ArrayDataType::kBool>().data;
return builder->CreateVector(src_data);
}
template <ArrayDataType T>
DataBuffer::FlatBufferOffset CopyBuffer(
const Array& array, flatbuffers::FlatBufferBuilder* builder) {
using NativeT = ::toco::DataType<T>;
const auto& src_data = array.GetBuffer<T>().data;
const uint8_t* dst_data = reinterpret_cast<const uint8_t*>(src_data.data());
auto size = src_data.size() * sizeof(NativeT);
return builder->CreateVector(dst_data, size);
}
void CopyStringFromBuffer(const ::tflite::Buffer& buffer, Array* array) {
auto* src_data = reinterpret_cast<const char*>(buffer.data()->data());
std::vector<std::string>* dst_data =
&array->GetMutableBuffer<ArrayDataType::kString>().data;
int32_t num_strings = ::tflite::GetStringCount(src_data);
for (int i = 0; i < num_strings; i++) {
::tflite::StringRef str_ref = ::tflite::GetString(src_data, i);
std::string this_str(str_ref.str, str_ref.len);
dst_data->push_back(this_str);
}
}
template <ArrayDataType T>
void CopyBuffer(const ::tflite::Buffer& buffer, Array* array) {
using NativeT = ::toco::DataType<T>;
auto* src_buffer = buffer.data();
const NativeT* src_data =
reinterpret_cast<const NativeT*>(src_buffer->data());
int num_items = src_buffer->size() / sizeof(NativeT);
std::vector<NativeT>* dst_data = &array->GetMutableBuffer<T>().data;
for (int i = 0; i < num_items; ++i) {
dst_data->push_back(*src_data);
++src_data;
}
}
} // namespace
::tflite::TensorType DataType::Serialize(ArrayDataType array_data_type) {
switch (array_data_type) {
case ArrayDataType::kFloat:
return ::tflite::TensorType_FLOAT32;
case ArrayDataType::kInt16:
return ::tflite::TensorType_INT16;
case ArrayDataType::kInt32:
return ::tflite::TensorType_INT32;
case ArrayDataType::kUint32:
return ::tflite::TensorType_UINT32;
case ArrayDataType::kInt64:
return ::tflite::TensorType_INT64;
case ArrayDataType::kUint8:
return ::tflite::TensorType_UINT8;
case ArrayDataType::kUint16:
return ::tflite::TensorType_UINT16;
case ArrayDataType::kString:
return ::tflite::TensorType_STRING;
case ArrayDataType::kBool:
return ::tflite::TensorType_BOOL;
case ArrayDataType::kComplex64:
return ::tflite::TensorType_COMPLEX64;
default:
// FLOAT32 is filled for unknown data types.
return ::tflite::TensorType_FLOAT32;
}
}
ArrayDataType DataType::Deserialize(int tensor_type) {
switch (::tflite::TensorType(tensor_type)) {
case ::tflite::TensorType_FLOAT32:
return ArrayDataType::kFloat;
case ::tflite::TensorType_INT16:
return ArrayDataType::kInt16;
case ::tflite::TensorType_INT32:
return ArrayDataType::kInt32;
case ::tflite::TensorType_UINT32:
return ArrayDataType::kUint32;
case ::tflite::TensorType_INT64:
return ArrayDataType::kInt64;
case ::tflite::TensorType_STRING:
return ArrayDataType::kString;
case ::tflite::TensorType_UINT8:
return ArrayDataType::kUint8;
case ::tflite::TensorType_UINT16:
return ArrayDataType::kUint16;
case ::tflite::TensorType_BOOL:
return ArrayDataType::kBool;
case ::tflite::TensorType_COMPLEX64:
return ArrayDataType::kComplex64;
default:
LOG(FATAL) << "Unhandled tensor type '" << tensor_type << "'.";
}
}
flatbuffers::Offset<flatbuffers::Vector<uint8_t>> DataBuffer::Serialize(
const Array& array, flatbuffers::FlatBufferBuilder* builder) {
if (!array.buffer) return 0; // an empty buffer, usually an output.
switch (array.data_type) {
case ArrayDataType::kFloat:
return CopyBuffer<ArrayDataType::kFloat>(array, builder);
case ArrayDataType::kInt16:
return CopyBuffer<ArrayDataType::kInt16>(array, builder);
case ArrayDataType::kInt32:
return CopyBuffer<ArrayDataType::kInt32>(array, builder);
case ArrayDataType::kUint32:
return CopyBuffer<ArrayDataType::kUint32>(array, builder);
case ArrayDataType::kInt64:
return CopyBuffer<ArrayDataType::kInt64>(array, builder);
case ArrayDataType::kString:
return CopyStringToBuffer(array, builder);
case ArrayDataType::kUint8:
return CopyBuffer<ArrayDataType::kUint8>(array, builder);
case ArrayDataType::kBool:
return CopyBoolToBuffer(array, builder);
case ArrayDataType::kComplex64:
return CopyBuffer<ArrayDataType::kComplex64>(array, builder);
default:
LOG(FATAL) << "Unhandled array data type.";
}
}
void DataBuffer::Deserialize(const ::tflite::Tensor& tensor,
const ::tflite::Buffer& buffer, Array* array) {
if (tensor.buffer() == 0) return; // an empty buffer, usually an output.
if (buffer.data() == nullptr) return; // a non-defined buffer.
switch (tensor.type()) {
case ::tflite::TensorType_FLOAT32:
return CopyBuffer<ArrayDataType::kFloat>(buffer, array);
case ::tflite::TensorType_INT16:
return CopyBuffer<ArrayDataType::kInt16>(buffer, array);
case ::tflite::TensorType_INT32:
return CopyBuffer<ArrayDataType::kInt32>(buffer, array);
case ::tflite::TensorType_UINT32:
return CopyBuffer<ArrayDataType::kUint32>(buffer, array);
case ::tflite::TensorType_INT64:
return CopyBuffer<ArrayDataType::kInt64>(buffer, array);
case ::tflite::TensorType_STRING:
return CopyStringFromBuffer(buffer, array);
case ::tflite::TensorType_UINT8:
return CopyBuffer<ArrayDataType::kUint8>(buffer, array);
case ::tflite::TensorType_BOOL:
return CopyBuffer<ArrayDataType::kBool>(buffer, array);
case ::tflite::TensorType_COMPLEX64:
return CopyBuffer<ArrayDataType::kComplex64>(buffer, array);
default:
LOG(FATAL) << "Unhandled tensor type.";
}
}
::tflite::Padding Padding::Serialize(PaddingType padding_type) {
switch (padding_type) {
case PaddingType::kSame:
return ::tflite::Padding_SAME;
case PaddingType::kValid:
return ::tflite::Padding_VALID;
default:
LOG(FATAL) << "Unhandled padding type.";
}
}
PaddingType Padding::Deserialize(int padding) {
switch (::tflite::Padding(padding)) {
case ::tflite::Padding_SAME:
return PaddingType::kSame;
case ::tflite::Padding_VALID:
return PaddingType::kValid;
default:
LOG(FATAL) << "Unhandled padding.";
}
}
::tflite::ActivationFunctionType ActivationFunction::Serialize(
FusedActivationFunctionType faf_type) {
switch (faf_type) {
case FusedActivationFunctionType::kNone:
return ::tflite::ActivationFunctionType_NONE;
case FusedActivationFunctionType::kRelu:
return ::tflite::ActivationFunctionType_RELU;
case FusedActivationFunctionType::kRelu6:
return ::tflite::ActivationFunctionType_RELU6;
case FusedActivationFunctionType::kRelu1:
return ::tflite::ActivationFunctionType_RELU_N1_TO_1;
default:
LOG(FATAL) << "Unhandled fused activation function type.";
}
}
FusedActivationFunctionType ActivationFunction::Deserialize(
int activation_function) {
switch (::tflite::ActivationFunctionType(activation_function)) {
case ::tflite::ActivationFunctionType_NONE:
return FusedActivationFunctionType::kNone;
case ::tflite::ActivationFunctionType_RELU:
return FusedActivationFunctionType::kRelu;
case ::tflite::ActivationFunctionType_RELU6:
return FusedActivationFunctionType::kRelu6;
case ::tflite::ActivationFunctionType_RELU_N1_TO_1:
return FusedActivationFunctionType::kRelu1;
default:
LOG(FATAL) << "Unhandled fused activation function type.";
}
}
} // namespace tflite
} // namespace toco
tensorflow/lite/toco/tflite/types.h
+64
View File
@@ -0,0 +1,64 @@
/* Copyright 2017 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_TOCO_TFLITE_TYPES_H_
#define TENSORFLOW_LITE_TOCO_TFLITE_TYPES_H_
#include <cstdint>
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "flatbuffers/vector.h" // from @flatbuffers
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/toco/model.h"
#include "tensorflow/lite/toco/runtime/types.h"
namespace toco {
namespace tflite {
struct DataType {
static ::tflite::TensorType Serialize(ArrayDataType array_data_type);
static ArrayDataType Deserialize(int tensor_type);
};
struct DataBuffer {
using FlatBufferOffset = flatbuffers::Offset<flatbuffers::Vector<uint8_t>>;
// Build the flatbuffer representation of a toco's Array and return the
// corresponding offset into the flatbuffer. Note that data from the array
// will be copied into the flatbuffer.
static FlatBufferOffset Serialize(const Array& array,
flatbuffers::FlatBufferBuilder* builder);
// Copy data from the given tensor into toco's Array.
static void Deserialize(const ::tflite::Tensor& tensor,
const ::tflite::Buffer& buffer, Array* array);
};
struct Padding {
static ::tflite::Padding Serialize(PaddingType padding_type);
static PaddingType Deserialize(int padding);
};
struct ActivationFunction {
static ::tflite::ActivationFunctionType Serialize(
FusedActivationFunctionType faf_type);
static FusedActivationFunctionType Deserialize(int activation_function);
};
} // namespace tflite
} // namespace toco
#endif // TENSORFLOW_LITE_TOCO_TFLITE_TYPES_H_
tensorflow/lite/toco/tflite/types_test.cc
+239
View File
@@ -0,0 +1,239 @@
/* Copyright 2017 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/types.h"
#include <complex>
#include <cstdint>
#include <initializer_list>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "flatbuffers/vector.h" // from @flatbuffers
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/toco/model.h"
#include "tensorflow/lite/toco/runtime/types.h"
namespace toco {
namespace tflite {
namespace {
using flatbuffers::Offset;
using flatbuffers::Vector;
// These are types that exist in TF Mini but don't have a correspondence
// in TF Lite.
static const ArrayDataType kUnsupportedTocoTypes[] = {ArrayDataType::kNone};
// These are TF Lite types for which there is no correspondence in TF Mini.
static const ::tflite::TensorType kUnsupportedTfLiteTypes[] = {
::tflite::TensorType_FLOAT16};
// A little helper to match flatbuffer offsets.
MATCHER_P(HasOffset, value, "") { return arg.o == value; }
// Helper function that creates an array, writes it into a flatbuffer, and then
// reads it back in.
template <ArrayDataType T>
Array ToFlatBufferAndBack(std::initializer_list<::toco::DataType<T>> items) {
// NOTE: This test does not construct the full buffers list. Since
// Deserialize normally takes a buffer, we need to synthesize one and provide
// an index that is non-zero so the buffer is not assumed to be empty.
Array src;
src.data_type = T;
src.GetMutableBuffer<T>().data = items;
Array result;
flatbuffers::FlatBufferBuilder builder;
builder.Finish(CreateTensor(builder, 0, DataType::Serialize(T),
/*buffer*/ 1)); // Can't use 0 which means empty.
flatbuffers::FlatBufferBuilder buffer_builder;
Offset<Vector<uint8_t>> data_buffer =
DataBuffer::Serialize(src, &buffer_builder);
buffer_builder.Finish(::tflite::CreateBuffer(buffer_builder, data_buffer));
auto* tensor =
flatbuffers::GetRoot<::tflite::Tensor>(builder.GetBufferPointer());
auto* buffer =
flatbuffers::GetRoot<::tflite::Buffer>(buffer_builder.GetBufferPointer());
DataBuffer::Deserialize(*tensor, *buffer, &result);
return result;
}
TEST(DataType, SupportedTypes) {
std::vector<std::pair<ArrayDataType, ::tflite::TensorType>> testdata = {
{ArrayDataType::kUint8, ::tflite::TensorType_UINT8},
{ArrayDataType::kInt32, ::tflite::TensorType_INT32},
{ArrayDataType::kUint32, ::tflite::TensorType_UINT32},
{ArrayDataType::kInt64, ::tflite::TensorType_INT64},
{ArrayDataType::kFloat, ::tflite::TensorType_FLOAT32},
{ArrayDataType::kBool, ::tflite::TensorType_BOOL},
{ArrayDataType::kComplex64, ::tflite::TensorType_COMPLEX64}};
for (auto x : testdata) {
EXPECT_EQ(x.second, DataType::Serialize(x.first));
EXPECT_EQ(x.first, DataType::Deserialize(x.second));
}
}
TEST(DataType, UnsupportedTypes) {
for (::tflite::TensorType t : kUnsupportedTfLiteTypes) {
EXPECT_DEATH(DataType::Deserialize(t), "Unhandled tensor type.");
}
// Unsupported types are all serialized as FLOAT32 currently.
for (ArrayDataType t : kUnsupportedTocoTypes) {
EXPECT_EQ(::tflite::TensorType_FLOAT32, DataType::Serialize(t));
}
}
TEST(DataBuffer, EmptyBuffers) {
flatbuffers::FlatBufferBuilder builder;
Array array;
EXPECT_THAT(DataBuffer::Serialize(array, &builder), HasOffset(0));
builder.Finish(::tflite::CreateTensor(builder));
auto* tensor =
flatbuffers::GetRoot<::tflite::Tensor>(builder.GetBufferPointer());
flatbuffers::FlatBufferBuilder buffer_builder;
Offset<Vector<uint8_t>> v = buffer_builder.CreateVector<uint8_t>({});
buffer_builder.Finish(::tflite::CreateBuffer(buffer_builder, v));
auto* buffer =
flatbuffers::GetRoot<::tflite::Buffer>(buffer_builder.GetBufferPointer());
DataBuffer::Deserialize(*tensor, *buffer, &array);
EXPECT_EQ(nullptr, array.buffer);
}
TEST(DataBuffer, UnsupportedTypes) {
for (ArrayDataType t : kUnsupportedTocoTypes) {
flatbuffers::FlatBufferBuilder builder;
Array array;
array.data_type = t;
array.GetMutableBuffer<ArrayDataType::kFloat>(); // This is OK.
EXPECT_DEATH(DataBuffer::Serialize(array, &builder),
"Unhandled array data type.");
}
for (::tflite::TensorType t : kUnsupportedTfLiteTypes) {
flatbuffers::FlatBufferBuilder builder;
builder.Finish(::tflite::CreateTensor(builder, 0, t, /*buffer*/ 1));
flatbuffers::FlatBufferBuilder buffer_builder;
Offset<Vector<uint8_t>> v = buffer_builder.CreateVector<uint8_t>({1});
buffer_builder.Finish(::tflite::CreateBuffer(buffer_builder, v));
auto* buffer = flatbuffers::GetRoot<::tflite::Buffer>(
buffer_builder.GetBufferPointer());
auto* tensor =
flatbuffers::GetRoot<::tflite::Tensor>(builder.GetBufferPointer());
Array array;
EXPECT_DEATH(DataBuffer::Deserialize(*tensor, *buffer, &array),
"Unhandled tensor type.");
}
}
TEST(DataBuffer, Float) {
Array recovered = ToFlatBufferAndBack<ArrayDataType::kFloat>({1.0f, 2.0f});
EXPECT_THAT(recovered.GetBuffer<ArrayDataType::kFloat>().data,
::testing::ElementsAre(1.0f, 2.0f));
}
TEST(DataBuffer, Uint8) {
Array recovered = ToFlatBufferAndBack<ArrayDataType::kUint8>({127, 244});
EXPECT_THAT(recovered.GetBuffer<ArrayDataType::kUint8>().data,
::testing::ElementsAre(127, 244));
}
TEST(DataBuffer, Int32) {
Array recovered = ToFlatBufferAndBack<ArrayDataType::kInt32>({1, 1 << 30});
EXPECT_THAT(recovered.GetBuffer<ArrayDataType::kInt32>().data,
::testing::ElementsAre(1, 1 << 30));
}
TEST(DataBuffer, Uint32) {
Array recovered = ToFlatBufferAndBack<ArrayDataType::kUint32>({1, 1U << 31});
EXPECT_THAT(recovered.GetBuffer<ArrayDataType::kUint32>().data,
::testing::ElementsAre(1, 1U << 31));
}
TEST(DataBuffer, Int16) {
Array recovered = ToFlatBufferAndBack<ArrayDataType::kInt16>({1, 1 << 14});
EXPECT_THAT(recovered.GetBuffer<ArrayDataType::kInt16>().data,
::testing::ElementsAre(1, 1 << 14));
}
TEST(DataBuffer, String) {
Array recovered = ToFlatBufferAndBack<ArrayDataType::kString>(
{"AA", "BBB", "Best. String. Ever."});
EXPECT_THAT(recovered.GetBuffer<ArrayDataType::kString>().data,
::testing::ElementsAre("AA", "BBB", "Best. String. Ever."));
}
TEST(DataBuffer, Bool) {
Array recovered =
ToFlatBufferAndBack<ArrayDataType::kBool>({true, false, true});
EXPECT_THAT(recovered.GetBuffer<ArrayDataType::kBool>().data,
::testing::ElementsAre(true, false, true));
}
TEST(DataBuffer, Complex64) {
Array recovered = ToFlatBufferAndBack<ArrayDataType::kComplex64>(
{std::complex<float>(1.0f, 2.0f), std::complex<float>(3.0f, 4.0f)});
EXPECT_THAT(recovered.GetBuffer<ArrayDataType::kComplex64>().data,
::testing::ElementsAre(std::complex<float>(1.0f, 2.0f),
std::complex<float>(3.0f, 4.0f)));
}
TEST(Padding, All) {
EXPECT_EQ(::tflite::Padding_SAME, Padding::Serialize(PaddingType::kSame));
EXPECT_EQ(PaddingType::kSame, Padding::Deserialize(::tflite::Padding_SAME));
EXPECT_EQ(::tflite::Padding_VALID, Padding::Serialize(PaddingType::kValid));
EXPECT_EQ(PaddingType::kValid, Padding::Deserialize(::tflite::Padding_VALID));
EXPECT_DEATH(Padding::Serialize(static_cast<PaddingType>(10000)),
"Unhandled padding type.");
EXPECT_DEATH(Padding::Deserialize(10000), "Unhandled padding.");
}
TEST(ActivationFunction, All) {
std::vector<
std::pair<FusedActivationFunctionType, ::tflite::ActivationFunctionType>>
testdata = {{FusedActivationFunctionType::kNone,
::tflite::ActivationFunctionType_NONE},
{FusedActivationFunctionType::kRelu,
::tflite::ActivationFunctionType_RELU},
{FusedActivationFunctionType::kRelu6,
::tflite::ActivationFunctionType_RELU6},
{FusedActivationFunctionType::kRelu1,
::tflite::ActivationFunctionType_RELU_N1_TO_1}};
for (auto x : testdata) {
EXPECT_EQ(x.second, ActivationFunction::Serialize(x.first));
EXPECT_EQ(x.first, ActivationFunction::Deserialize(x.second));
}
EXPECT_DEATH(ActivationFunction::Serialize(
static_cast<FusedActivationFunctionType>(10000)),
"Unhandled fused activation function type.");
EXPECT_DEATH(ActivationFunction::Deserialize(10000),
"Unhandled fused activation function type.");
}
} // namespace
} // namespace tflite
} // namespace toco