Skip to content

Commit

Permalink
[lite] Fix for signature def export in TFLiteConverter to handle case…
Browse files Browse the repository at this point in the history 
…s when output was deduplicated and multiple outputs point to same output buffer/tensor.

PiperOrigin-RevId: 366908259
Change-Id: I58b231e316419d89bd98a8cfb7a577c4e50f67fd
  • Loading branch information
@karimnosseir @tensorflower-gardener
karimnosseir authored and tensorflower-gardener committed Apr 6, 2021
1 parent be1238c commit 258528e
Show file tree
Hide file tree
Showing 2 changed files with 116 additions and 3 deletions.
12 changes: 9 additions & 3 deletions tensorflow/compiler/mlir/lite/flatbuffer_export.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1574,7 +1574,8 @@ std::vector<std::string> GetStringsFromDictionaryAttr(
}

std::vector<SignatureDefData> BuildSignaturedef(
FuncOp main_op, const std::string& saved_model_tag) {
FuncOp main_op, const std::string& saved_model_tag,
tensorflow::OpOrArgNameMapper& name_mapper) {
static const char kSignatureDefIndexPath[] = "tf_saved_model.index_path";
static const char kEntryFunctionAttributes[] = "tf.entry_function";

Expand Down Expand Up @@ -1638,7 +1639,11 @@ std::vector<SignatureDefData> BuildSignaturedef(
result[0].inputs[sig_def_inputs[i]] = input_names[i].str();
}
for (int i = 0; i < output_names.size(); ++i) {
result[0].outputs[sig_def_outputs[i]] = output_names[i].str();
// Fetch the name from the actual operand and not rely on names from
// outputs as deduping can make them invalid after conversion.
auto& operand = term->getOpOperand(i);
auto unique_name = std::string(name_mapper.GetUniqueName(operand.get()));
result[0].outputs[sig_def_outputs[i]] = unique_name;
}
if (auto name_attr = exported_name[0].dyn_cast_or_null<StringAttr>())
result[0].method_name = name_attr.getValue().str();
Expand Down Expand Up @@ -1830,7 +1835,8 @@ Optional<std::string> Translator::TranslateInternal() {
// Build SignatureDef
// We only have 1 entry point 'main' function, so build only 1 signature def.
auto main_fn_signature_def = BuildSignaturedef(
main_fn, saved_model_tags_.empty() ? "" : *saved_model_tags_.begin());
main_fn, saved_model_tags_.empty() ? "" : *saved_model_tags_.begin(),
name_mapper_);
auto signature_defs = CreateSignatureDefs(main_fn_signature_def);

auto model = tflite::CreateModel(builder_, TFLITE_SCHEMA_VERSION,
Expand Down
107 changes: 107 additions & 0 deletions tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/signature_def_output_override.mlir
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_to_string - | FileCheck %s

// CHECK: {
// CHECK-NEXT: version: 3,
// CHECK-NEXT: operator_codes: [ {
// CHECK-NEXT: deprecated_builtin_code: 9,
// CHECK-NEXT: version: 1,
// CHECK-NEXT: builtin_code: FULLY_CONNECTED
// CHECK-NEXT: } ],
// CHECK-NEXT: subgraphs: [ {
// CHECK-NEXT: tensors: [ {
// CHECK-NEXT: shape: [ 1, 3 ],
// CHECK-NEXT: buffer: 1,
// CHECK-NEXT: name: "serving_default_input2:0",
// CHECK-NEXT: quantization: {
// CHECK-EMPTY:
// CHECK-NEXT: },
// CHECK-NEXT: shape_signature: [ -1, 3 ]
// CHECK-NEXT: }, {
// CHECK-NEXT: shape: [ 1, 3 ],
// CHECK-NEXT: buffer: 2,
// CHECK-NEXT: name: "serving_default_input1:0",
// CHECK-NEXT: quantization: {
// CHECK-EMPTY:
// CHECK-NEXT: },
// CHECK-NEXT: shape_signature: [ -1, 3 ]
// CHECK-NEXT: }, {
// CHECK-NEXT: shape: [ 5 ],
// CHECK-NEXT: buffer: 3,
// CHECK-NEXT: name: "std.constant",
// CHECK-NEXT: quantization: {
// CHECK-EMPTY:
// CHECK-NEXT: }
// CHECK-NEXT: }, {
// CHECK-NEXT: shape: [ 5, 3 ],
// CHECK-NEXT: buffer: 4,
// CHECK-NEXT: name: "std.constant1",
// CHECK-NEXT: quantization: {
// CHECK-EMPTY:
// CHECK-NEXT: }
// CHECK-NEXT: }, {
// CHECK-NEXT: shape: [ 1, 5 ],
// CHECK-NEXT: buffer: 5,
// CHECK-NEXT: name: "StatefulPartitionedCall:1",
// CHECK-NEXT: quantization: {
// CHECK-EMPTY:
// CHECK-NEXT: },
// CHECK-NEXT: shape_signature: [ -1, 5 ]
// CHECK-NEXT: } ],
// CHECK-NEXT: inputs: [ 0, 1 ],
// CHECK-NEXT: outputs: [ 4, 4 ],
// CHECK-NEXT: operators: [ {
// CHECK-NEXT: inputs: [ 0, 3, 2 ],
// CHECK-NEXT: outputs: [ 4 ],
// CHECK-NEXT: builtin_options_type: FullyConnectedOptions,
// CHECK-NEXT: builtin_options: {
// CHECK-EMPTY:
// CHECK-NEXT: }
// CHECK-NEXT: } ],
// CHECK-NEXT: name: "main"
// CHECK-NEXT: } ],
// CHECK-NEXT: description: "MLIR Converted.",
// CHECK-NEXT: buffers: [ {
// CHECK-EMPTY:
// CHECK-NEXT: }, {
// CHECK-EMPTY:
// CHECK-NEXT: }, {
// CHECK-EMPTY:
// CHECK-NEXT: }, {
// CHECK-NEXT: data: [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
// CHECK-NEXT: }, {
// CHECK-NEXT: data: [ 0, 0, 128, 63, 0, 0, 128, 63, 0, 0, 128, 63, 0, 0, 128, 63, 0, 0, 128, 63, 0, 0, 128, 63, 0, 0, 128, 63, 0, 0, 128, 63, 0, 0, 128, 63, 0, 0, 128, 63, 0, 0, 128, 63, 0, 0, 128, 63, 0, 0, 128, 63, 0, 0, 128, 63, 0, 0, 128, 63 ]
// CHECK-NEXT: }, {
// CHECK-EMPTY:
// CHECK-NEXT: }, {
// CHECK-NEXT: data: [ 49, 46, 53, 46, 48, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
// CHECK-NEXT: } ],
// CHECK-NEXT: metadata: [ {
// CHECK-NEXT: name: "min_runtime_version",
// CHECK-NEXT: buffer: 6
// CHECK-NEXT: } ],
// CHECK-NEXT: signature_defs: [ {
// CHECK-NEXT: inputs: [ {
// CHECK-NEXT: name: "input1",
// CHECK-NEXT: tensor_index: 1
// CHECK-NEXT: }, {
// CHECK-NEXT: name: "input2"
// CHECK-NEXT: } ],
// CHECK-NEXT: outputs: [ {
// CHECK-NEXT: name: "end_logits",
// CHECK-NEXT: tensor_index: 4
// CHECK-NEXT: }, {
// CHECK-NEXT: name: "start_logits",
// CHECK-NEXT: tensor_index: 4
// CHECK-NEXT: } ],
// CHECK-NEXT: method_name: "serving_default",
// CHECK-NEXT: key: ""
// CHECK-NEXT: } ]
// CHECK-NEXT:}
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 554 : i32}, tf_saved_model.semantics} {
func @main(%arg0: tensor<?x3xf32> {tf_saved_model.index_path = ["input2"]}, %arg1: tensor<?x3xf32> {tf_saved_model.index_path = ["input1"]}) -> (tensor<?x5xf32> {tf_saved_model.index_path = ["start_logits"]}, tensor<?x5xf32> {tf_saved_model.index_path = ["end_logits"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input2:0,serving_default_input1:0", outputs = "StatefulPartitionedCall:1,StatefulPartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
%cst = constant dense<0.000000e+00> : tensor<5xf32>
%cst_0 = constant dense<1.0> : tensor<5x3xf32>
%0 = "tfl.fully_connected"(%arg0, %cst_0, %cst) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<?x3xf32>, tensor<5x3xf32>, tensor<5xf32>) -> tensor<?x5xf32>
return %0, %0 : tensor<?x5xf32>, tensor<?x5xf32>
}
}

0 comments on commit 258528e

Please sign in to comment.