From 99fc7367714e7c9efbc717160b10cece2ff39f44 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Tue, 7 May 2024 07:25:02 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 631415620 --- mediapipe/calculators/tensor/BUILD | 62 ++ .../tensor/inference_calculator.cc | 12 + .../calculators/tensor/inference_calculator.h | 3 + .../tensor/inference_calculator.proto | 20 + .../tensor/inference_calculator_cpu.cc | 4 +- .../tensor/inference_calculator_gl.cc | 1 + .../inference_calculator_gl_advanced.cc | 1 + .../tensor/inference_calculator_metal.cc | 1 + .../tensor/inference_calculator_xnnpack.cc | 7 +- .../tensor/inference_feedback_manager.cc | 191 +++++ .../tensor/inference_feedback_manager.h | 83 ++ .../tensor/inference_feedback_manager_test.cc | 710 ++++++++++++++++++ .../inference_interpreter_delegate_runner.cc | 79 +- .../inference_interpreter_delegate_runner.h | 5 +- .../calculators/tensor/inference_io_mapper.cc | 88 ++- .../calculators/tensor/inference_io_mapper.h | 1 + .../tensor/inference_io_mapper_test.cc | 36 +- .../calculators/tensor/inference_runner.h | 4 +- .../feedback_tensor_test_model.tflite | Bin 0 -> 1696 bytes ...edback_tensor_with_state_copy_model.tflite | Bin 0 -> 1172 bytes mediapipe/util/tflite/BUILD | 11 + mediapipe/util/tflite/utils.cc | 15 + mediapipe/util/tflite/utils.h | 25 + 23 files changed, 1308 insertions(+), 51 deletions(-) create mode 100644 mediapipe/calculators/tensor/inference_feedback_manager.cc create mode 100644 mediapipe/calculators/tensor/inference_feedback_manager.h create mode 100644 mediapipe/calculators/tensor/inference_feedback_manager_test.cc create mode 100644 mediapipe/calculators/tensor/testdata/feedback_tensor_test_model.tflite create mode 100644 mediapipe/calculators/tensor/testdata/feedback_tensor_with_state_copy_model.tflite create mode 100644 mediapipe/util/tflite/utils.cc create mode 100644 mediapipe/util/tflite/utils.h diff --git a/mediapipe/calculators/tensor/BUILD b/mediapipe/calculators/tensor/BUILD index 70496f2270..bda2ce8298 100644 --- a/mediapipe/calculators/tensor/BUILD +++ b/mediapipe/calculators/tensor/BUILD @@ -518,6 +518,66 @@ cc_test( ], ) +cc_library_with_tflite( + name = "inference_feedback_manager", + srcs = ["inference_feedback_manager.cc"], + hdrs = ["inference_feedback_manager.h"], + tflite_deps = [ + ":inference_io_mapper", + "//mediapipe/util/tflite:utils", + "//mediapipe/util/tflite:tflite_signature_reader", + "@org_tensorflow//tensorflow/lite:framework_stable", + "@org_tensorflow//tensorflow/lite:namespace", + "@org_tensorflow//tensorflow/lite/c:common", + "@org_tensorflow//tensorflow/lite/kernels:builtin_ops", + ], + deps = [ + ":inference_calculator_cc_proto", + ":inference_calculator_utils", + "//mediapipe/framework/port:ret_check", + "//mediapipe/framework/port:status", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/log:absl_log", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", + ], +) + +cc_test( + name = "inference_feedback_manager_test", + srcs = ["inference_feedback_manager_test.cc"], + data = [ + ":testdata/feedback_tensor_test_model.tflite", + ":testdata/feedback_tensor_with_state_copy_model.tflite", + ], + deps = [ + ":inference_calculator", + ":inference_calculator_cc_proto", + ":inference_calculator_cpu", + ":inference_calculator_interface", + ":inference_feedback_manager", + ":inference_io_mapper", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework/api2:packet", + "//mediapipe/framework/formats:tensor", + "//mediapipe/framework/port:gtest_main", + "//mediapipe/framework/port:parse_text_proto", + "//mediapipe/framework/port:status_matchers", + "//mediapipe/framework/tool:sink", + "//mediapipe/util/tflite:tflite_model_loader", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + "@org_tensorflow//tensorflow/lite:framework_stable", + "@org_tensorflow//tensorflow/lite/core:framework", + "@org_tensorflow//tensorflow/lite/core/api:op_resolver", + "@org_tensorflow//tensorflow/lite/delegates/xnnpack:xnnpack_delegate_hdrs_only", + "@org_tensorflow//tensorflow/lite/kernels:builtin_ops", + ], +) + cc_library( name = "inference_calculator_gl", srcs = ["inference_calculator_gl.cc"], @@ -645,6 +705,7 @@ cc_library_with_tflite( ":inference_runner", ":tflite_delegate_ptr", ":inference_io_mapper", + ":inference_feedback_manager", "//mediapipe/util/tflite:tflite_model_loader", "@org_tensorflow//tensorflow/lite:framework_stable", "@org_tensorflow//tensorflow/lite/c:c_api_types", @@ -658,6 +719,7 @@ cc_library_with_tflite( "//mediapipe/framework/formats:tensor", "//mediapipe/framework/port:ret_check", "//mediapipe/framework/port:status", + "@com_google_absl//absl/base:nullability", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@org_tensorflow//tensorflow/lite:string_util", diff --git a/mediapipe/calculators/tensor/inference_calculator.cc b/mediapipe/calculators/tensor/inference_calculator.cc index 677b46401e..bc8fc08086 100644 --- a/mediapipe/calculators/tensor/inference_calculator.cc +++ b/mediapipe/calculators/tensor/inference_calculator.cc @@ -119,5 +119,17 @@ InferenceCalculator::GetOpResolverAsPacket(CalculatorContext* cc) { tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates>()); } +void InferenceCalculator::WarnFeedbackTensorsUnsupported( + CalculatorContract* cc) { + const auto& options = cc->Options(); + if (options.has_input_output_config() && + !options.input_output_config().feedback_tensor_links().empty()) { + ABSL_LOG(WARNING) + << "Feedback tensor support is only available for CPU and " + << "XNNPACK inference. Ignoring " + "input_output_config.feedback_tensor_links option."; + } +} + } // namespace api2 } // namespace mediapipe diff --git a/mediapipe/calculators/tensor/inference_calculator.h b/mediapipe/calculators/tensor/inference_calculator.h index 219cc0fab5..b7d727f79b 100644 --- a/mediapipe/calculators/tensor/inference_calculator.h +++ b/mediapipe/calculators/tensor/inference_calculator.h @@ -157,6 +157,9 @@ class InferenceCalculator : public NodeIntf { static absl::StatusOr> GetOpResolverAsPacket( CalculatorContext* cc); + + // Checks if feedback tensor support is available and warns otherwise. + static void WarnFeedbackTensorsUnsupported(CalculatorContract* cc); }; struct InferenceCalculatorSelector : public InferenceCalculator { diff --git a/mediapipe/calculators/tensor/inference_calculator.proto b/mediapipe/calculators/tensor/inference_calculator.proto index f588eb9425..0e2d49b3ea 100644 --- a/mediapipe/calculators/tensor/inference_calculator.proto +++ b/mediapipe/calculators/tensor/inference_calculator.proto @@ -281,6 +281,26 @@ message InferenceCalculatorOptions { TensorIndicesMap output_tensor_indices_map = 2; TensorNamesMap output_tensor_names_map = 4; } + + // Feedback tensor links are pairs of model input / output tensors where + // the output should be set as inputs in the next model invocation. This + // allows to manage a notion of temporal state by continuously feeding from + // the model's output to the model's input during each inference step. Note + // that these feedback tensors must be excluded from the input/output + // tensor maps above as they are not used as regular inputs/outputs of the + // inference calculator. + message FeedbackTensorLink { + // TfLite output tensor name from default TfLite signature to use as + // source. + optional string from_output_tensor_name = 1; + // TfLite tensor name from default TfLitesignature to pass input + // tensor to. + optional string to_input_tensor_name = 2; + } + + // Defines a mapping between output tensors that should be + // used as input tensors during the next inference invocation. + repeated FeedbackTensorLink feedback_tensor_links = 5; } // Optionally remaps input and output tensors to align with TfLite model and diff --git a/mediapipe/calculators/tensor/inference_calculator_cpu.cc b/mediapipe/calculators/tensor/inference_calculator_cpu.cc index fb306228d2..16ce4782f4 100644 --- a/mediapipe/calculators/tensor/inference_calculator_cpu.cc +++ b/mediapipe/calculators/tensor/inference_calculator_cpu.cc @@ -89,12 +89,14 @@ absl::StatusOr> InferenceCalculatorCpuImpl::CreateInferenceRunner(CalculatorContext* cc) { MP_ASSIGN_OR_RETURN(auto model_packet, GetModelAsPacket(cc)); MP_ASSIGN_OR_RETURN(auto op_resolver_packet, GetOpResolverAsPacket(cc)); + const auto& options = cc->Options(); const int interpreter_num_threads = cc->Options().cpu_num_thread(); MP_ASSIGN_OR_RETURN(TfLiteDelegatePtr delegate, MaybeCreateDelegate(cc)); return CreateInferenceInterpreterDelegateRunner( std::move(model_packet), std::move(op_resolver_packet), - std::move(delegate), interpreter_num_threads); + std::move(delegate), interpreter_num_threads, + &options.input_output_config()); } absl::StatusOr diff --git a/mediapipe/calculators/tensor/inference_calculator_gl.cc b/mediapipe/calculators/tensor/inference_calculator_gl.cc index 82a9e8dbcc..cbf0cf6b4e 100644 --- a/mediapipe/calculators/tensor/inference_calculator_gl.cc +++ b/mediapipe/calculators/tensor/inference_calculator_gl.cc @@ -283,6 +283,7 @@ absl::Status InferenceCalculatorGlImpl::UpdateContract(CalculatorContract* cc) { RET_CHECK(!options.model_path().empty() ^ kSideInModel(cc).IsConnected()) << "Either model as side packet or model path in options is required."; + WarnFeedbackTensorsUnsupported(cc); return mediapipe::GlCalculatorHelper::UpdateContract(cc); } diff --git a/mediapipe/calculators/tensor/inference_calculator_gl_advanced.cc b/mediapipe/calculators/tensor/inference_calculator_gl_advanced.cc index b7e2fc4f9a..95368798f6 100644 --- a/mediapipe/calculators/tensor/inference_calculator_gl_advanced.cc +++ b/mediapipe/calculators/tensor/inference_calculator_gl_advanced.cc @@ -396,6 +396,7 @@ absl::Status InferenceCalculatorGlAdvancedImpl::UpdateContract( RET_CHECK(!options.model_path().empty() ^ kSideInModel(cc).IsConnected()) << "Either model as side packet or model path in options is required."; + WarnFeedbackTensorsUnsupported(cc); MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc)); return absl::OkStatus(); } diff --git a/mediapipe/calculators/tensor/inference_calculator_metal.cc b/mediapipe/calculators/tensor/inference_calculator_metal.cc index ebdc010060..6c9634509f 100644 --- a/mediapipe/calculators/tensor/inference_calculator_metal.cc +++ b/mediapipe/calculators/tensor/inference_calculator_metal.cc @@ -132,6 +132,7 @@ absl::Status InferenceCalculatorMetalImpl::UpdateContract( RET_CHECK(!options.model_path().empty() ^ kSideInModel(cc).IsConnected()) << "Either model as side packet or model path in options is required."; + WarnFeedbackTensorsUnsupported(cc); MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]); return absl::OkStatus(); } diff --git a/mediapipe/calculators/tensor/inference_calculator_xnnpack.cc b/mediapipe/calculators/tensor/inference_calculator_xnnpack.cc index 2a76443deb..61fb0970b0 100644 --- a/mediapipe/calculators/tensor/inference_calculator_xnnpack.cc +++ b/mediapipe/calculators/tensor/inference_calculator_xnnpack.cc @@ -86,12 +86,13 @@ absl::StatusOr> InferenceCalculatorXnnpackImpl::CreateInferenceRunner(CalculatorContext* cc) { MP_ASSIGN_OR_RETURN(auto model_packet, GetModelAsPacket(cc)); MP_ASSIGN_OR_RETURN(auto op_resolver_packet, GetOpResolverAsPacket(cc)); - const int interpreter_num_threads = - cc->Options().cpu_num_thread(); + const auto& options = cc->Options(); + const int interpreter_num_threads = options.cpu_num_thread(); MP_ASSIGN_OR_RETURN(TfLiteDelegatePtr delegate, CreateDelegate(cc)); return CreateInferenceInterpreterDelegateRunner( std::move(model_packet), std::move(op_resolver_packet), - std::move(delegate), interpreter_num_threads); + std::move(delegate), interpreter_num_threads, + &options.input_output_config()); } absl::StatusOr diff --git a/mediapipe/calculators/tensor/inference_feedback_manager.cc b/mediapipe/calculators/tensor/inference_feedback_manager.cc new file mode 100644 index 0000000000..42efb7425e --- /dev/null +++ b/mediapipe/calculators/tensor/inference_feedback_manager.cc @@ -0,0 +1,191 @@ +#include "mediapipe/calculators/tensor/inference_feedback_manager.h" + +#include +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "absl/log/absl_log.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_format.h" +#include "absl/strings/str_join.h" +#include "mediapipe/calculators/tensor/inference_calculator.pb.h" +#include "mediapipe/calculators/tensor/inference_io_mapper.h" +#include "mediapipe/framework/port/ret_check.h" +#include "mediapipe/framework/port/status_macros.h" +#include "mediapipe/util/tflite/tflite_signature_reader.h" +#include "mediapipe/util/tflite/utils.h" +#include "tensorflow/lite/c/common.h" +#include "tensorflow/lite/interpreter.h" + +namespace mediapipe { + +namespace { + +bool TfLiteTensorSpecEqual(const TfLiteTensor& a, const TfLiteTensor& b) { + return a.type == b.type && TfLiteIntArrayEqual(a.dims, b.dims) && + a.params.scale == b.params.scale && + a.params.zero_point == b.params.zero_point && + a.allocation_type == b.allocation_type && a.bytes == b.bytes; +} + +absl::flat_hash_map CreateNameToIndexMap( + const std::vector& names) { + absl::flat_hash_map name_to_index_map; + for (int i = 0; i < names.size(); ++i) { + name_to_index_map[names[i]] = i; + } + return name_to_index_map; +} + +} // namespace + +absl::Status InferenceFeedbackManager::Init( + const InferenceCalculatorOptions::InputOutputConfig& io_config, + const InputOutputTensorNames& input_output_tensor_names, + tflite::Interpreter* interpreter) { + interpreter_ = interpreter; + MP_ASSIGN_OR_RETURN(feedback_tensor_indices_links_, + ConvertSignatureTensorNamesToModelIndices( + io_config, input_output_tensor_names)); + + for (const auto& link : feedback_tensor_indices_links_) { + const auto [output_unused_iter, output_was_inserted] = + feedback_output_indices_.insert(link.from_idx); + RET_CHECK(output_was_inserted) << "Feedback output tensors must be unique."; + TfLiteTensor* from_tensor = + interpreter_->tensor(interpreter->outputs()[link.from_idx]); + RET_CHECK(!util::tflite::IsDynamicTensor(*from_tensor)) + << "Feedback output tensors must not be dynamic."; + const auto [input_unused_iter, input_was_inserted] = + feedback_input_indices_.insert(link.to_idx); + RET_CHECK(input_was_inserted) << "Feedback input tensors must be unique."; + TfLiteTensor* to_tensor = + interpreter_->tensor(interpreter->inputs()[link.to_idx]); + RET_CHECK(!util::tflite::IsDynamicTensor(*to_tensor)) + << "Feedback input tensors must not be dynamic."; + RET_CHECK(TfLiteTensorSpecEqual(*from_tensor, *to_tensor)) + << "Feedback tensors must have the same spec."; + // Since the TfLite API isn't specific about the initialization of newly + // allocated Tensor memory, we initialize the input to_tensor tensor with + // zeros. + memset(to_tensor->data.raw, 0, to_tensor->bytes); + } + + // Populate input_tensor_to_model_indices_ which maps InferenceRunner input + // tensors indices to the model input indices. + input_tensor_to_model_indices_.reserve(interpreter_->inputs().size()); + for (int i = 0; i < interpreter_->inputs().size(); ++i) { + if (!feedback_input_indices_.contains(i)) { + input_tensor_to_model_indices_.push_back(i); + } + } + return absl::OkStatus(); +} + +void InferenceFeedbackManager::SwapFeedbackTensors() { + for (const auto& link : feedback_tensor_indices_links_) { + TfLiteTensor* from_tensor = + interpreter_->tensor(interpreter_->outputs()[link.from_idx]); + TfLiteTensor* to_tensor = + interpreter_->tensor(interpreter_->inputs()[link.to_idx]); + { + using std::swap; + // TODO b/338023494 - Use TfLite CustomAllocator to manage memory of + // feedback tensors (replace std::swap) + swap(*from_tensor, *to_tensor); + } + } +} + +// static +absl::StatusOr> +InferenceFeedbackManager::ConvertSignatureTensorNamesToModelIndices( + const InferenceCalculatorOptions::InputOutputConfig& io_config, + const InputOutputTensorNames& input_output_tensor_names_map) { + std::vector indices_links; + if (input_output_tensor_names_map.empty() || + input_output_tensor_names_map.size() > 1) { + // Fail gracefully by returning an empty TensorFeedbackIndicesLink list if + // SignatureDef is not available or not supported. + ABSL_LOG(WARNING) + << "Feedback manager requires a model with a single signature " + "inference. Disabling support for feedback tensors."; + return indices_links; + } + // Obtain reference to single-signature in input_output_tensor_names_map. + const auto& input_output_tensor_names = + input_output_tensor_names_map.begin()->second; + + const auto input_name_to_index_map = + CreateNameToIndexMap(input_output_tensor_names.input_tensor_names); + const auto output_name_to_index_map = + CreateNameToIndexMap(input_output_tensor_names.output_tensor_names); + + // Create a set of all input/output tensor names used for InferenceCalculator + // I/O mapping. + absl::flat_hash_set input_output_mapping_tensor_names; + for (const auto& name : io_config.input_tensor_names_map().tensor_names()) { + input_output_mapping_tensor_names.insert(name); + } + for (const auto& name : io_config.output_tensor_names_map().tensor_names()) { + input_output_mapping_tensor_names.insert(name); + } + + for (const auto& link : io_config.feedback_tensor_links()) { + RET_CHECK(!input_output_mapping_tensor_names.contains( + link.from_output_tensor_name())) + << absl::StrFormat( + "Feedback output tensor [%s] cannot be used for input/output " + "mapping. Input/output mapping tensor names: [%s]", + link.from_output_tensor_name(), + absl::StrJoin(input_output_mapping_tensor_names, ", ")); + RET_CHECK(!input_output_mapping_tensor_names.contains( + link.to_input_tensor_name())) + << absl::StrFormat( + "Feedback input tensor [%s] cannot be used for input/output " + "mapping. Input/output mapping tensor names: [%s]", + link.to_input_tensor_name(), + absl::StrJoin(input_output_mapping_tensor_names, ", ")); + TensorFeedbackIndicesLink indices_link; + auto from_it = + output_name_to_index_map.find(link.from_output_tensor_name()); + RET_CHECK(from_it != output_name_to_index_map.end()) + << "Output tensor name not found: " << link.from_output_tensor_name(); + auto to_it = input_name_to_index_map.find(link.to_input_tensor_name()); + RET_CHECK(to_it != input_name_to_index_map.end()) + << "Input tensor name not found: " << link.to_input_tensor_name(); + indices_link.from_idx = from_it->second; + indices_link.to_idx = to_it->second; + indices_links.push_back(indices_link); + } + return indices_links; +} + +bool InferenceFeedbackManager::IsFeedbackInputTensorAtIndex(int idx) const { + return feedback_input_indices_.contains(idx); +} + +bool InferenceFeedbackManager::IsFeedbackOutputTensorAtIndex(int idx) const { + return feedback_output_indices_.contains(idx); +} + +absl::StatusOr InferenceFeedbackManager::MapInputTensorToModelIndex( + int input_idx) const { + RET_CHECK(input_idx >= 0 && + input_idx <= input_tensor_to_model_indices_.size()) + << "Invalid input tensor index: " << input_idx; + return input_tensor_to_model_indices_[input_idx]; +} + +int InferenceFeedbackManager::GetNumberOfNonFeedbackInputTensors() const { + return input_tensor_to_model_indices_.size(); +} + +int InferenceFeedbackManager::GetNumberOfFeedbackTensors() const { + return feedback_tensor_indices_links_.size(); +} +} // namespace mediapipe diff --git a/mediapipe/calculators/tensor/inference_feedback_manager.h b/mediapipe/calculators/tensor/inference_feedback_manager.h new file mode 100644 index 0000000000..a978d00685 --- /dev/null +++ b/mediapipe/calculators/tensor/inference_feedback_manager.h @@ -0,0 +1,83 @@ +#ifndef MEDIAPIPE_CALCULATORS_TENSOR_Inference_feedback_manager_H_ +#define MEDIAPIPE_CALCULATORS_TENSOR_Inference_feedback_manager_H_ + +#include + +#include "absl/container/flat_hash_set.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "mediapipe/calculators/tensor/inference_calculator.pb.h" +#include "mediapipe/calculators/tensor/inference_io_mapper.h" +#include "tensorflow/lite/interpreter.h" + +namespace mediapipe { + +// Feedback tensors are pairs of model output / input tensors where the +// model output is used as model input in the next model invocation. This allows +// to manage a notion of temporal state by continuously feeding from the model's +// output to the model's input during each inference step. The +// InferenceFeedbackManager initializes the feedback input tensors with zeros +// and efficiently swaps them from output to input with zero copies. +class InferenceFeedbackManager { + public: + // Initializes the feedback tensors with zeros and generates + // feedback_tensor_indices_links_. The provided interpreter must outlive the + // InferenceFeedbackManager instance. + absl::Status Init( + const mediapipe::InferenceCalculatorOptions::InputOutputConfig& io_config, + const InputOutputTensorNames& input_output_tensor_names_map, + tflite::Interpreter* interpreter); + + // Swaps the feedback tensors from output to input. + void SwapFeedbackTensors(); + + // Returns the number of expected non-feedback tensors. This can be used to + // confirm the number of input tensors to the InferenceRunner implementation. + int GetNumberOfNonFeedbackInputTensors() const; + + // Returns the number of feedback tensor pairs. + int GetNumberOfFeedbackTensors() const; + + // Since feedback tensors are excluded from InferenceRunner input, This method + // maps the tensor index from the InferenceRunner input to the TfLite model + // tensor input. + absl::StatusOr MapInputTensorToModelIndex(int input_idx) const; + + // Returns true if the tensor at the given index is a feedback input tensor. + bool IsFeedbackInputTensorAtIndex(int idx) const; + + // Returns true if the tensor at the given index is a feedback output tensor. + bool IsFeedbackOutputTensorAtIndex(int idx) const; + + private: + // Links between feedback tensors defined by model tensor indices. + struct TensorFeedbackIndicesLink { + int from_idx; + int to_idx; + }; + + // Translates the tensor names from the input/output config into the + // corresponding TfLite tensor indices. + static absl::StatusOr> + ConvertSignatureTensorNamesToModelIndices( + const mediapipe::InferenceCalculatorOptions::InputOutputConfig& io_config, + const InputOutputTensorNames& input_output_tensor_names_map); + + // Non-owning reference to the TfLite interpreter. + tflite::Interpreter* interpreter_ = nullptr; + + // List of tensor feedback pairs defined by model tensor indices. + std::vector feedback_tensor_indices_links_; + + // Maps InferenceRunner input indices to TfLiteModel input indices. + std::vector input_tensor_to_model_indices_; + + // Set of feedback input model tensor indices. + absl::flat_hash_set feedback_input_indices_; + + // Set of feedback output model tensor indices. + absl::flat_hash_set feedback_output_indices_; +}; +} // namespace mediapipe + +#endif // MEDIAPIPE_CALCULATORS_TENSOR_Inference_feedback_manager_H_ diff --git a/mediapipe/calculators/tensor/inference_feedback_manager_test.cc b/mediapipe/calculators/tensor/inference_feedback_manager_test.cc new file mode 100644 index 0000000000..36f74f4070 --- /dev/null +++ b/mediapipe/calculators/tensor/inference_feedback_manager_test.cc @@ -0,0 +1,710 @@ +// Copyright 2024 The MediaPipe Authors. +// +// 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 "mediapipe/calculators/tensor/inference_feedback_manager.h" + +#include +#include +#include +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/strings/str_replace.h" +#include "mediapipe/calculators/tensor/inference_calculator.pb.h" +#include "mediapipe/calculators/tensor/inference_io_mapper.h" +#include "mediapipe/framework/api2/packet.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/tensor.h" +#include "mediapipe/framework/port/gmock.h" +#include "mediapipe/framework/port/gtest.h" +#include "mediapipe/framework/port/parse_text_proto.h" +#include "mediapipe/framework/port/status_matchers.h" +#include "mediapipe/util/tflite/tflite_model_loader.h" +#include "tensorflow/lite/core/api/op_resolver.h" +#include "tensorflow/lite/core/interpreter_builder.h" +#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h" +#include "tensorflow/lite/interpreter.h" +#include "tensorflow/lite/kernels/register.h" + +namespace mediapipe { +namespace api2 { +namespace { + +// feedback_tensor_test_model.iflite model passes through stateless/non-feedback +// tensors and increments "stateful" tensors by one during inference: +// +// class Wrapper(tf.Module): +// @tf.function(input_signature= +// (tf.TensorSpec(shape=[1,1], +// dtype=tf.float32, +// name="regular_input0"), +// tf.TensorSpec(shape=[1,1], +// dtype=tf.float32, +// name="feedback_float_input"), +// tf.TensorSpec(shape=[1,2], +// dtype=tf.float32, +// name="regular_input1"), +// tf.TensorSpec(shape=[1,1], dtype=tf.int32, +// name="feedback_int_input"))) +// def model(self, ...): +// return {"regular_output0": regular_input0, +// "feedback_incremented_float_output": feedback_float_input + 1, +// "regular_output1": regular_input1, +// "feedback_incremented_int_output": feedback_int_input + 1} +// +constexpr char kFeedbackTestModelPath[] = + "mediapipe/calculators/tensor/testdata/" + "feedback_tensor_test_model.tflite"; + +// feedback_tensor_with_state_copy_model.tflite model passes through +// stateless/non-feedback tensors and increments "stateful" tensors by one. It +// also copies the stateful tensor to a second tensor to enable to observe its +// state. +// +// class Wrapper(tf.Module): +// @tf.function(input_signature=(tf.TensorSpec(shape=[1,1], +// dtype=tf.int32, +// name="regular_int_input"), +// tf.TensorSpec(shape=[1,1], +// dtype=tf.int32, +// name="feedback_int_input"))) +// def model(self, regular_int_input, feedback_int_input): +// return {"regular_int_output": regular_int_input, +// "feedback_incremented_int_output": feedback_int_input + 1, +// "feedback_incremented_int_copy": feedback_int_input + 0} +constexpr char kFeedbackTestWithStateCopyModelPath[] = + "mediapipe/calculators/tensor/testdata/" + "feedback_tensor_with_state_copy_model.tflite"; + +using ::mediapipe::Packet; +using ::mediapipe::tool::AddVectorSink; +using ::testing::HasSubstr; +using ::tflite::InterpreterBuilder; +using TfLiteDelegatePtr = + std::unique_ptr>; + +static Tensor CreateSingleIntTensor(int value) { + std::vector dims = {1, 1}; + Tensor tensor(Tensor::ElementType::kInt32, Tensor::Shape(dims)); + auto write_view = tensor.GetCpuWriteView(); + *write_view.buffer() = value; + return tensor; +} + +class InferenceFeedbackManagerTest : public ::testing::Test { + protected: + void InitModelAndInterpreter(const std::string& model_path) { + MP_ASSERT_OK_AND_ASSIGN(model_, + TfLiteModelLoader::LoadFromPath(model_path)); + op_resolver_ = std::make_unique( + tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates()); + InterpreterBuilder builder(*model_.Get(), *op_resolver_); + auto xnnpack_opts = TfLiteXNNPackDelegateOptionsDefault(); + xnnpack_opts.num_threads = 1; + delegate_ = TfLiteDelegatePtr(TfLiteXNNPackDelegateCreate(&xnnpack_opts), + &TfLiteXNNPackDelegateDelete); + builder.AddDelegate(delegate_.get()); + builder.SetNumThreads(1); + ASSERT_EQ(builder(&interpreter_), kTfLiteOk); + ASSERT_NE(interpreter_, nullptr); + ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk); + } + + // Helper methods to access tensors of feedback_tensor_test_model.tflite + + // ~~~~~~~~~~ INPUTS ~~~~~~~~~~ + // 0 : regular_input0 : [1 1] : F32 + // 1 : feedback_float_input : [1 1] : F32 + // 2 : feedback_int_input : [1 1] : I32 + // 3 : regular_input1 : [1 2] : F32 + // ~~~~~~~~~~ OUTPUTS ~~~~~~~~~ + // 0 : feedback_incremented_float_output : [1 1] : F32 + // 1 : regular_output1 : [1 2] : F32 + // 2 : feedback_incremented_int_output : [1 1] : I32 + // 3 : regular_output0 : [1 1] : F32 + void PopulateRegularSingleFloatInputTensor(float value) { + std::vector input_buffer = {value}; + CopyTensorBufferToInterpreter(input_buffer, /*input_tensor_index=*/0); + } + + void PopulateFeedbackFloatInputTensor(float value) { + std::vector input_buffer = {value}; + CopyTensorBufferToInterpreter(input_buffer, /*input_tensor_index=*/1); + } + + void PopulateFeedbackIntInputTensor(int value) { + std::vector input_buffer = {value}; + CopyTensorBufferToInterpreter(input_buffer, /*input_tensor_index=*/2); + } + + void PopulateRegularTwoFloatInputTensor(std::vector input_buffer) { + CopyTensorBufferToInterpreter(input_buffer, /*input_tensor_index=*/3); + } + + void RunInference() { ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk); } + + float GetFeedbackFloatOutput() { + std::vector output_buffer = CopyTensorBufferFromInterpreter( + /*output_tensor_index=*/0, /*num_elements=*/1); + return output_buffer[0]; + } + + std::vector GetRegularTwoFloatsOutput() { + return CopyTensorBufferFromInterpreter(/*output_tensor_index=*/1, + /*num_elements=*/2); + } + int GetFeedbackIntOutput() { + std::vector output_buffer = CopyTensorBufferFromInterpreter( + /*output_tensor_index=*/2, /*num_elements=*/1); + return output_buffer[0]; + } + + float GetRegularSingleFloatOutput() { + std::vector output_buffer = CopyTensorBufferFromInterpreter( + /*output_tensor_index=*/3, /*num_elements=*/1); + return output_buffer[0]; + } + + api2::Packet model_; + std::unique_ptr op_resolver_; + TfLiteDelegatePtr delegate_; + std::unique_ptr interpreter_; + + private: + template + void CopyTensorBufferToInterpreter(const std::vector& input_buffer, + int input_tensor_index) { + EXPECT_LT(input_tensor_index, interpreter_->inputs().size()); + T* local_tensor_buffer = + interpreter_->typed_input_tensor(input_tensor_index); + std::memcpy(local_tensor_buffer, + static_cast(input_buffer.data()), + input_buffer.size() * sizeof(T)); + } + + template + std::vector CopyTensorBufferFromInterpreter(int output_tensor_index, + int num_elements) { + EXPECT_LT(output_tensor_index, interpreter_->outputs().size()); + std::vector result(num_elements); + T* local_tensor_buffer = + interpreter_->typed_output_tensor(output_tensor_index); + std::memcpy(static_cast(result.data()), local_tensor_buffer, + result.size() * sizeof(T)); + return result; + } +}; + +TEST_F(InferenceFeedbackManagerTest, ModelShouldIncreaseStatefulTensorsByOne) { + // Test the test model. + InitModelAndInterpreter(kFeedbackTestModelPath); + // Initialize "stateful" input tensors with zero values. + PopulateFeedbackFloatInputTensor(0); + PopulateRegularTwoFloatInputTensor({1.0f, 2.0f}); + PopulateFeedbackIntInputTensor(0); + PopulateRegularSingleFloatInputTensor(3.14f); + RunInference(); + EXPECT_EQ(GetFeedbackIntOutput(), 1); + EXPECT_FLOAT_EQ(GetFeedbackFloatOutput(), 1.0f); + EXPECT_FLOAT_EQ(GetRegularSingleFloatOutput(), 3.14f); + EXPECT_THAT(GetRegularTwoFloatsOutput(), testing::ElementsAre(1.0f, 2.0f)); +} + +TEST_F(InferenceFeedbackManagerTest, + ShouldInitializeFeedbackTensorInputWithZeros) { + InitModelAndInterpreter(kFeedbackTestModelPath); + + constexpr int kFeedbackFloatInputIndex = 1; + constexpr int kFeedbackIntInputIndex = 2; + EXPECT_STREQ(interpreter_->GetInputName(kFeedbackFloatInputIndex), + "serving_default_feedback_float_input:0"); + EXPECT_STREQ(interpreter_->GetInputName(kFeedbackIntInputIndex), + "serving_default_feedback_int_input:0"); + // Initialize "stateful" input tensors with non-zero values. + *interpreter_->typed_input_tensor(kFeedbackFloatInputIndex) = 123.0f; + *interpreter_->typed_input_tensor(kFeedbackIntInputIndex) = 123; + + InferenceFeedbackManager feedback_manager; + InferenceCalculatorOptions::InputOutputConfig config = + ParseTextProtoOrDie(R"pb( + feedback_tensor_links { + from_output_tensor_name: "feedback_incremented_float_output", + to_input_tensor_name: "feedback_float_input", + } + feedback_tensor_links { + from_output_tensor_name: "feedback_incremented_int_output", + to_input_tensor_name: "feedback_int_input", + } + )pb"); + + MP_ASSERT_OK_AND_ASSIGN( + const auto input_output_tensor_names, + InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( + *interpreter_)); + MP_ASSERT_OK(feedback_manager.Init(config, input_output_tensor_names, + interpreter_.get())); + + EXPECT_EQ(*interpreter_->typed_input_tensor(kFeedbackFloatInputIndex), + 0.0f); + EXPECT_EQ(*interpreter_->typed_input_tensor(kFeedbackIntInputIndex), 0); +} + +TEST_F(InferenceFeedbackManagerTest, + ShouldAllowToQueryForFeedbackTensorIndices) { + InitModelAndInterpreter(kFeedbackTestModelPath); + + // ~~~~~~~~~~ INPUTS ~~~~~~~~~~ + // 0 : regular_input0 : [1 1] : F32 + // 1 : feedback_float_input : [1 1] : F32 + // 2 : feedback_int_input : [1 1] : I32 + // 3 : regular_input1 : [1 2] : F32 + // ~~~~~~~~~~ OUTPUTS ~~~~~~~~~ + // 0 : feedback_incremented_float_output : [1 1] : F32 + // 1 : regular_output1 : [1 2] : F32 + // 2 : feedback_incremented_int_output : [1 1] : I32 + // 3 : regular_output0 : [1 1] : F32 + + // Confirm input signatures. + constexpr int kRegularInput0Index = 0; + constexpr int kFeedbackFloatInputIndex = 1; + constexpr int kFeedbackIntInputIndex = 2; + constexpr int kRegularInput1Index = 3; + + // Confirm output signatures. + constexpr int kFeedbackIncrementedFloatOutputIndex = 0; + constexpr int kRegularOutput0Index = 1; + constexpr int kFeedbackIncrementedIntOutputIndex = 2; + constexpr int kRegularOutput1Index = 3; + + InferenceFeedbackManager feedback_manager; + InferenceCalculatorOptions::InputOutputConfig config = + ParseTextProtoOrDie(R"pb( + feedback_tensor_links { + from_output_tensor_name: "feedback_incremented_float_output", + to_input_tensor_name: "feedback_float_input", + } + feedback_tensor_links { + from_output_tensor_name: "feedback_incremented_int_output", + to_input_tensor_name: "feedback_int_input", + } + )pb"); + MP_ASSERT_OK_AND_ASSIGN( + const auto input_output_tensor_names, + InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( + *interpreter_)); + MP_ASSERT_OK(feedback_manager.Init(config, input_output_tensor_names, + interpreter_.get())); + + EXPECT_TRUE( + feedback_manager.IsFeedbackInputTensorAtIndex(kFeedbackFloatInputIndex)); + EXPECT_TRUE( + feedback_manager.IsFeedbackInputTensorAtIndex(kFeedbackIntInputIndex)); + EXPECT_FALSE( + feedback_manager.IsFeedbackInputTensorAtIndex(kRegularInput0Index)); + EXPECT_FALSE( + feedback_manager.IsFeedbackInputTensorAtIndex(kRegularInput1Index)); + + EXPECT_TRUE(feedback_manager.IsFeedbackOutputTensorAtIndex( + kFeedbackIncrementedIntOutputIndex)); + EXPECT_TRUE(feedback_manager.IsFeedbackOutputTensorAtIndex( + kFeedbackIncrementedFloatOutputIndex)); + EXPECT_FALSE( + feedback_manager.IsFeedbackOutputTensorAtIndex(kRegularOutput0Index)); + EXPECT_FALSE( + feedback_manager.IsFeedbackOutputTensorAtIndex(kRegularOutput1Index)); +} + +TEST_F(InferenceFeedbackManagerTest, ShouldMapInputTensorToModelTensorIndices) { + InitModelAndInterpreter(kFeedbackTestModelPath); + + // First two input tensors are stateful / feedback tensors. + constexpr int kRegularInput0Index = 0; + constexpr int kRegularInput1Index = 3; + + InferenceFeedbackManager feedback_manager; + InferenceCalculatorOptions::InputOutputConfig config = + ParseTextProtoOrDie(R"pb( + feedback_tensor_links { + from_output_tensor_name: "feedback_incremented_float_output", + to_input_tensor_name: "feedback_float_input", + } + feedback_tensor_links { + from_output_tensor_name: "feedback_incremented_int_output", + to_input_tensor_name: "feedback_int_input", + } + )pb"); + MP_ASSERT_OK_AND_ASSIGN( + const auto input_output_tensor_names, + InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( + *interpreter_)); + MP_ASSERT_OK(feedback_manager.Init(config, input_output_tensor_names, + interpreter_.get())); + + // Feedback tensors are skipped in InferenceRunner input. Therefore the first + // two InferenceRunner input tensors must point to the first two non-feedback + // model input tensors. + MP_ASSERT_OK_AND_ASSIGN(const int inference_runner_index0, + feedback_manager.MapInputTensorToModelIndex(0)); + MP_ASSERT_OK_AND_ASSIGN(const int inference_runner_index1, + feedback_manager.MapInputTensorToModelIndex(1)); + EXPECT_EQ(inference_runner_index0, kRegularInput0Index); + EXPECT_EQ(inference_runner_index1, kRegularInput1Index); +} + +TEST_F(InferenceFeedbackManagerTest, ShouldDetectLinksWithDifferentTypes) { + InitModelAndInterpreter(kFeedbackTestModelPath); + + constexpr int kFeedbackFloatInputIndex = 1; + EXPECT_STREQ(interpreter_->GetInputName(kFeedbackFloatInputIndex), + "serving_default_feedback_float_input:0"); + + InferenceFeedbackManager feedback_manager; + InferenceCalculatorOptions::InputOutputConfig config = + ParseTextProtoOrDie(R"pb( + feedback_tensor_links { + from_output_tensor_name: "feedback_incremented_int_output", + to_input_tensor_name: "feedback_float_input", + } + )pb"); + + MP_ASSERT_OK_AND_ASSIGN( + const auto input_output_tensor_names, + InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( + *interpreter_)); + EXPECT_THAT(feedback_manager.Init(config, input_output_tensor_names, + interpreter_.get()), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("Feedback tensors must have the same spec"))); +} + +TEST_F(InferenceFeedbackManagerTest, ShouldDetectDynamicInputTensors) { + InitModelAndInterpreter(kFeedbackTestModelPath); + + InferenceFeedbackManager feedback_manager; + InferenceCalculatorOptions::InputOutputConfig config = + ParseTextProtoOrDie(R"pb( + feedback_tensor_links { + from_output_tensor_name: "feedback_incremented_float_output", + to_input_tensor_name: "feedback_float_input", + } + )pb"); + + // Mark feedback tensor as dynamic by setting one dimension to -1. + interpreter_->tensor(interpreter_->inputs()[1])->dims->data[0] = -1; + + MP_ASSERT_OK_AND_ASSIGN( + const auto input_output_tensor_names, + InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( + *interpreter_)); + EXPECT_THAT( + feedback_manager.Init(config, input_output_tensor_names, + interpreter_.get()), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("Feedback input tensors must not be dynamic"))); +} + +TEST_F(InferenceFeedbackManagerTest, ShouldDetectDynamicOutputTensors) { + InitModelAndInterpreter(kFeedbackTestModelPath); + + InferenceFeedbackManager feedback_manager; + InferenceCalculatorOptions::InputOutputConfig config = + ParseTextProtoOrDie(R"pb( + feedback_tensor_links { + from_output_tensor_name: "feedback_incremented_float_output", + to_input_tensor_name: "feedback_float_input", + } + )pb"); + + // Mark feedback tensor as dynamic by setting one dimension to -1. + interpreter_->tensor(interpreter_->outputs()[0])->dims->data[0] = -1; + + MP_ASSERT_OK_AND_ASSIGN( + const auto input_output_tensor_names, + InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( + *interpreter_)); + EXPECT_THAT( + feedback_manager.Init(config, input_output_tensor_names, + interpreter_.get()), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("Feedback output tensors must not be dynamic"))); +} + +TEST_F(InferenceFeedbackManagerTest, ShouldDetectLinksWithDifferentDimensions) { + InitModelAndInterpreter(kFeedbackTestModelPath); + + constexpr int kFeedbackFloatInputIndex = 1; + EXPECT_STREQ(interpreter_->GetInputName(kFeedbackFloatInputIndex), + "serving_default_feedback_float_input:0"); + + InferenceFeedbackManager feedback_manager; + InferenceCalculatorOptions::InputOutputConfig config = + ParseTextProtoOrDie(R"pb( + feedback_tensor_links { + from_output_tensor_name: "regular_output1", # has dimension [1 2] + to_input_tensor_name: "feedback_float_input", # has dimension [1 1] + } + )pb"); + + MP_ASSERT_OK_AND_ASSIGN( + const auto input_output_tensor_names, + InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( + *interpreter_)); + EXPECT_THAT(feedback_manager.Init(config, input_output_tensor_names, + interpreter_.get()), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("Feedback tensors must have the same spec"))); +} + +TEST_F(InferenceFeedbackManagerTest, + ShouldDetectMismatchBetweenStatefulAndRegularTensors) { + InitModelAndInterpreter(kFeedbackTestModelPath); + + InferenceFeedbackManager feedback_manager; + InferenceCalculatorOptions::InputOutputConfig config = + ParseTextProtoOrDie(R"pb( + input_tensor_names_map { tensor_names: "feedback_float_input" } + feedback_tensor_links { + from_output_tensor_name: "regular_output1", # has dimension [1 2] + to_input_tensor_name: "feedback_float_input", # has dimension [1 1] + } + )pb"); + + MP_ASSERT_OK_AND_ASSIGN( + const auto input_output_tensor_names, + InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( + *interpreter_)); + EXPECT_THAT(feedback_manager.Init(config, input_output_tensor_names, + interpreter_.get()), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("Feedback input tensor [feedback_float_input] " + "cannot be used for input/output mapping"))); +} + +TEST_F(InferenceFeedbackManagerTest, ShouldSwapFeedbackTensors) { + InitModelAndInterpreter(kFeedbackTestModelPath); + InferenceFeedbackManager feedback_manager; + InferenceCalculatorOptions::InputOutputConfig config = + ParseTextProtoOrDie(R"pb( + feedback_tensor_links { + from_output_tensor_name: "feedback_incremented_float_output", + to_input_tensor_name: "feedback_float_input", + } + feedback_tensor_links { + from_output_tensor_name: "feedback_incremented_int_output", + to_input_tensor_name: "feedback_int_input", + } + )pb"); + MP_ASSERT_OK_AND_ASSIGN( + const auto input_output_tensor_names, + InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( + *interpreter_)); + MP_ASSERT_OK(feedback_manager.Init(config, input_output_tensor_names, + interpreter_.get())); + // Initialize "stateful" input tensors with zero values. + PopulateFeedbackFloatInputTensor(0); + PopulateFeedbackIntInputTensor(0); + PopulateRegularSingleFloatInputTensor(3.14f); + PopulateRegularTwoFloatInputTensor({1.0f, 2.0f}); + RunInference(); + EXPECT_EQ(GetFeedbackIntOutput(), 1); + EXPECT_FLOAT_EQ(GetFeedbackFloatOutput(), 1.0f); + EXPECT_FLOAT_EQ(GetRegularSingleFloatOutput(), 3.14f); + EXPECT_THAT(GetRegularTwoFloatsOutput(), testing::ElementsAre(1.0f, 2.0f)); + + feedback_manager.SwapFeedbackTensors(); + + RunInference(); + EXPECT_EQ(GetFeedbackIntOutput(), 2); + EXPECT_FLOAT_EQ(GetFeedbackFloatOutput(), 2.0f); + EXPECT_FLOAT_EQ(GetRegularSingleFloatOutput(), 3.14f); + EXPECT_THAT(GetRegularTwoFloatsOutput(), testing::ElementsAre(1.0f, 2.0f)); +} + +TEST_F(InferenceFeedbackManagerTest, ShouldRunE2ESmokeTest) { + CalculatorGraph graph; + CalculatorGraphConfig graph_config = + ParseTextProtoOrDie(absl::StrReplaceAll( + R"pb( + input_stream: "regular_int_input" + input_stream: "feedback_int_input" + output_stream: "regular_int_output" + output_stream: "feedback_incremented_int_copy" + node { + calculator: "InferenceCalculator" + # ~~~~~~~~~~ INPUTS ~~~~~~~~~~ + # 0 : feedback_int_input : [1 1] : I32 + # 1 : regular_int_input : [1 1] : I32 + # ~~~~~~~~~~ OUTPUTS ~~~~~~~~~ + # 0 : feedback_incremented_int_copy : [1 1] : I32 + # 1 : regular_int_output : [1 1] : I32 + # 2 : feedback_incremented_int_output : [1 1] : I32 + # (copy of feedback_incremented_int_output) + input_stream: "TENSOR:0:regular_int_input" + output_stream: "TENSOR:0:feedback_incremented_int_copy" + output_stream: "TENSOR:1:regular_int_output" + options { + [mediapipe.InferenceCalculatorOptions.ext] { + model_path: "$model" + delegate {} # empty delegate message enables CPU inference. + input_output_config { + input_tensor_names_map { tensor_names: "regular_int_input" } + output_tensor_names_map { + tensor_names: "feedback_incremented_int_copy" + tensor_names: "regular_int_output" + } + feedback_tensor_links { + from_output_tensor_name: "feedback_incremented_int_output" + to_input_tensor_name: "feedback_int_input" + } + } + } + } + } + )pb", + {{"$model", kFeedbackTestWithStateCopyModelPath}})); + + std::vector regular_int_output; + AddVectorSink("regular_int_output", &graph_config, ®ular_int_output); + std::vector feedback_incremented_int_copy; + AddVectorSink("feedback_incremented_int_copy", &graph_config, + &feedback_incremented_int_copy); + + MP_ASSERT_OK(graph.Initialize(graph_config)); + MP_ASSERT_OK(graph.StartRun({})); + + const std::vector kRegularInputTensorValues = {100, 200, 300}; + // Simulate 3 inference steps. + for (int n = 0; n < 3; ++n) { + Tensor input_tensor = CreateSingleIntTensor(kRegularInputTensorValues[n]); + MP_ASSERT_OK(graph.AddPacketToInputStream( + "regular_int_input", + mediapipe::MakePacket(std::move(input_tensor)) + .At(Timestamp(n)))); + } + MP_ASSERT_OK(graph.CloseAllInputStreams()); + MP_ASSERT_OK(graph.WaitUntilDone()); + + EXPECT_EQ(regular_int_output.size(), 3); + EXPECT_EQ(feedback_incremented_int_copy.size(), 3); + for (int i = 0; i < regular_int_output.size(); ++i) { + const auto regular_read_view = + regular_int_output[i].Get().GetCpuReadView(); + EXPECT_EQ(regular_read_view.buffer()[0], kRegularInputTensorValues[i]); + + // Stateful tensor are initialized with zero and incremented by one in + // every iteration. + const auto feedback_read_view = + feedback_incremented_int_copy[i].Get().GetCpuReadView(); + EXPECT_EQ(feedback_read_view.buffer()[0], i); + } +} + +TEST_F(InferenceFeedbackManagerTest, + ShouldRunE2EWithWithoutFeedbackManagerConfigSmokeTest) { + CalculatorGraph graph; + CalculatorGraphConfig graph_config = + ParseTextProtoOrDie(absl::StrReplaceAll( + R"pb( + input_stream: "regular_int_input" + input_stream: "feedback_int_input" + output_stream: "regular_int_output" + output_stream: "feedback_incremented_int_output" + output_stream: "feedback_incremented_int_copy" + node { + calculator: "InferenceCalculator" + # ~~~~~~~~~~ INPUTS ~~~~~~~~~~ + # 0 : feedback_int_input : [1 1] : I32 + # 1 : regular_int_input : [1 1] : I32 + # ~~~~~~~~~~ OUTPUTS ~~~~~~~~~ + # 0 : feedback_incremented_int_copy : [1 1] : I32 + # 1 : regular_int_output : [1 1] : I32 + # 2 : feedback_incremented_int_output : [1 1] : I32 + # (copy of feedback_incremented_int_output) + input_stream: "TENSOR:0:feedback_int_input" + input_stream: "TENSOR:1:regular_int_input" + output_stream: "TENSOR:0:feedback_incremented_int_copy" + output_stream: "TENSOR:1:regular_int_output" + output_stream: "TENSOR:2:feedback_incremented_int_output" + options { + [mediapipe.InferenceCalculatorOptions.ext] { + model_path: "$model" + delegate {} # empty delegate message enables CPU inference. + } + } + } + )pb", + {{"$model", kFeedbackTestWithStateCopyModelPath}})); + + std::vector regular_int_output; + AddVectorSink("regular_int_output", &graph_config, ®ular_int_output); + std::vector feedback_incremented_int_output; + AddVectorSink("feedback_incremented_int_output", &graph_config, + &feedback_incremented_int_output); + std::vector feedback_incremented_int_copy; + AddVectorSink("feedback_incremented_int_copy", &graph_config, + &feedback_incremented_int_copy); + + MP_ASSERT_OK(graph.Initialize(graph_config)); + MP_ASSERT_OK(graph.StartRun({})); + + const std::vector kRegularInputTensorValues = {100, 200, 300}; + const std::vector kFeedbackInputTensorValues = {111, 222, 333}; + // Simulate 3 inference steps. + for (int n = 0; n < 3; ++n) { + Tensor regular_int_input_tensor = + CreateSingleIntTensor(kRegularInputTensorValues[n]); + MP_ASSERT_OK(graph.AddPacketToInputStream( + "regular_int_input", + mediapipe::MakePacket(std::move(regular_int_input_tensor)) + .At(Timestamp(n)))); + Tensor feedback_int_input_tensor = + CreateSingleIntTensor(kFeedbackInputTensorValues[n]); + MP_ASSERT_OK(graph.AddPacketToInputStream( + "feedback_int_input", + mediapipe::MakePacket(std::move(feedback_int_input_tensor)) + .At(Timestamp(n)))); + } + + MP_ASSERT_OK(graph.CloseAllInputStreams()); + MP_ASSERT_OK(graph.WaitUntilDone()); + + EXPECT_EQ(regular_int_output.size(), 3); + EXPECT_EQ(feedback_incremented_int_copy.size(), 3); + EXPECT_EQ(feedback_incremented_int_output.size(), 3); + for (int i = 0; i < regular_int_output.size(); ++i) { + const auto regular_read_view = + regular_int_output[i].Get().GetCpuReadView(); + EXPECT_EQ(regular_read_view.buffer()[0], kRegularInputTensorValues[i]); + + // Stateful tensor are initialized with zero and incremented by one in + // every iteration. + { + const auto read_view = + feedback_incremented_int_output[i].Get().GetCpuReadView(); + EXPECT_EQ(read_view.buffer()[0], kFeedbackInputTensorValues[i] + 1); + } + { + const auto read_view = + feedback_incremented_int_copy[i].Get().GetCpuReadView(); + EXPECT_EQ(read_view.buffer()[0], kFeedbackInputTensorValues[i]); + } + } +} + +} // namespace +} // namespace api2 +} // namespace mediapipe diff --git a/mediapipe/calculators/tensor/inference_interpreter_delegate_runner.cc b/mediapipe/calculators/tensor/inference_interpreter_delegate_runner.cc index 76a8bf7818..10d270381f 100644 --- a/mediapipe/calculators/tensor/inference_interpreter_delegate_runner.cc +++ b/mediapipe/calculators/tensor/inference_interpreter_delegate_runner.cc @@ -15,12 +15,14 @@ #include "mediapipe/calculators/tensor/inference_interpreter_delegate_runner.h" #include +#include #include #include #include #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "mediapipe/calculators/tensor/inference_feedback_manager.h" #include "mediapipe/calculators/tensor/inference_io_mapper.h" #include "mediapipe/calculators/tensor/tensor_span.h" #include "mediapipe/calculators/tensor/tflite_delegate_ptr.h" @@ -86,11 +88,13 @@ class InferenceInterpreterDelegateRunner : public InferenceRunner { InferenceInterpreterDelegateRunner( api2::Packet model, std::unique_ptr interpreter, TfLiteDelegatePtr delegate, - InputOutputTensorNames&& input_output_tensor_names) + InputOutputTensorNames&& input_output_tensor_names, + std::unique_ptr feedback_manager) : model_(std::move(model)), interpreter_(std::move(interpreter)), delegate_(std::move(delegate)), - input_output_tensor_names_(std::move(input_output_tensor_names)) {} + input_output_tensor_names_(std::move(input_output_tensor_names)), + feedback_manager_(std::move(feedback_manager)) {} absl::StatusOr> Run( CalculatorContext* cc, const TensorSpan& tensor_span) override; @@ -104,27 +108,42 @@ class InferenceInterpreterDelegateRunner : public InferenceRunner { std::unique_ptr interpreter_; TfLiteDelegatePtr delegate_; InputOutputTensorNames input_output_tensor_names_; + std::unique_ptr feedback_manager_; }; absl::StatusOr> InferenceInterpreterDelegateRunner::Run( CalculatorContext* cc, const TensorSpan& tensor_span) { - // Read CPU input into tensors. - RET_CHECK_EQ(interpreter_->inputs().size(), tensor_span.size()); + const int num_feedback_tensors = + feedback_manager_ ? feedback_manager_->GetNumberOfFeedbackTensors() : 0; + + RET_CHECK_EQ(tensor_span.size() + num_feedback_tensors, + interpreter_->inputs().size()); // If the input tensors have dynamic shape, then the tensors need to be // resized and reallocated before we can copy the tensor values. bool resized_tensor_shapes = false; for (int i = 0; i < tensor_span.size(); ++i) { + int input_model_index; + if (feedback_manager_) { + // Feedback tensors are stripped from the InferenceRunner input. Calling + // MapInputTensorToModelIndex assigns the input tensors to the correct + // model index. + MP_ASSIGN_OR_RETURN(input_model_index, + feedback_manager_->MapInputTensorToModelIndex(i)); + } else { + input_model_index = i; + } const Tensor& input_tensor = tensor_span[i]; if (input_tensor.shape().is_dynamic) { const TfLiteTensor* interpreter_tensor = - interpreter_->tensor(interpreter_->inputs()[i]); + interpreter_->tensor(interpreter_->inputs()[input_model_index]); // TODO: Can avoid copying even these <= 4 values in the future. std::vector interpreter_dims{ interpreter_tensor->dims->data, interpreter_tensor->dims->data + interpreter_tensor->dims->size}; if (interpreter_dims != input_tensor.shape().dims) { - interpreter_->ResizeInputTensorStrict(i, input_tensor.shape().dims); + interpreter_->ResizeInputTensorStrict(input_model_index, + input_tensor.shape().dims); resized_tensor_shapes = true; } } @@ -135,39 +154,49 @@ absl::StatusOr> InferenceInterpreterDelegateRunner::Run( // TODO: Replace this using the util function in // inference_calculator_utils. for (int i = 0; i < tensor_span.size(); ++i) { + int input_model_index; + if (feedback_manager_) { + // Feedback tensors are stripped from the InferenceRunner input. Calling + // MapInputTensorToModelIndex assigns the input tensors to the correct + // model index. + MP_ASSIGN_OR_RETURN(input_model_index, + feedback_manager_->MapInputTensorToModelIndex(i)); + } else { + input_model_index = i; + } const TfLiteType input_tensor_type = - interpreter_->tensor(interpreter_->inputs()[i])->type; + interpreter_->tensor(interpreter_->inputs()[input_model_index])->type; const Tensor& input_tensor = tensor_span[i]; switch (input_tensor_type) { case TfLiteType::kTfLiteFloat16: case TfLiteType::kTfLiteFloat32: { CopyTensorBufferToInterpreter(input_tensor, interpreter_.get(), - i); + input_model_index); break; } case TfLiteType::kTfLiteUInt8: { CopyTensorBufferToInterpreter(input_tensor, interpreter_.get(), - i); + input_model_index); break; } case TfLiteType::kTfLiteInt8: { CopyTensorBufferToInterpreter(input_tensor, interpreter_.get(), - i); + input_model_index); break; } case TfLiteType::kTfLiteInt32: { CopyTensorBufferToInterpreter(input_tensor, interpreter_.get(), - i); + input_model_index); break; } case TfLiteType::kTfLiteString: { CopyTensorBufferToInterpreter(input_tensor, interpreter_.get(), - i); + input_model_index); break; } case TfLiteType::kTfLiteBool: { CopyTensorBufferToInterpreter(input_tensor, interpreter_.get(), - i); + input_model_index); break; } default: @@ -184,8 +213,13 @@ absl::StatusOr> InferenceInterpreterDelegateRunner::Run( // Output result tensors (CPU). const auto& tensor_indexes = interpreter_->outputs(); std::vector output_tensors; - output_tensors.reserve(tensor_indexes.size()); + output_tensors.reserve(tensor_indexes.size() - num_feedback_tensors); for (int i = 0; i < tensor_indexes.size(); ++i) { + if (feedback_manager_ && + feedback_manager_->IsFeedbackOutputTensorAtIndex(i)) { + // Exclude feedback tensors from InferenceRunner output. + continue; + } TfLiteTensor* tensor = interpreter_->tensor(tensor_indexes[i]); Tensor::Shape shape{std::vector{ tensor->dims->data, tensor->dims->data + tensor->dims->size}}; @@ -232,6 +266,9 @@ absl::StatusOr> InferenceInterpreterDelegateRunner::Run( TfLiteTypeGetName(tensor->type))); } } + if (feedback_manager_) { + feedback_manager_->SwapFeedbackTensors(); + } return output_tensors; } @@ -239,7 +276,9 @@ absl::StatusOr> CreateInferenceInterpreterDelegateRunner( api2::Packet model, api2::Packet op_resolver, TfLiteDelegatePtr delegate, - int interpreter_num_threads) { + int interpreter_num_threads, + const mediapipe::InferenceCalculatorOptions::InputOutputConfig* + input_output_config) { InterpreterBuilder interpreter_builder(*model.Get(), op_resolver.Get()); if (delegate) { interpreter_builder.AddDelegate(delegate.get()); @@ -257,9 +296,17 @@ CreateInferenceInterpreterDelegateRunner( auto input_output_tensor_names, InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( *interpreter)); + std::unique_ptr inference_feedback_manager; + if (input_output_config) { + // Create inference_feedback_manager if input_output_config is available. + inference_feedback_manager = std::make_unique(); + MP_RETURN_IF_ERROR(inference_feedback_manager->Init( + *input_output_config, input_output_tensor_names, interpreter.get())); + } return std::make_unique( std::move(model), std::move(interpreter), std::move(delegate), - std::move(input_output_tensor_names)); + std::move(input_output_tensor_names), + std::move(inference_feedback_manager)); } } // namespace mediapipe diff --git a/mediapipe/calculators/tensor/inference_interpreter_delegate_runner.h b/mediapipe/calculators/tensor/inference_interpreter_delegate_runner.h index ca6d79851a..448e8c729d 100644 --- a/mediapipe/calculators/tensor/inference_interpreter_delegate_runner.h +++ b/mediapipe/calculators/tensor/inference_interpreter_delegate_runner.h @@ -33,11 +33,14 @@ namespace mediapipe { // // `delegate` can be nullptr, in that case newly initialized interpreter will // use what is available by default. +// `input_output_config` optional config to enable feedback tensors. absl::StatusOr> CreateInferenceInterpreterDelegateRunner( api2::Packet model, api2::Packet op_resolver, TfLiteDelegatePtr delegate, - int interpreter_num_threads); + int interpreter_num_threads, + const mediapipe::InferenceCalculatorOptions::InputOutputConfig* + input_output_config = nullptr); } // namespace mediapipe diff --git a/mediapipe/calculators/tensor/inference_io_mapper.cc b/mediapipe/calculators/tensor/inference_io_mapper.cc index 3a952bb6c5..81ce1018ca 100644 --- a/mediapipe/calculators/tensor/inference_io_mapper.cc +++ b/mediapipe/calculators/tensor/inference_io_mapper.cc @@ -106,6 +106,46 @@ static absl::StatusOr> MapTensorNamesToIndices( return result; }; +// Feedback tensors are excluded from the InferenceRunner input and output +// accordingly (since they are class-internally handled by the +// InferenceFeedbackManager). This means that the input and output Tensor orders +// of the InferenceRunner don't match the model I/O tensors anymore and +// therefore tensor I/O indices need to be adjusted accordingly. +absl::Status ExcludeFeedbackTensorsFromRemappingIndicesVector( + const InferenceCalculatorOptions::InputOutputConfig& io_config, + const std::vector& model_tensor_names, + std::vector& remapping_tensor_indices) { + // Create set of all feedback tensor names. + absl::flat_hash_set feedback_tensor_names; + for (const auto& link : io_config.feedback_tensor_links()) { + { + // No need to check for name collisions. Inference feedback manager + // confirms validity of feedback tensor names. + feedback_tensor_names.insert(link.from_output_tensor_name()); + feedback_tensor_names.insert(link.to_input_tensor_name()); + } + } + // Built model index translation vector which maps InferenceRunner I/O tensor + // indices to InferenceRunner I/O indices with excluded feedback tensors. + std::vector indices_translation(model_tensor_names.size(), -1); + int model_output_idx = 0; + for (int i = 0; i < model_tensor_names.size(); ++i) { + if (!feedback_tensor_names.contains(model_tensor_names[i])) { + indices_translation[i] = model_output_idx; + ++model_output_idx; + } + } + // Adjust remapping_tensor_indices. + for (int i = 0; i < remapping_tensor_indices.size(); ++i) { + const int model_index = remapping_tensor_indices[i]; + RET_CHECK(model_index >= 0 && model_index < indices_translation.size()) + << "Index " << model_index << " out of range."; + remapping_tensor_indices[i] = + indices_translation[remapping_tensor_indices[i]]; + } + return absl::OkStatus(); +} + } // namespace // static @@ -145,6 +185,18 @@ InferenceIoMapper::GetInputOutputTensorNamesFromModel( absl::Status InferenceIoMapper::UpdateIoMap( const InferenceCalculatorOptions::InputOutputConfig& io_config, const InputOutputTensorNames& input_output_tensor_names) { + num_feedback_tensors_ = io_config.feedback_tensor_links().size(); + + if ((io_config.has_input_tensor_indices_map() || + io_config.has_output_tensor_indices_map()) && + num_feedback_tensors_ > 0) { + // TODO b/336767692 - remove this check once indices-based feedback + // tensors are supported. + return absl::FailedPreconditionError( + "Feedback tensors are not supported with tensor index-based I/O " + "mapping."); + } + input_tensor_indices_.clear(); output_tensor_indices_.clear(); @@ -186,6 +238,9 @@ absl::Status InferenceIoMapper::UpdateIoMap( input_output_tensor_names.begin()->second; if (io_config.has_input_tensor_names_map()) { + // Read number of model inputs directly from the signature. + const int num_model_input_tensors = + input_output_tensor_names_default_signature.input_tensor_names.size(); input_tensor_indices_.reserve( io_config.input_tensor_names_map().tensor_names().size()); MP_ASSIGN_OR_RETURN( @@ -193,18 +248,38 @@ absl::Status InferenceIoMapper::UpdateIoMap( MapTensorNamesToIndices( input_output_tensor_names_default_signature.input_tensor_names, io_config.input_tensor_names_map())); + if (num_feedback_tensors_ > 0) { + MP_RETURN_IF_ERROR(ExcludeFeedbackTensorsFromRemappingIndicesVector( + io_config, + input_output_tensor_names_default_signature.input_tensor_names, + input_tensor_indices_)); + } + // Feedback tensors are excluded from the input_tensor_indices_. + RET_CHECK_EQ(input_tensor_indices_.size() + num_feedback_tensors_, + num_model_input_tensors) + << "Unexpected number of input tensors."; } if (io_config.has_output_tensor_names_map()) { - output_tensor_indices_.reserve( - io_config.output_tensor_names_map().tensor_names().size()); + const int num_model_output_tensors = + input_output_tensor_names_default_signature.output_tensor_names.size(); + output_tensor_indices_.reserve(num_model_output_tensors); MP_ASSIGN_OR_RETURN( output_tensor_indices_, MapTensorNamesToIndices( input_output_tensor_names_default_signature.output_tensor_names, io_config.output_tensor_names_map())); + if (num_feedback_tensors_ > 0) { + MP_RETURN_IF_ERROR(ExcludeFeedbackTensorsFromRemappingIndicesVector( + io_config, + input_output_tensor_names_default_signature.output_tensor_names, + output_tensor_indices_)); + } + // Feedback tensors are excluded from the output_tensor_indices_. + RET_CHECK_EQ(output_tensor_indices_.size() + num_feedback_tensors_, + num_model_output_tensors) + << "Unexpected number of output tensors."; } - return absl::OkStatus(); } @@ -214,8 +289,7 @@ absl::StatusOr InferenceIoMapper::RemapInputTensors( return unmapped_tensors; } RET_CHECK_EQ(unmapped_tensors.size(), input_tensor_indices_.size()) - << "Number of input tensors does not match number indices in the " - "provided mapping."; + << "Unexpected number of input tensors."; std::vector mapped_tensors(unmapped_tensors.size()); for (int i = 0; i < unmapped_tensors.size(); ++i) { const int index = input_tensor_indices_[i]; @@ -233,8 +307,7 @@ absl::StatusOr> InferenceIoMapper::RemapOutputTensors( return std::move(unmapped_tensors); } RET_CHECK_EQ(unmapped_tensors.size(), output_tensor_indices_.size()) - << "Number of output tensors does not match number indices in the " - "provided mapping."; + << "Unexpected number of output tensors."; std::vector mapped_tensors; mapped_tensors.reserve(unmapped_tensors.size()); for (int i = 0; i < unmapped_tensors.size(); ++i) { @@ -242,6 +315,7 @@ absl::StatusOr> InferenceIoMapper::RemapOutputTensors( RET_CHECK(index < unmapped_tensors.size()) << "Index " << index << " out of range" << ". Size of TensorIndicesMap: " << unmapped_tensors.size() << "."; + mapped_tensors.emplace_back(std::move(unmapped_tensors[index])); } return mapped_tensors; diff --git a/mediapipe/calculators/tensor/inference_io_mapper.h b/mediapipe/calculators/tensor/inference_io_mapper.h index 42fc79125d..0b1e017e15 100644 --- a/mediapipe/calculators/tensor/inference_io_mapper.h +++ b/mediapipe/calculators/tensor/inference_io_mapper.h @@ -68,6 +68,7 @@ class InferenceIoMapper { std::vector&& unmapped_tensors); private: + int num_feedback_tensors_ = 0; std::vector input_tensor_indices_; std::vector output_tensor_indices_; }; diff --git a/mediapipe/calculators/tensor/inference_io_mapper_test.cc b/mediapipe/calculators/tensor/inference_io_mapper_test.cc index 04ba50da8e..76b8c6a098 100644 --- a/mediapipe/calculators/tensor/inference_io_mapper_test.cc +++ b/mediapipe/calculators/tensor/inference_io_mapper_test.cc @@ -453,11 +453,9 @@ TEST_F(InferenceIoMapperTest, ShouldReportTooFewInputMappingIndices) { InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( *interpreter_)); MP_EXPECT_OK(mapper.UpdateIoMap(map, input_output_tensor_names)); - EXPECT_THAT( - mapper.RemapInputTensors(MakeTensorSpan(input_tensors_unmapped)), - StatusIs( - absl::StatusCode::kInternal, - HasSubstr("Number of input tensors does not match number indices"))); + EXPECT_THAT(mapper.RemapInputTensors(MakeTensorSpan(input_tensors_unmapped)), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("Unexpected number of input tensors"))); } TEST_F(InferenceIoMapperTest, ShouldReportTooFewOutputMappingIndices) { @@ -480,11 +478,9 @@ TEST_F(InferenceIoMapperTest, ShouldReportTooFewOutputMappingIndices) { InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( *interpreter_)); MP_EXPECT_OK(mapper.UpdateIoMap(map, input_output_tensor_names)); - EXPECT_THAT( - mapper.RemapOutputTensors(std::move(output_tensors_unmapped)), - StatusIs( - absl::StatusCode::kInternal, - HasSubstr("Number of output tensors does not match number indices"))); + EXPECT_THAT(mapper.RemapOutputTensors(std::move(output_tensors_unmapped)), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("Unexpected number of output tensors"))); } TEST_F(InferenceIoMapperTest, ShouldReportTooManyMappingInputIndices) { @@ -509,11 +505,9 @@ TEST_F(InferenceIoMapperTest, ShouldReportTooManyMappingInputIndices) { InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( *interpreter_)); MP_EXPECT_OK(mapper.UpdateIoMap(map, input_output_tensor_names)); - EXPECT_THAT( - mapper.RemapInputTensors(MakeTensorSpan(input_tensors_unmapped)), - StatusIs( - absl::StatusCode::kInternal, - HasSubstr("Number of input tensors does not match number indices"))); + EXPECT_THAT(mapper.RemapInputTensors(MakeTensorSpan(input_tensors_unmapped)), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("Unexpected number of input tensors"))); } TEST_F(InferenceIoMapperTest, ShouldReportTooManyMappingOutputIndices) { @@ -538,11 +532,9 @@ TEST_F(InferenceIoMapperTest, ShouldReportTooManyMappingOutputIndices) { InferenceIoMapper::GetInputOutputTensorNamesFromInterpreter( *interpreter_)); MP_EXPECT_OK(mapper.UpdateIoMap(map, input_output_tensor_names)); - EXPECT_THAT( - mapper.RemapOutputTensors(std::move(output_tensors_unmapped)), - StatusIs( - absl::StatusCode::kInternal, - HasSubstr("Number of output tensors does not match number indices"))); + EXPECT_THAT(mapper.RemapOutputTensors(std::move(output_tensors_unmapped)), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("Unexpected number of output tensors"))); } TEST_F(InferenceIoMapperTest, ShouldReportDuplicatedMappingIndices) { @@ -681,8 +673,10 @@ class InferenceIoMapperSmokeTest absl::StrCat("input", n), MakePacket(std::move(input_tensor)).At(Timestamp(0)))); } - MP_EXPECT_OK(graph.WaitUntilIdle()); + MP_EXPECT_OK(graph.CloseAllInputStreams()); + MP_EXPECT_OK(graph.WaitUntilDone()); + EXPECT_EQ(output_packets.size(), expected_order.size()); for (int i = 0; i < output_packets.size(); ++i) { EXPECT_EQ(output_packets[i].size(), 1); const auto read_view = diff --git a/mediapipe/calculators/tensor/inference_runner.h b/mediapipe/calculators/tensor/inference_runner.h index 41834091b1..65856cc484 100644 --- a/mediapipe/calculators/tensor/inference_runner.h +++ b/mediapipe/calculators/tensor/inference_runner.h @@ -18,8 +18,8 @@ class InferenceRunner { virtual absl::StatusOr> Run( CalculatorContext* cc, const TensorSpan& tensor_span) = 0; - // Returns the input/output tensor names from the TfLite model signature. This - // enables tensor name-based I/O mapping. + // Returns the TfLite model's input/output tensor names. This enables tensor + // name based I/O mapping in the InferenceCalculator base class. virtual const InputOutputTensorNames& GetInputOutputTensorNames() const = 0; }; diff --git a/mediapipe/calculators/tensor/testdata/feedback_tensor_test_model.tflite b/mediapipe/calculators/tensor/testdata/feedback_tensor_test_model.tflite new file mode 100644 index 0000000000000000000000000000000000000000..3c0d360ca1220308c86d73abb44cd241c3f3fa73 GIT binary patch literal 1696 zcmZ`(J!n%=6h4Wytxc$Zr8)$RAmUJBq7|hGC8RAWq^ZQvMd9_ed1)UeFXi>6D0JwM z!NH+}ql1Hkh&VflI5>!ih;v6L5vL9vs?YDcxi8I+x#7#XKj-H==bm?CB2t{oT_2OU z#N@n;$tj6qZ&(H}7wm~V4~ZNQVeB6hxj!fZA=rX8$B43`%?MVsr`;FeB(Q0@Teh=i zl&rE@uXxZB*#@$}6@crfe|?`mIj4%ofu#>0%mp8&&r^jPnz1x>egT+`8l0Rm<42+)caYNCfAHz&`K`_yOzz zpMclE3*Z^R`((@|#HU;bn!UkZ8g0d69uok2hAc#eRL@`%q8&$6_aQw-v2IK+=jZP2 zmU$M`jQ4jfk-C{cEkEm+jF(E>;8ts74}Cue{MI^*!J3K;Do#NkG-K~fKZNc!umd!J z#{m0rAva%$XKD`Wf^JGEAIycf8+vf32m7d?SlSceD1Md>6fFKb5n=9JmZZ_VDS!y z%fLR?AAwyU2hip#`WhlcE(C2b!g2h)Mf}r{cTCr_J#>&&%9xePREi1!`4D;SgZv3{ ziUS9I?2|Zaw=bY>*lKZLZWv+mxvk#=vs5xtS0AZ^_g(Q=PwK8?pL6JHO_85ny3CpJ z`9{SdyWTUS^Shs-YKtp5i;Szc^_;pnT`gohuFj5fVEr+|ll(`1mms=hwcY@}XX*v$ zjK|u((zRh-338y%BK9vs_VwBJuFz|Hy_p01yxD!dIRvp7AExpFvEq>3u{vs%?Bz~- JdLK|fk$;W*NjLxi literal 0 HcmV?d00001 diff --git a/mediapipe/calculators/tensor/testdata/feedback_tensor_with_state_copy_model.tflite b/mediapipe/calculators/tensor/testdata/feedback_tensor_with_state_copy_model.tflite new file mode 100644 index 0000000000000000000000000000000000000000..e6489d5174685b205495f69eb8d59898c88e0a36 GIT binary patch literal 1172 zcmaJ>%}x_h6h0lR5NJ@5APWplBynSEt0)P(reH!+puxr!$vAc{I6A#;rgw-C7A9_7 znD_!b03kkt3zx183l}U1&%lC(8y&xIm^Q$+G)LmQ7cewrHyja z4#D3jqV2u^*|s{}9q7^0pq9_6T5dz0uEo=FDe^Ok&014+803!)7;g?>{u9Kzk(}>< z*Va~FF0H>LVm6%R*dRV<%6sR`m{kywGXv%5M=B(Hz=)jX$1leRN4@R?>ew6Y z*wf=uWxiAvIZbjhU#0Yvq5=K-6GgPy_o?9o#xt0ta%+z7FH~p%Feb5^iILP`F^size; ++i) { + if (tensor.dims->data[i] == -1) { + return true; + } + } + return false; +} +} // namespace mediapipe::util::tflite diff --git a/mediapipe/util/tflite/utils.h b/mediapipe/util/tflite/utils.h new file mode 100644 index 0000000000..7c7f0bc5ef --- /dev/null +++ b/mediapipe/util/tflite/utils.h @@ -0,0 +1,25 @@ +// Copyright 2024 The MediaPipe Authors. +// +// 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 MEDIAPIPE_UTIL_TFLITE_UTILS_H_ +#define MEDIAPIPE_UTIL_TFLITE_UTILS_H_ + +#include "tensorflow/lite/c/common.h" + +namespace mediapipe::util::tflite { + +// Returns +bool IsDynamicTensor(const TfLiteTensor& tensor); +} // namespace mediapipe::util::tflite + +#endif // MEDIAPIPE_UTIL_TFLITE_UTILS_H_