219 lines
8.6 KiB
C++
219 lines
8.6 KiB
C++
// Copyright 2022 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_interpreter_delegate_runner.h"
|
|
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
#include "absl/status/status.h"
|
|
#include "absl/status/statusor.h"
|
|
#include "mediapipe/framework/formats/tensor.h"
|
|
#include "mediapipe/framework/mediapipe_profiling.h"
|
|
#include "mediapipe/framework/port/ret_check.h"
|
|
#include "tensorflow/lite/c/c_api_types.h"
|
|
#include "tensorflow/lite/interpreter.h"
|
|
#include "tensorflow/lite/interpreter_builder.h"
|
|
#include "tensorflow/lite/string_util.h"
|
|
|
|
#define PERFETTO_TRACK_EVENT_NAMESPACE mediapipe
|
|
|
|
namespace mediapipe {
|
|
|
|
namespace {
|
|
|
|
template <typename T>
|
|
void CopyTensorBufferToInterpreter(const Tensor& input_tensor,
|
|
tflite::Interpreter* interpreter,
|
|
int input_tensor_index) {
|
|
auto input_tensor_view = input_tensor.GetCpuReadView();
|
|
auto input_tensor_buffer = input_tensor_view.buffer<T>();
|
|
T* local_tensor_buffer =
|
|
interpreter->typed_input_tensor<T>(input_tensor_index);
|
|
std::memcpy(local_tensor_buffer, input_tensor_buffer, input_tensor.bytes());
|
|
}
|
|
|
|
template <>
|
|
void CopyTensorBufferToInterpreter<char>(const Tensor& input_tensor,
|
|
tflite::Interpreter* interpreter,
|
|
int input_tensor_index) {
|
|
const char* input_tensor_buffer =
|
|
input_tensor.GetCpuReadView().buffer<char>();
|
|
tflite::DynamicBuffer dynamic_buffer;
|
|
dynamic_buffer.AddString(input_tensor_buffer,
|
|
input_tensor.shape().num_elements());
|
|
dynamic_buffer.WriteToTensorAsVector(
|
|
interpreter->tensor(interpreter->inputs()[input_tensor_index]));
|
|
}
|
|
|
|
template <typename T>
|
|
void CopyTensorBufferFromInterpreter(tflite::Interpreter* interpreter,
|
|
int output_tensor_index,
|
|
Tensor* output_tensor) {
|
|
auto output_tensor_view = output_tensor->GetCpuWriteView();
|
|
auto output_tensor_buffer = output_tensor_view.buffer<T>();
|
|
T* local_tensor_buffer =
|
|
interpreter->typed_output_tensor<T>(output_tensor_index);
|
|
std::memcpy(output_tensor_buffer, local_tensor_buffer,
|
|
output_tensor->bytes());
|
|
}
|
|
|
|
} // namespace
|
|
|
|
class InferenceInterpreterDelegateRunner : public InferenceRunner {
|
|
public:
|
|
InferenceInterpreterDelegateRunner(
|
|
api2::Packet<TfLiteModelPtr> model,
|
|
std::unique_ptr<tflite::Interpreter> interpreter,
|
|
TfLiteDelegatePtr delegate)
|
|
: model_(std::move(model)),
|
|
interpreter_(std::move(interpreter)),
|
|
delegate_(std::move(delegate)) {}
|
|
|
|
absl::StatusOr<std::vector<Tensor>> Run(
|
|
CalculatorContext* cc, const std::vector<Tensor>& input_tensors) override;
|
|
|
|
private:
|
|
api2::Packet<TfLiteModelPtr> model_;
|
|
std::unique_ptr<tflite::Interpreter> interpreter_;
|
|
TfLiteDelegatePtr delegate_;
|
|
};
|
|
|
|
absl::StatusOr<std::vector<Tensor>> InferenceInterpreterDelegateRunner::Run(
|
|
CalculatorContext* cc, const std::vector<Tensor>& input_tensors) {
|
|
// Read CPU input into tensors.
|
|
RET_CHECK_EQ(interpreter_->inputs().size(), input_tensors.size());
|
|
for (int i = 0; i < input_tensors.size(); ++i) {
|
|
const TfLiteType input_tensor_type =
|
|
interpreter_->tensor(interpreter_->inputs()[i])->type;
|
|
switch (input_tensor_type) {
|
|
case TfLiteType::kTfLiteFloat16:
|
|
case TfLiteType::kTfLiteFloat32: {
|
|
CopyTensorBufferToInterpreter<float>(input_tensors[i],
|
|
interpreter_.get(), i);
|
|
break;
|
|
}
|
|
case TfLiteType::kTfLiteUInt8: {
|
|
CopyTensorBufferToInterpreter<uint8_t>(input_tensors[i],
|
|
interpreter_.get(), i);
|
|
break;
|
|
}
|
|
case TfLiteType::kTfLiteInt8: {
|
|
CopyTensorBufferToInterpreter<int8_t>(input_tensors[i],
|
|
interpreter_.get(), i);
|
|
break;
|
|
}
|
|
case TfLiteType::kTfLiteInt32: {
|
|
CopyTensorBufferToInterpreter<int32_t>(input_tensors[i],
|
|
interpreter_.get(), i);
|
|
break;
|
|
}
|
|
case TfLiteType::kTfLiteString: {
|
|
CopyTensorBufferToInterpreter<char>(input_tensors[i],
|
|
interpreter_.get(), i);
|
|
break;
|
|
}
|
|
case TfLiteType::kTfLiteBool:
|
|
// No current use-case for copying MediaPipe Tensors with bool type to
|
|
// TfLiteTensors.
|
|
default:
|
|
return absl::InvalidArgumentError(
|
|
absl::StrCat("Unsupported input tensor type:", input_tensor_type));
|
|
}
|
|
}
|
|
|
|
// Run inference.
|
|
{
|
|
MEDIAPIPE_PROFILING(CPU_TASK_INVOKE, cc);
|
|
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
|
|
}
|
|
// Output result tensors (CPU).
|
|
const auto& tensor_indexes = interpreter_->outputs();
|
|
std::vector<Tensor> output_tensors;
|
|
output_tensors.reserve(tensor_indexes.size());
|
|
for (int i = 0; i < tensor_indexes.size(); ++i) {
|
|
TfLiteTensor* tensor = interpreter_->tensor(tensor_indexes[i]);
|
|
Tensor::Shape shape{std::vector<int>{
|
|
tensor->dims->data, tensor->dims->data + tensor->dims->size}};
|
|
switch (tensor->type) {
|
|
case TfLiteType::kTfLiteFloat16:
|
|
case TfLiteType::kTfLiteFloat32:
|
|
output_tensors.emplace_back(Tensor::ElementType::kFloat32, shape);
|
|
CopyTensorBufferFromInterpreter<float>(interpreter_.get(), i,
|
|
&output_tensors.back());
|
|
break;
|
|
case TfLiteType::kTfLiteUInt8:
|
|
output_tensors.emplace_back(
|
|
Tensor::ElementType::kUInt8, shape,
|
|
Tensor::QuantizationParameters{tensor->params.scale,
|
|
tensor->params.zero_point});
|
|
CopyTensorBufferFromInterpreter<uint8>(interpreter_.get(), i,
|
|
&output_tensors.back());
|
|
break;
|
|
case TfLiteType::kTfLiteInt8:
|
|
output_tensors.emplace_back(
|
|
Tensor::ElementType::kInt8, shape,
|
|
Tensor::QuantizationParameters{tensor->params.scale,
|
|
tensor->params.zero_point});
|
|
CopyTensorBufferFromInterpreter<int8>(interpreter_.get(), i,
|
|
&output_tensors.back());
|
|
break;
|
|
case TfLiteType::kTfLiteInt32:
|
|
output_tensors.emplace_back(Tensor::ElementType::kInt32, shape);
|
|
CopyTensorBufferFromInterpreter<int32_t>(interpreter_.get(), i,
|
|
&output_tensors.back());
|
|
break;
|
|
case TfLiteType::kTfLiteBool:
|
|
output_tensors.emplace_back(Tensor::ElementType::kBool, shape,
|
|
Tensor::QuantizationParameters{1.0f, 0});
|
|
CopyTensorBufferFromInterpreter<bool>(interpreter_.get(), i,
|
|
&output_tensors.back());
|
|
break;
|
|
case TfLiteType::kTfLiteString:
|
|
// No current use-case for copying TfLiteTensors with string type to
|
|
// MediaPipe Tensors.
|
|
default:
|
|
return absl::InvalidArgumentError(
|
|
absl::StrCat("Unsupported output tensor type:",
|
|
TfLiteTypeGetName(tensor->type)));
|
|
}
|
|
}
|
|
return output_tensors;
|
|
}
|
|
|
|
absl::StatusOr<std::unique_ptr<InferenceRunner>>
|
|
CreateInferenceInterpreterDelegateRunner(
|
|
api2::Packet<TfLiteModelPtr> model,
|
|
api2::Packet<tflite::OpResolver> op_resolver, TfLiteDelegatePtr delegate,
|
|
int interpreter_num_threads) {
|
|
tflite::InterpreterBuilder interpreter_builder(*model.Get(),
|
|
op_resolver.Get());
|
|
if (delegate) {
|
|
interpreter_builder.AddDelegate(delegate.get());
|
|
}
|
|
#if defined(__EMSCRIPTEN__)
|
|
interpreter_builder.SetNumThreads(1);
|
|
#else
|
|
interpreter_builder.SetNumThreads(interpreter_num_threads);
|
|
#endif // __EMSCRIPTEN__
|
|
std::unique_ptr<tflite::Interpreter> interpreter;
|
|
RET_CHECK_EQ(interpreter_builder(&interpreter), kTfLiteOk);
|
|
RET_CHECK(interpreter);
|
|
RET_CHECK_EQ(interpreter->AllocateTensors(), kTfLiteOk);
|
|
return std::make_unique<InferenceInterpreterDelegateRunner>(
|
|
std::move(model), std::move(interpreter), std::move(delegate));
|
|
}
|
|
|
|
} // namespace mediapipe
|