Skip to content

Validate TFLite inputs before copying tensor data #4148

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
x-Spartacus wants to merge 1 commit into
base: master
Choose a base branch
from
+187 −3
Open

Validate TFLite inputs before copying tensor data #4148

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 50 additions & 1 deletion tensorflow_serving/servables/tensorflow/tflite_session.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ limitations under the License.
#include "tensorflow/cc/saved_model/signature_constants.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/tensor_util.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/notification.h"
#include "tensorflow/core/lib/gtl/cleanup.h"
Expand Down Expand Up @@ -201,6 +202,17 @@ absl::Status SetInputAndInvokeMiniBatch(
int tflite_input_idx = tflite_input_indices[i];
auto tflite_input_tensor = interpreter->tensor(tflite_input_idx);
const auto& tf_input_tensors = inputs[i];
DataType expected_tf_type;
TF_RETURN_IF_ERROR(
TfLiteTypeToTfType(tflite_input_tensor->type, &expected_tf_type));
for (const Tensor* tf_input_tensor : tf_input_tensors) {
if (tf_input_tensor->dtype() != expected_tf_type) {
return absl::InvalidArgumentError(absl::StrCat(
"Expected input '", tflite_input_tensor->name, "' to have type ",
DataTypeString(expected_tf_type), ", but got ",
DataTypeString(tf_input_tensor->dtype()), "."));
}
}
if (tflite_input_tensor->type != kTfLiteString) {
const Tensor* tf_input_tensor = tf_input_tensors[0];
// concated.tensor_data() may be accessed later.
Expand Down Expand Up @@ -231,6 +243,13 @@ absl::Status SetInputAndInvokeMiniBatch(
if (interpreter->AllocateTensors() != kTfLiteOk) {
return absl::InternalError("Failed to allocate tensors");
}
tflite_input_tensor = interpreter->tensor(tflite_input_idx);
}
if (tensor_bytes.size() != tflite_input_tensor->bytes) {
return absl::InternalError(absl::StrCat(
"Input tensor byte size mismatch for '", tflite_input_tensor->name,
"': source has ", tensor_bytes.size(), " bytes, destination has ",
tflite_input_tensor->bytes, " bytes."));
}
std::memcpy(tflite_input_tensor->data.raw, tensor_bytes.data(),
tensor_bytes.size());
Expand All @@ -246,6 +265,7 @@ absl::Status SetInputAndInvokeMiniBatch(
if (interpreter->AllocateTensors() != kTfLiteOk) {
return absl::InternalError("Failed to allocate tensors");
}
tflite_input_tensor = interpreter->tensor(tflite_input_idx);
}
if (fixed_batch_size) {
*fixed_batch_size = interpreter_wrapper->GetBatchSize();
Expand Down Expand Up @@ -663,7 +683,36 @@ absl::Status MergeInputTensors(
return absl::InternalError(absl::StrCat(
"Batch size expected to be positive; was ", batch.num_tasks()));
}
const int tensors_per_task = batch.task(0).inputs.size();
const TfLiteBatchTask& reference_task = batch.task(0);
const int tensors_per_task = reference_task.inputs.size();
if (reference_task.input_indices.size() != tensors_per_task) {
return absl::InvalidArgumentError(
"TFLite batch task input tensors and indices are not aligned.");
}
if (reference_task.output_tensor_names == nullptr) {
return absl::InvalidArgumentError(
"TFLite batch task is missing output tensor names.");
}
for (int i = 1; i < batch.num_tasks(); ++i) {
const TfLiteBatchTask& task = batch.task(i);
if (task.inputs.size() != tensors_per_task) {
return absl::InvalidArgumentError(
"TFLite batch tasks have different input tensor counts.");
}
if (task.input_indices.size() != task.inputs.size()) {
return absl::InvalidArgumentError(
"TFLite batch task input tensors and indices are not aligned.");
}
if (task.input_indices != reference_task.input_indices) {
return absl::InvalidArgumentError(
"TFLite batch tasks have different input tensor indices.");
}
if (task.output_tensor_names == nullptr ||
*task.output_tensor_names != *reference_task.output_tensor_names) {
return absl::InvalidArgumentError(
"TFLite batch tasks have different output tensor names.");
}
}
*batch_size = 0;
// each entry in merged_inputs is a list of task tensors.
for (int i = 0; i < tensors_per_task; ++i) {
Expand Down
139 changes: 137 additions & 2 deletions tensorflow_serving/servables/tensorflow/tflite_session_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ namespace serving {
namespace {

using ::testing::_;
using ::testing::HasSubstr;
using ::testing::Pair;
using ::testing::SizeIs;
using ::testing::UnorderedElementsAre;
Expand Down Expand Up @@ -141,6 +142,39 @@ TEST(TfLiteSession, BasicTest) {
}
}

TEST(TfLiteSession, RejectsMismatchedInputTypes) {
std::string model_bytes;
TF_ASSERT_OK(ReadFileToString(tensorflow::Env::Default(),
test_util::TestSrcDirPath(kTestModel),
&model_bytes));

::google::protobuf::Map<std::string, SignatureDef> signatures;
std::unique_ptr<TfLiteSession> session;
tensorflow::SessionOptions options;
TF_ASSERT_OK(TfLiteSession::Create(
std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
absl::GetFlag(FLAGS_num_tflite_interpreters), &session, &signatures));

std::vector<Tensor> outputs;
Tensor int64_input(DT_INT64, TensorShape({1}));
absl::Status status = session->Run({{"x", int64_input}}, {"y"}, {}, &outputs);
EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
EXPECT_THAT(status.message(), HasSubstr("Expected input 'x'"));

Tensor string_input =
test::AsTensor<tstring>({"not a float"}, TensorShape({1}));
status = session->Run({{"x", string_input}}, {"y"}, {}, &outputs);
EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
EXPECT_THAT(status.message(), HasSubstr("Expected input 'x'"));

Tensor valid_input = test::AsTensor<float>({1.0}, TensorShape({1}));
outputs.clear();
TF_EXPECT_OK(session->Run({{"x", valid_input}}, {"y"}, {}, &outputs));
ASSERT_EQ(outputs.size(), 1);
test::ExpectTensorEqual<float>(
outputs[0], test::AsTensor<float>({2.5}, TensorShape({1})));
}

TEST(TfLiteSession, ResizeWithSameNumElementsTest) {
std::string model_bytes;
TF_ASSERT_OK(ReadFileToString(tensorflow::Env::Default(),
Expand Down Expand Up @@ -629,10 +663,18 @@ TEST(TfLiteSession, SimpleSignatureDefAndRun) {
ASSERT_EQ(sigdef.outputs().at(kSignatureOutput).name(), kTestModelOutput);
ASSERT_EQ(sigdef.method_name(), kClassifyMethodName);

Tensor input_list =
test::AsTensor<tstring>({"a", "b", "c", "d"}, TensorShape({4}));
Tensor invalid_input_list =
test::AsTensor<int32_t>({1, 2, 3, 4}, TensorShape({4}));
Tensor input_shape = test::AsTensor<int32_t>({2, 2}, TensorShape({2}));
std::vector<Tensor> outputs;
absl::Status status = session->Run({{kTestModelInputList, invalid_input_list},
{kTestModelInputShape, input_shape}},
{kTestModelOutput}, {}, &outputs);
EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
EXPECT_THAT(status.message(), HasSubstr("to have type string"));

Tensor input_list =
test::AsTensor<tstring>({"a", "b", "c", "d"}, TensorShape({4}));
TF_EXPECT_OK(session->Run(
{{kTestModelInputList, input_list}, {kTestModelInputShape, input_shape}},
{kTestModelOutput}, {}, &outputs));
Expand All @@ -642,6 +684,88 @@ TEST(TfLiteSession, SimpleSignatureDefAndRun) {
test::AsTensor<tstring>({"a", "b", "c", "d"}, TensorShape({2, 2})));
}

absl::Status BuildTwoInputBatchedSession(
std::unique_ptr<TfLiteSession>* session) {
auto model_signature_def_map = GetTestSignatureDefMap();
std::string model_bytes =
BuildTestModel(tflite::TensorType_STRING, /*use_flex_op=*/false,
&model_signature_def_map);
::google::protobuf::Map<std::string, SignatureDef> signatures;
tensorflow::SessionOptions options;
TF_RETURN_IF_ERROR(TfLiteSession::Create(std::move(model_bytes), options, 1,
1, session, &signatures));

BasicBatchScheduler<TfLiteBatchTask>::Options scheduler_options;
scheduler_options.num_batch_threads = 1;
scheduler_options.max_batch_size = 2;
scheduler_options.batch_timeout_micros = 10 * 1000 * 1000;
return (*session)->SetScheduler(
TfLiteSession::CreateDefaultBasicBatchScheduler, scheduler_options);
}

TEST(TfLiteSession, BatchedRequestsRejectDifferentInputCounts) {
std::unique_ptr<TfLiteSession> session;
TF_ASSERT_OK(BuildTwoInputBatchedSession(&session));

Tensor input_list = test::AsTensor<tstring>({"a"}, TensorShape({1}));
Tensor input_shape = test::AsTensor<int32_t>({1}, TensorShape({1}));
std::vector<Tensor> first_outputs;
std::vector<Tensor> second_outputs;
absl::Status first_status;
absl::Status second_status;
std::unique_ptr<Thread> first_request_thread(
Env::Default()->StartThread(ThreadOptions(), "first_request", [&] {
first_status = session->Run({{kTestModelInputList, input_list},
{kTestModelInputShape, input_shape}},
{kTestModelOutput}, {}, &first_outputs);
}));
std::unique_ptr<Thread> second_request_thread(
Env::Default()->StartThread(ThreadOptions(), "second_request", [&] {
second_status = session->Run({{kTestModelInputList, input_list}},
{kTestModelOutput}, {}, &second_outputs);
}));
first_request_thread.reset();
second_request_thread.reset();

EXPECT_EQ(first_status.code(), absl::StatusCode::kInvalidArgument);
EXPECT_EQ(second_status.code(), absl::StatusCode::kInvalidArgument);
EXPECT_THAT(first_status.message(),
HasSubstr("different input tensor counts"));
EXPECT_THAT(second_status.message(),
HasSubstr("different input tensor counts"));
}

TEST(TfLiteSession, BatchedRequestsRejectDifferentInputIndices) {
std::unique_ptr<TfLiteSession> session;
TF_ASSERT_OK(BuildTwoInputBatchedSession(&session));

Tensor input_list = test::AsTensor<tstring>({"a"}, TensorShape({1}));
Tensor input_shape = test::AsTensor<int32_t>({1}, TensorShape({1}));
std::vector<Tensor> first_outputs;
std::vector<Tensor> second_outputs;
absl::Status first_status;
absl::Status second_status;
std::unique_ptr<Thread> first_request_thread(
Env::Default()->StartThread(ThreadOptions(), "first_request", [&] {
first_status = session->Run({{kTestModelInputList, input_list}},
{kTestModelOutput}, {}, &first_outputs);
}));
std::unique_ptr<Thread> second_request_thread(
Env::Default()->StartThread(ThreadOptions(), "second_request", [&] {
second_status = session->Run({{kTestModelInputShape, input_shape}},
{kTestModelOutput}, {}, &second_outputs);
}));
first_request_thread.reset();
second_request_thread.reset();

EXPECT_EQ(first_status.code(), absl::StatusCode::kInvalidArgument);
EXPECT_EQ(second_status.code(), absl::StatusCode::kInvalidArgument);
EXPECT_THAT(first_status.message(),
HasSubstr("different input tensor indices"));
EXPECT_THAT(second_status.message(),
HasSubstr("different input tensor indices"));
}

absl::Status BuildSessionInBatch(std::unique_ptr<TfLiteSession>* sess,
bool use_model_batch_size,
const std::string& model_path) {
Expand Down Expand Up @@ -729,6 +853,17 @@ TEST_P(TfLiteSessionBatchSizeTest, TestBatchParallelismForFloat) {
EXPECT_TRUE(outputs[0].shape().IsSameSize(TensorShape({kBatchSize, 1})));
}

TEST_P(TfLiteSessionBatchSizeTest, RejectsMismatchedInputType) {
std::unique_ptr<TfLiteSession> sess;
TF_ASSERT_OK(BuildSessionInBatch(&sess, GetParam(), kTestModel));

Tensor invalid_input(DT_INT64, TensorShape({kBatchSize, 1}));
std::vector<Tensor> outputs;
absl::Status status = sess->Run({{"x", invalid_input}}, {"y"}, {}, &outputs);
EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
EXPECT_THAT(status.message(), HasSubstr("Expected input 'x'"));
}

TEST_P(TfLiteSessionBatchSizeTest, TestBatchParallelismForString) {
std::unique_ptr<TfLiteSession> sess;
TF_ASSERT_OK(BuildSessionInBatch(&sess, GetParam(), kParseExampleModel));
Expand Down