TensorFlow Serving C++ API Documentation
tflite_session_test.cc
1 /* Copyright 2019 Google Inc. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7  http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #include "tensorflow_serving/servables/tensorflow/tflite_session.h"
17 
18 #include <map>
19 #include <memory>
20 #include <random>
21 #include <set>
22 #include <string>
23 #include <utility>
24 #include <vector>
25 
26 #include <gtest/gtest.h>
27 #include "absl/flags/flag.h"
28 #include "absl/functional/bind_front.h"
29 #include "flatbuffers/flexbuffers.h"
30 #include "tensorflow/cc/saved_model/signature_constants.h"
31 #include "tensorflow/core/example/example.pb.h"
32 #include "tensorflow/core/example/feature.pb.h"
33 #include "tensorflow/core/framework/attr_value.pb.h"
34 #include "tensorflow/core/framework/node_def.pb.h"
35 #include "tensorflow/core/framework/tensor_testutil.h"
36 #include "tensorflow/core/framework/types.pb.h"
37 #include "tensorflow/core/lib/core/status_test_util.h"
38 #include "tensorflow/core/platform/env.h"
39 #include "tensorflow/core/platform/logging.h"
40 #include "tensorflow/core/platform/protobuf.h"
41 #include "tensorflow/core/platform/strcat.h"
42 #include "tensorflow/core/platform/test_benchmark.h"
43 #include "tensorflow/core/platform/threadpool_options.h"
44 #include "tensorflow/core/protobuf/config.pb.h"
45 #include "tensorflow/core/protobuf/error_codes.pb.h"
46 #include "tensorflow/lite/string_util.h"
47 #include "tensorflow/lite/tools/signature/signature_def_util.h"
48 #include "tensorflow/lite/util.h"
49 #include "tensorflow/lite/version.h"
50 #include "tensorflow_serving/test_util/test_util.h"
51 
52 ABSL_FLAG(int, num_pools, 1, "Number of interpreter pools of a TfLiteSession.");
53 
54 ABSL_FLAG(int, num_tflite_interpreters, 1,
55  "Number of TFLite interpreters "
56  "in an interpreter pool of a TfLiteSession.");
57 
58 namespace tensorflow {
59 namespace serving {
60 namespace {
61 
62 using ::testing::_;
63 using ::testing::Pair;
64 using ::testing::SizeIs;
65 using ::testing::UnorderedElementsAre;
66 
67 constexpr char kTestModel[] =
68  "/servables/tensorflow/testdata/saved_model_half_plus_two_tflite/00000123/"
69  "model.tflite";
70 
71 constexpr char kTestModelWithSigdef[] =
72  "/servables/tensorflow/testdata/"
73  "saved_model_half_plus_two_tflite_with_sigdef/00000123/model.tflite";
74 
75 constexpr char kMobileNetModel[] =
76  "/servables/tensorflow/testdata/mobilenet_v1_quant_tflite/00000123/"
77  "model.tflite";
78 
79 constexpr char kParseExampleModel[] =
80  "/servables/tensorflow/testdata/parse_example_tflite/00000123/"
81  "model.tflite";
82 
83 TEST(TfLiteSession, BasicTest) {
84  string model_bytes;
85  TF_ASSERT_OK(ReadFileToString(tensorflow::Env::Default(),
86  test_util::TestSrcDirPath(kTestModel),
87  &model_bytes));
88 
89  ::google::protobuf::Map<string, SignatureDef> signatures;
90  std::unique_ptr<TfLiteSession> session;
91  tensorflow::SessionOptions options;
92  TF_ASSERT_OK(TfLiteSession::Create(
93  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
94  absl::GetFlag(FLAGS_num_tflite_interpreters), &session, &signatures));
95  EXPECT_EQ(signatures.size(), 1);
96  EXPECT_EQ(signatures.begin()->first, "serving_default");
97  EXPECT_THAT(signatures.begin()->second, test_util::EqualsProto(R"(
98  inputs {
99  key: "x"
100  value {
101  name: "x"
102  dtype: DT_FLOAT
103  tensor_shape {
104  dim { size: 1 }
105  dim { size: 1 }
106  }
107  }
108  }
109  outputs {
110  key: "y"
111  value {
112  name: "y"
113  dtype: DT_FLOAT
114  tensor_shape {
115  dim { size: 1 }
116  dim { size: 1 }
117  }
118  }
119  }
120  method_name: "tensorflow/serving/predict"
121  )"));
122  Tensor input = test::AsTensor<float>({1.0, 2.0, 3.0}, TensorShape({3}));
123  {
124  // Use TF Lite tensor names.
125  std::vector<Tensor> outputs;
126  TF_EXPECT_OK(session->Run({{"x", input}}, {"y"}, {}, &outputs));
127  ASSERT_EQ(outputs.size(), 1);
128  test::ExpectTensorEqual<float>(
129  outputs[0], test::AsTensor<float>({2.5, 3, 3.5}, TensorShape({3})));
130  }
131  {
132  // Use TF tensor names (with `:0` suffix).
133  std::vector<Tensor> outputs;
134  TF_EXPECT_OK(session->Run({{"x:0", input}}, {"y:0"}, {}, &outputs));
135  ASSERT_EQ(outputs.size(), 1);
136  test::ExpectTensorEqual<float>(
137  outputs[0], test::AsTensor<float>({2.5, 3, 3.5}, TensorShape({3})));
138  }
139 }
140 
141 TEST(TfLiteSession, ResizeWithSameNumElementsTest) {
142  string model_bytes;
143  TF_ASSERT_OK(ReadFileToString(tensorflow::Env::Default(),
144  test_util::TestSrcDirPath(kTestModel),
145  &model_bytes));
146 
147  ::google::protobuf::Map<string, SignatureDef> signatures;
148  std::unique_ptr<TfLiteSession> session;
149  tensorflow::SessionOptions options;
150  TF_ASSERT_OK(TfLiteSession::Create(
151  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
152  absl::GetFlag(FLAGS_num_tflite_interpreters), &session, &signatures));
153  EXPECT_EQ(signatures.size(), 1);
154  EXPECT_EQ(signatures.begin()->first, "serving_default");
155  EXPECT_THAT(signatures.begin()->second, test_util::EqualsProto(R"pb(
156  inputs {
157  key: "x"
158  value {
159  name: "x"
160  dtype: DT_FLOAT
161  tensor_shape {
162  dim { size: 1 }
163  dim { size: 1 }
164  }
165  }
166  }
167  outputs {
168  key: "y"
169  value {
170  name: "y"
171  dtype: DT_FLOAT
172  tensor_shape {
173  dim { size: 1 }
174  dim { size: 1 }
175  }
176  }
177  }
178  method_name: "tensorflow/serving/predict"
179  )pb"));
180  Tensor input = test::AsTensor<float>({2.0}, TensorShape({1}));
181  {
182  // Use TF Lite tensor names.
183  std::vector<Tensor> outputs;
184  TF_EXPECT_OK(session->Run({{"x", input}}, {"y"}, {}, &outputs));
185  ASSERT_EQ(outputs.size(), 1);
186  test::ExpectTensorEqual<float>(
187  outputs[0], test::AsTensor<float>({3.0}, TensorShape({1})));
188  }
189 }
190 
191 TEST(TfLiteSession, ModelFromLegacyConverterWithSigdef) {
192  // A model converted with TF v1 converter, having a signature def.
193  // The signature def references an input tensor named "tflite_input:0", but
194  // the converter striped the tensor name to "tflite_input".
195  string model_bytes;
196  TF_ASSERT_OK(ReadFileToString(tensorflow::Env::Default(),
197  test_util::TestSrcDirPath(kTestModelWithSigdef),
198  &model_bytes));
199 
200  ::google::protobuf::Map<string, SignatureDef> signatures;
201  std::unique_ptr<TfLiteSession> session;
202  tensorflow::SessionOptions options;
203  TF_ASSERT_OK(TfLiteSession::Create(
204  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
205  absl::GetFlag(FLAGS_num_tflite_interpreters), &session, &signatures));
206  EXPECT_EQ(signatures.size(), 1);
207  EXPECT_EQ(signatures.begin()->first, "serving_default");
208  // While, in the model, the tensor name of input "x" is "tflite_input:0". in
209  // the output signature the name must have falled back to "tflite_input".
210  EXPECT_THAT(signatures.begin()->second, test_util::EqualsProto(R"(
211  inputs {
212  key: "x"
213  value {
214  name: "tflite_input"
215  dtype: DT_FLOAT
216  tensor_shape {
217  dim { size: 1 }
218  dim { size: 1 }
219  }
220  }
221  }
222  outputs {
223  key: "y"
224  value {
225  name: "y"
226  dtype: DT_FLOAT
227  tensor_shape {
228  dim { size: 1 }
229  dim { size: 1 }
230  }
231  }
232  }
233  method_name: "tensorflow/serving/predict"
234  )"));
235 
236  Tensor input = test::AsTensor<float>({1.0, 2.0, 3.0}, TensorShape({3}));
237  {
238  // Use TF Lite tensor names.
239  std::vector<Tensor> outputs;
240  TF_EXPECT_OK(session->Run({{"tflite_input", input}}, {"y"}, {}, &outputs));
241  ASSERT_EQ(outputs.size(), 1);
242  test::ExpectTensorEqual<float>(
243  outputs[0], test::AsTensor<float>({2.5, 3, 3.5}, TensorShape({3})));
244  }
245  {
246  // Use TF tensor names (with `:0` suffix).
247  std::vector<Tensor> outputs;
248  TF_EXPECT_OK(
249  session->Run({{"tflite_input:0", input}}, {"y:0"}, {}, &outputs));
250  ASSERT_EQ(outputs.size(), 1);
251  test::ExpectTensorEqual<float>(
252  outputs[0], test::AsTensor<float>({2.5, 3, 3.5}, TensorShape({3})));
253  }
254 }
255 
256 constexpr char kTestModelInputList[] = "list";
257 constexpr char kTestModelInputShape[] = "shape";
258 constexpr char kTestModelOutput[] = "output";
259 
260 constexpr char kSignatureInputList[] = "input_list";
261 constexpr char kSignatureInputShape[] = "input_shape";
262 constexpr char kSignatureOutput[] = "sigdef_output";
263 
264 constexpr int kBatchSize = 500;
265 
266 std::map<string, SignatureDef> GetTestSignatureDefMap() {
267  auto signature_def = SignatureDef();
268  TensorInfo input_list_tensor;
269  TensorInfo input_shape_tensor;
270  TensorInfo output_tensor;
271  *input_list_tensor.mutable_name() = absl::StrCat(kTestModelInputList, ":0");
272  *input_shape_tensor.mutable_name() = absl::StrCat(kTestModelInputShape, ":0");
273  *output_tensor.mutable_name() = absl::StrCat(kTestModelOutput, ":0");
274  *signature_def.mutable_method_name() = kClassifyMethodName;
275  (*signature_def.mutable_inputs())[kSignatureInputList] = input_list_tensor;
276  (*signature_def.mutable_inputs())[kSignatureInputShape] = input_shape_tensor;
277  (*signature_def.mutable_outputs())[kSignatureOutput] = output_tensor;
278  std::map<string, SignatureDef> signature_def_map = {
279  {kDefaultServingSignatureDefKey, signature_def}};
280  return signature_def_map;
281 }
282 
283 tensorflow::DataType ToTfTensorType(tflite::TensorType tflite_type) {
284  switch (tflite_type) {
285  case tflite::TensorType_INT32:
286  return tensorflow::DT_INT32;
287  case tflite::TensorType_STRING:
288  return tensorflow::DT_STRING;
289  default:
290  LOG(FATAL) << "Unsupported tflite type: " << tflite_type;
291  }
292 }
293 
294 string BuildTestModel(tflite::TensorType tensor_type,
295  const string& input1_tensor_name,
296  const string& input2_tensor_name, bool use_flex_op,
297  std::map<string, SignatureDef>* signature_def_map) {
298  std::vector<int32_t> inputs;
299  std::vector<int32_t> outputs;
300  std::vector<flatbuffers::Offset<tflite::Tensor>> tensors;
301  std::vector<flatbuffers::Offset<tflite::OperatorCode>> opcodes;
302  std::vector<flatbuffers::Offset<tflite::Operator>> operators;
303  std::vector<flatbuffers::Offset<tflite::Buffer>> buffers;
304  flatbuffers::FlatBufferBuilder builder;
305 
306  // Input list: 1D tensor for list of `tensor_type` elements.
307  inputs.push_back(tensors.size());
308  tensors.push_back(CreateTensor(builder, builder.CreateVector<int>({1}),
309  tensor_type, /*buffer=*/0,
310  builder.CreateString(input1_tensor_name),
311  /*quantization=*/0, /*is_variable=*/false));
312 
313  // Input shape: 1D tensor for shape.
314  inputs.push_back(tensors.size());
315  tensors.push_back(CreateTensor(builder, builder.CreateVector<int>({1}),
316  tflite::TensorType_INT32, /*buffer=*/0,
317  builder.CreateString(input2_tensor_name),
318  /*quantization=*/0, /*is_variable=*/false));
319 
320  // Output: Reshaped list to shape.
321  outputs.push_back(tensors.size());
322  tensors.push_back(CreateTensor(builder, builder.CreateVector<int>({1}),
323  tensor_type, /*buffer=*/0,
324  builder.CreateString(kTestModelOutput),
325  /*quantization=*/0, /*is_variable=*/false));
326 
327  // Add reshape operator.
328  tflite::BuiltinOptions builtin_opts_type =
329  tflite::BuiltinOptions_ReshapeOptions;
330  flatbuffers::Offset<void> reshape_opts =
331  tflite::CreateReshapeOptions(builder, builder.CreateVector<int>({}))
332  .Union();
333  flatbuffers::Offset<flatbuffers::Vector<uint8_t>> custom_opts = 0;
334  if (use_flex_op) {
335  string flexop = std::string(tflite::kFlexCustomCodePrefix) + "Reshape";
336  opcodes.push_back(CreateOperatorCodeDirect(
337  builder, tflite::BuiltinOperator_CUSTOM, flexop.data()));
338  builtin_opts_type = tflite::BuiltinOptions_NONE;
339  reshape_opts = 0;
340  NodeDef node_def;
341  node_def.set_name("Reshape");
342  node_def.set_op("Reshape");
343  (*node_def.mutable_attr())["T"].set_type(ToTfTensorType(tensor_type));
344  string node_def_str;
345  CHECK(node_def.SerializeToString(&node_def_str));
346  auto flex_builder = absl::make_unique<flexbuffers::Builder>();
347  flex_builder->Vector([&]() {
348  flex_builder->String(node_def.op());
349  flex_builder->String(node_def_str);
350  });
351  flex_builder->Finish();
352  custom_opts = builder.CreateVector(flex_builder->GetBuffer());
353  } else {
354  opcodes.push_back(
355  CreateOperatorCode(builder, tflite::BuiltinOperator_RESHAPE, 0));
356  }
357 
358  operators.push_back(CreateOperator(
359  builder, /*opcode_index=*/0, builder.CreateVector<int32_t>(inputs),
360  builder.CreateVector<int32_t>(outputs), builtin_opts_type, reshape_opts,
361  custom_opts, tflite::CustomOptionsFormat_FLEXBUFFERS));
362 
363  auto subgraph = CreateSubGraph(builder, builder.CreateVector(tensors),
364  builder.CreateVector<int32_t>(inputs),
365  builder.CreateVector<int32_t>(outputs),
366  builder.CreateVector(operators));
367  builder.Finish(CreateModel(
368  builder, TFLITE_SCHEMA_VERSION, builder.CreateVector(opcodes),
369  builder.CreateVector(&subgraph, 1), builder.CreateString("testmodel"),
370  builder.CreateVector(buffers)));
371 
372  if (signature_def_map) {
373  std::string model_buffer = string(
374  reinterpret_cast<char*>(builder.GetBufferPointer()), builder.GetSize());
375  std::string model_buffer_with_signature_def;
376  auto model = tflite::FlatBufferModel::BuildFromModel(
377  flatbuffers::GetRoot<tflite::Model>(model_buffer.data()));
378  TF_CHECK_OK(tflite::SetSignatureDefMap(model->GetModel(),
379  *signature_def_map,
380  &model_buffer_with_signature_def));
381  return model_buffer_with_signature_def;
382  }
383 
384  return string(reinterpret_cast<char*>(builder.GetBufferPointer()),
385  builder.GetSize());
386 }
387 
388 // Returns a serialized FlatBuffer tflite model.
389 //
390 // The model has two inputs (kTestModelInputList|Shape) and one output
391 // kTestModelOutput. The output is list that is reshaped to shape via
392 // tf.reshape operator.
393 //
394 // Elements of list are expected to be of `tensor_type` type. `use_flex_op`
395 // sets up the model to use the `Reshape` *flex* op as opposed to using the
396 // builtin `Reshape` op from TF Lite.
397 string BuildTestModel(tflite::TensorType tensor_type, bool use_flex_op,
398  std::map<string, SignatureDef>* signature_def_map) {
399  return BuildTestModel(tensor_type, kTestModelInputList, kTestModelInputShape,
400  use_flex_op, signature_def_map);
401 }
402 
403 TEST(TfLiteSession, ProcessStrings) {
404  auto model_signature_def_map = GetTestSignatureDefMap();
405  string model_bytes =
406  BuildTestModel(tflite::TensorType_STRING, /*use_flex_op=*/false,
407  &model_signature_def_map);
408  ::google::protobuf::Map<string, SignatureDef> signatures;
409  std::unique_ptr<TfLiteSession> session;
410  tensorflow::SessionOptions options;
411  TF_ASSERT_OK(TfLiteSession::Create(
412  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
413  absl::GetFlag(FLAGS_num_tflite_interpreters), &session, &signatures));
414  Tensor input_list =
415  test::AsTensor<tstring>({"a", "b", "c", "d"}, TensorShape({4}));
416  Tensor input_shape = test::AsTensor<int32>({2, 2}, TensorShape({2}));
417  std::vector<Tensor> outputs;
418  TF_EXPECT_OK(session->Run(
419  {{kTestModelInputList, input_list}, {kTestModelInputShape, input_shape}},
420  {kTestModelOutput}, {}, &outputs));
421  ASSERT_EQ(outputs.size(), 1);
422  test::ExpectTensorEqual<tstring>(
423  outputs[0],
424  test::AsTensor<tstring>({"a", "b", "c", "d"}, TensorShape({2, 2})));
425 }
426 
427 TEST(TfLiteSession, ProcessStringsFlex) {
428  auto model_signature_def_map = GetTestSignatureDefMap();
429  string model_bytes =
430  BuildTestModel(tflite::TensorType_STRING, /*use_flex_op=*/true,
431  &model_signature_def_map);
432  ::google::protobuf::Map<string, SignatureDef> signatures;
433  std::unique_ptr<TfLiteSession> session;
434  tensorflow::SessionOptions options;
435  TF_ASSERT_OK(TfLiteSession::Create(
436  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
437  absl::GetFlag(FLAGS_num_tflite_interpreters), &session, &signatures));
438  Tensor input_list =
439  test::AsTensor<tstring>({"a", "b", "c", "d"}, TensorShape({4}));
440  Tensor input_shape = test::AsTensor<int32>({2, 2}, TensorShape({2}));
441  std::vector<Tensor> outputs;
442  TF_EXPECT_OK(session->Run(
443  {{kTestModelInputList, input_list}, {kTestModelInputShape, input_shape}},
444  {kTestModelOutput}, {}, &outputs));
445  ASSERT_EQ(outputs.size(), 1);
446  test::ExpectTensorEqual<tstring>(
447  outputs[0],
448  test::AsTensor<tstring>({"a", "b", "c", "d"}, TensorShape({2, 2})));
449 }
450 
451 TEST(TfLiteSession, ThreadPoolOptions) {
452  auto model_signature_def_map = GetTestSignatureDefMap();
453  string model_bytes =
454  BuildTestModel(tflite::TensorType_STRING, /*use_flex_op=*/false,
455  &model_signature_def_map);
456  ::google::protobuf::Map<string, SignatureDef> signatures;
457  std::unique_ptr<TfLiteSession> session;
458  tensorflow::SessionOptions options;
459  TF_ASSERT_OK(TfLiteSession::Create(
460  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
461  absl::GetFlag(FLAGS_num_tflite_interpreters), &session, &signatures));
462  Tensor input_list =
463  test::AsTensor<tstring>({"a", "b", "c", "d"}, TensorShape({4}));
464  Tensor input_shape = test::AsTensor<int32>({2, 2}, TensorShape({2}));
465  std::vector<Tensor> outputs;
466  RunMetadata run_metadata;
467  thread::ThreadPoolOptions thread_pool_options;
468  test_util::CountingThreadPool inter_op_threadpool(Env::Default(), "InterOp",
469  /*num_threads=*/1);
470  test_util::CountingThreadPool intra_op_threadpool(Env::Default(), "IntraOp",
471  /*num_threads=*/1);
472  thread_pool_options.inter_op_threadpool = &inter_op_threadpool;
473  thread_pool_options.intra_op_threadpool = &intra_op_threadpool;
474  TF_EXPECT_OK(session->Run(
475  RunOptions(),
476  {{kTestModelInputList, input_list}, {kTestModelInputShape, input_shape}},
477  {kTestModelOutput}, {}, &outputs, &run_metadata, thread_pool_options));
478  ASSERT_EQ(outputs.size(), 1);
479  test::ExpectTensorEqual<tstring>(
480  outputs[0],
481  test::AsTensor<tstring>({"a", "b", "c", "d"}, TensorShape({2, 2})));
482  // TfLiteSession does not use the ThreadPoolOptions.
483  EXPECT_EQ(inter_op_threadpool.NumScheduled(), 0);
484  EXPECT_EQ(intra_op_threadpool.NumScheduled(), 0);
485 }
486 
487 TEST(TfLiteSession, SimpleSignatureDef) {
488  auto model_signature_def_map = GetTestSignatureDefMap();
489  string model_bytes =
490  BuildTestModel(tflite::TensorType_STRING, /*use_flex_op=*/false,
491  &model_signature_def_map);
492 
493  ::google::protobuf::Map<string, SignatureDef> signatures;
494  // Fill an entry in the output signatures map, to check that it gets cleared
495  string kResidualSignatureKey = "residual_signature";
496  signatures[kResidualSignatureKey] = SignatureDef();
497 
498  std::unique_ptr<TfLiteSession> session;
499  tensorflow::SessionOptions options;
500  TF_ASSERT_OK(TfLiteSession::Create(
501  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
502  absl::GetFlag(FLAGS_num_tflite_interpreters), &session, &signatures));
503 
504  ASSERT_THAT(signatures,
505  UnorderedElementsAre(Pair(kDefaultServingSignatureDefKey, _)));
506 
507  auto sigdef = signatures[kDefaultServingSignatureDefKey];
508  EXPECT_EQ(sigdef.inputs().at(kSignatureInputList).name(),
509  kTestModelInputList);
510  EXPECT_EQ(sigdef.inputs().at(kSignatureInputShape).name(),
511  kTestModelInputShape);
512  EXPECT_EQ(sigdef.outputs().at(kSignatureOutput).name(), kTestModelOutput);
513  EXPECT_EQ(sigdef.method_name(), kClassifyMethodName);
514 }
515 
516 TEST(TfLiteSession, MultipleSignatureDef) {
517  TensorInfo input_list_tensor;
518  TensorInfo input_shape_tensor;
519  TensorInfo output_tensor;
520  *input_list_tensor.mutable_name() = kTestModelInputList;
521  *input_shape_tensor.mutable_name() = kTestModelInputShape;
522  *output_tensor.mutable_name() = kTestModelOutput;
523  SignatureDef signature1 = SignatureDef();
524  *signature1.mutable_method_name() = kClassifyMethodName;
525  (*signature1.mutable_inputs())[kSignatureInputList] = input_list_tensor;
526  (*signature1.mutable_outputs())[kSignatureOutput] = output_tensor;
527  SignatureDef signature2 = SignatureDef();
528  *signature2.mutable_method_name() = kClassifyMethodName;
529  (*signature2.mutable_inputs())[kSignatureInputShape] = input_shape_tensor;
530  (*signature2.mutable_outputs())[kSignatureOutput] = output_tensor;
531  constexpr char kSignatureKey1[] = "signature1";
532  constexpr char kSignatureKey2[] = "signature2";
533  std::map<string, SignatureDef> signature_def_map = {
534  {kSignatureKey1, signature1}, {kSignatureKey2, signature2}};
535 
536  string model_bytes = BuildTestModel(
537  tflite::TensorType_STRING, /*use_flex_op=*/false, &signature_def_map);
538  ::google::protobuf::Map<string, SignatureDef> signatures;
539  std::unique_ptr<TfLiteSession> session;
540  tensorflow::SessionOptions options;
541  TF_EXPECT_OK(TfLiteSession::Create(
542  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
543  absl::GetFlag(FLAGS_num_tflite_interpreters), &session, &signatures));
544 
545  ASSERT_THAT(signatures, UnorderedElementsAre(Pair(kSignatureKey1, _),
546  Pair(kSignatureKey2, _)));
547  auto result_signature1 = signatures[kSignatureKey1];
548  EXPECT_THAT(result_signature1.inputs().at(kSignatureInputList).name(),
549  kTestModelInputList);
550  EXPECT_EQ(result_signature1.outputs().at(kSignatureOutput).name(),
551  kTestModelOutput);
552  EXPECT_EQ(result_signature1.method_name(), kClassifyMethodName);
553  auto result_signature2 = signatures[kSignatureKey2];
554  EXPECT_EQ(result_signature2.inputs().at(kSignatureInputShape).name(),
555  kTestModelInputShape);
556  EXPECT_EQ(result_signature2.outputs().at(kSignatureOutput).name(),
557  kTestModelOutput);
558  EXPECT_EQ(result_signature2.method_name(), kClassifyMethodName);
559 }
560 
561 TEST(TfLiteSession, SignatureDefWithCommonTensorPrefix) {
562  // Attempts to normalize behaviors of different TFLite converter versions
563  // (which at one point striped the :<index> suffix from tensor names),
564  // must not create name collisions when signatures have tensors with the same
565  // prefix.
566  SignatureDef signature;
567  protobuf::TextFormat::ParseFromString(R"(
568  inputs {
569  key: "x0"
570  value {
571  name: "myTensor:0"
572  dtype: DT_FLOAT
573  tensor_shape { dim { size: 1 } }
574  }
575  }
576  inputs {
577  key: "x1"
578  value {
579  name: "myTensor:1"
580  dtype: DT_FLOAT
581  tensor_shape { dim { size: 1 } }
582  }
583  }
584  method_name: "tensorflow/serving/predict"
585  )",
586  &signature);
587  std::map<string, SignatureDef> signature_def_map = {
588  {kDefaultServingSignatureDefKey, signature}};
589  string model_bytes =
590  BuildTestModel(tflite::TensorType_STRING, "myTensor:0", "myTensor:1",
591  /*use_flex_op=*/false, &signature_def_map);
592  ::google::protobuf::Map<string, SignatureDef> signatures;
593  std::unique_ptr<TfLiteSession> session;
594  tensorflow::SessionOptions options;
595  TF_ASSERT_OK(TfLiteSession::Create(
596  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
597  absl::GetFlag(FLAGS_num_tflite_interpreters), &session, &signatures));
598 
599  // Inputs must still be processed as two different tensors.
600  auto outputSigdef = signatures[kDefaultServingSignatureDefKey];
601  std::set<string> tensorNamesSet;
602  for (const auto& input : outputSigdef.inputs()) {
603  tensorNamesSet.insert(input.second.name());
604  }
605  EXPECT_THAT(tensorNamesSet, SizeIs(2));
606 }
607 
608 TEST(TfLiteSession, SimpleSignatureDefAndRun) {
609  auto model_signature_def_map = GetTestSignatureDefMap();
610  string model_bytes =
611  BuildTestModel(tflite::TensorType_STRING, /*use_flex_op=*/false,
612  &model_signature_def_map);
613  ::google::protobuf::Map<string, SignatureDef> signatures;
614  std::unique_ptr<TfLiteSession> session;
615  tensorflow::SessionOptions options;
616  TF_EXPECT_OK(TfLiteSession::Create(
617  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
618  absl::GetFlag(FLAGS_num_tflite_interpreters), &session, &signatures));
619 
620  auto sigdef = signatures[kDefaultServingSignatureDefKey];
621  ASSERT_EQ(sigdef.inputs().at(kSignatureInputList).name(),
622  kTestModelInputList);
623  ASSERT_EQ(sigdef.inputs().at(kSignatureInputShape).name(),
624  kTestModelInputShape);
625  ASSERT_EQ(sigdef.outputs().at(kSignatureOutput).name(), kTestModelOutput);
626  ASSERT_EQ(sigdef.method_name(), kClassifyMethodName);
627 
628  Tensor input_list =
629  test::AsTensor<tstring>({"a", "b", "c", "d"}, TensorShape({4}));
630  Tensor input_shape = test::AsTensor<int32>({2, 2}, TensorShape({2}));
631  std::vector<Tensor> outputs;
632  TF_EXPECT_OK(session->Run(
633  {{kTestModelInputList, input_list}, {kTestModelInputShape, input_shape}},
634  {kTestModelOutput}, {}, &outputs));
635  ASSERT_EQ(outputs.size(), 1);
636  test::ExpectTensorEqual<tstring>(
637  outputs[0],
638  test::AsTensor<tstring>({"a", "b", "c", "d"}, TensorShape({2, 2})));
639 }
640 
641 Status BuildSessionInBatch(std::unique_ptr<TfLiteSession>* sess,
642  bool use_model_batch_size,
643  const string& model_path) {
644  std::string model_bytes;
645  TF_RETURN_IF_ERROR(ReadFileToString(
646  Env::Default(), test_util::TestSrcDirPath(model_path), &model_bytes));
647  auto model = tflite::FlatBufferModel::BuildFromModel(
648  flatbuffers::GetRoot<tflite::Model>(model_bytes.data()));
649  const int model_batch_size = 5;
650  if (use_model_batch_size) {
651  const tflite::Model* tflite_model = model->GetModel();
652  auto mutable_model = absl::make_unique<tflite::ModelT>();
653  tflite_model->UnPackTo(mutable_model.get(), nullptr);
654 
655  if (mutable_model->subgraphs.size() != 1) {
656  return Status(
657  static_cast<absl::StatusCode>(absl::StatusCode::kInvalidArgument),
658  strings::StrCat("Model subgraph size ",
659  mutable_model->subgraphs.size(), " not equal to 1"));
660  }
661  auto* subgraph = mutable_model->subgraphs[0].get();
662  if (subgraph->inputs.size() != 1) {
663  return Status(
664  static_cast<absl::StatusCode>(absl::StatusCode::kInvalidArgument),
665  strings::StrCat("Model subgraph input size ",
666  mutable_model->subgraphs.size(), " not equal to 1"));
667  }
668  auto* tensor = subgraph->tensors[subgraph->inputs[0]].get();
669  if (tensor->shape[0] != 1) {
670  return Status(
671  static_cast<absl::StatusCode>(absl::StatusCode::kInvalidArgument),
672  strings::StrCat("Model subgraph input shape[0] ",
673  mutable_model->subgraphs.size(), " not equal to 1"));
674  }
675  tensor->shape[0] = model_batch_size;
676  flatbuffers::FlatBufferBuilder builder;
677  auto packed_model = tflite::Model::Pack(builder, mutable_model.get());
678  FinishModelBuffer(builder, packed_model);
679  model_bytes.assign(
680  reinterpret_cast<const char*>(builder.GetBufferPointer()),
681  builder.GetSize());
682  }
683  auto model_signature_def_map = GetTestSignatureDefMap();
684  ::google::protobuf::Map<string, SignatureDef> signatures;
685  tensorflow::SessionOptions options;
686  const int num_tflite_interpreters = 4;
687 
688  TF_RETURN_IF_ERROR(TfLiteSession::Create(std::move(model_bytes), options, 1,
689  num_tflite_interpreters, sess,
690  &signatures));
691  auto scheduler_options = (*sess)->GetSchedulerOptions();
692  const int expected_batch_size = use_model_batch_size
693  ? model_batch_size
694  : kBatchSize / num_tflite_interpreters;
695  if (scheduler_options.max_execution_batch_size != expected_batch_size) {
696  return Status(
697  static_cast<absl::StatusCode>(absl::StatusCode::kInvalidArgument),
698  strings::StrCat("Scheulder max_execution_batch_size ",
699  scheduler_options.max_execution_batch_size,
700  " not equal to expected ", expected_batch_size));
701  }
702  return Status();
703 }
704 
705 using TfLiteSessionBatchSizeTest = ::testing::TestWithParam<bool>;
706 TEST_P(TfLiteSessionBatchSizeTest, TestBatchParallelismForFloat) {
707  std::unique_ptr<TfLiteSession> sess;
708  TF_ASSERT_OK(BuildSessionInBatch(&sess, GetParam(), kTestModel));
709  std::vector<float> example_list;
710  std::vector<float> expected;
711  std::vector<tstring> expected_bytes;
712  std::vector<Tensor> outputs;
713 
714  std::mt19937 random_engine;
715  auto random_func = [&]() {
716  return std::uniform_real_distribution<float>(-0.5, 0.5)(random_engine);
717  };
718  for (int i = 0; i < kBatchSize; i++) {
719  example_list.push_back(random_func());
720  }
721 
722  Tensor example_list_tensor =
723  test::AsTensor<float>(example_list, TensorShape({kBatchSize, 1}));
724  TF_EXPECT_OK(sess->Run({{"x", example_list_tensor}}, {"y"}, {}, &outputs));
725  EXPECT_TRUE(outputs[0].shape().IsSameSize(TensorShape({kBatchSize, 1})));
726 }
727 
728 TEST_P(TfLiteSessionBatchSizeTest, TestBatchParallelismForString) {
729  std::unique_ptr<TfLiteSession> sess;
730  TF_ASSERT_OK(BuildSessionInBatch(&sess, GetParam(), kParseExampleModel));
731  const float default_value = 0;
732  std::vector<tstring> example_list;
733  std::vector<float> expected;
734  std::vector<tstring> expected_bytes;
735  std::vector<Tensor> outputs;
736 
737  std::mt19937 random_engine;
738  auto random_func = [&]() {
739  return std::uniform_real_distribution<float>(-0.5, 0.5)(random_engine);
740  };
741  const std::string kTestString = "test string";
742  const std::string kDefaultString = "missing";
743  for (int i = 0; i < kBatchSize; i++) {
744  float val = random_func();
745  tensorflow::Example example;
746  std::string str;
747  if (val < -1) {
748  expected.push_back(default_value);
749  expected_bytes.push_back(kDefaultString);
750  } else {
751  expected.push_back(val);
752  expected_bytes.push_back(kTestString);
753  auto* features = example.mutable_features();
754  (*features->mutable_feature())["x"].mutable_float_list()->add_value(val);
755  (*features->mutable_feature())["y"].mutable_bytes_list()->add_value(
756  kTestString);
757  }
758  example.SerializeToString(&str);
759  example_list.push_back(str);
760  }
761  Tensor example_list_tensor =
762  test::AsTensor<tstring>(example_list, TensorShape({kBatchSize}));
763  TF_EXPECT_OK(sess->Run(
764  {{"input", example_list_tensor}},
765  {"ParseExample/ParseExampleV2", "ParseExample/ParseExampleV2:1"}, {},
766  &outputs));
767  test::ExpectTensorEqual<float>(
768  outputs[0],
769  test::AsTensor<float>(expected, TensorShape({kBatchSize, 1})));
770  EXPECT_EQ(outputs.size(), 2);
771  test::ExpectTensorEqual<tstring>(
772  outputs[1],
773  test::AsTensor<tstring>(expected_bytes, TensorShape({kBatchSize, 1})));
774 }
775 
776 INSTANTIATE_TEST_SUITE_P(TfLiteSessionBatchSizeTests,
777  TfLiteSessionBatchSizeTest, ::testing::Bool());
778 
779 TEST(TfLiteSession, TestSetScheduler) {
780  std::string model_bytes;
781  TF_ASSERT_OK(ReadFileToString(Env::Default(),
782  test_util::TestSrcDirPath(kParseExampleModel),
783  &model_bytes));
784  auto model = tflite::FlatBufferModel::BuildFromModel(
785  flatbuffers::GetRoot<tflite::Model>(model_bytes.data()));
786  auto model_signature_def_map = GetTestSignatureDefMap();
787  ::google::protobuf::Map<string, SignatureDef> signatures;
788  std::unique_ptr<TfLiteSession> sess;
789  tensorflow::SessionOptions options;
790 
791  int split_called = 0;
792  auto TestSplitTfLiteInputTask =
793  [&split_called](
794  std::unique_ptr<TfLiteBatchTask>* input_task_ptr,
795  int open_batch_remaining_slot, int max_batch_size,
796  std::vector<std::unique_ptr<TfLiteBatchTask>>* output_tasks) {
797  split_called += 1;
798  auto status = TfLiteSession::SplitTfLiteInputTask(
799  input_task_ptr, open_batch_remaining_slot, max_batch_size,
800  output_tasks);
801  return status;
802  };
803 
804  BasicBatchScheduler<TfLiteBatchTask>::Options scheduler_options;
805  scheduler_options.num_batch_threads = 1;
806  scheduler_options.max_batch_size = internal::kInitialBatchSize;
807  scheduler_options.enable_large_batch_splitting = true;
808  scheduler_options.max_execution_batch_size = 130;
809  scheduler_options.max_enqueued_batches = 4;
810  scheduler_options.split_input_task_func = TestSplitTfLiteInputTask;
811 
812  TF_ASSERT_OK(TfLiteSession::Create(std::move(model_bytes), options, 1, 1,
813  &sess, &signatures));
814 
815  TF_ASSERT_OK(sess->SetScheduler(
816  TfLiteSession::CreateDefaultBasicBatchScheduler, scheduler_options));
817 
818  const int batch_size = 500;
819  const float default_value = 0;
820  std::vector<tstring> example_list;
821  std::vector<float> expected;
822  std::vector<tstring> expected_bytes;
823  std::vector<Tensor> outputs;
824 
825  std::mt19937 random_engine;
826  auto random_func = [&]() {
827  return std::uniform_real_distribution<float>(-0.5, 0.5)(random_engine);
828  };
829  const std::string kTestString = "test string";
830  const std::string kDefaultString = "missing";
831  for (int i = 0; i < batch_size; i++) {
832  float val = random_func();
833  tensorflow::Example example;
834  std::string str;
835  if (val < -1) {
836  expected.push_back(default_value);
837  expected_bytes.push_back(kDefaultString);
838  } else {
839  expected.push_back(val);
840  expected_bytes.push_back(kTestString);
841  auto* features = example.mutable_features();
842  (*features->mutable_feature())["x"].mutable_float_list()->add_value(val);
843  (*features->mutable_feature())["y"].mutable_bytes_list()->add_value(
844  kTestString);
845  }
846  example.SerializeToString(&str);
847  example_list.push_back(str);
848  }
849 
850  Tensor example_list_tensor =
851  test::AsTensor<tstring>(example_list, TensorShape({batch_size}));
852  TF_EXPECT_OK(sess->Run(
853  {{"input", example_list_tensor}},
854  {"ParseExample/ParseExampleV2", "ParseExample/ParseExampleV2:1"}, {},
855  &outputs));
856  test::ExpectTensorEqual<float>(
857  outputs[0],
858  test::AsTensor<float>(expected, TensorShape({batch_size, 1})));
859  EXPECT_EQ(outputs.size(), 2);
860  test::ExpectTensorEqual<tstring>(
861  outputs[1],
862  test::AsTensor<tstring>(expected_bytes, TensorShape({batch_size, 1})));
863  EXPECT_EQ(split_called, 1);
864 }
865 
866 #ifdef PLATFORM_GOOGLE
867 // These benchmarks rely on https://github.com/google/benchmark features,
868 // not available in open-sourced TF codebase.
869 
870 static void BM_Reshape(benchmark::State& state, bool use_flex_op) {
871  static TfLiteSession* session;
872  if (state.thread_index() == 0) {
873  auto model_signature_def_map = GetTestSignatureDefMap();
874  string model_bytes = BuildTestModel(tflite::TensorType_INT32, use_flex_op,
875  &model_signature_def_map);
876  ::google::protobuf::Map<string, SignatureDef> signatures;
877  std::unique_ptr<TfLiteSession> sess;
878  tensorflow::SessionOptions options;
879  TF_ASSERT_OK(TfLiteSession::Create(
880  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
881  absl::GetFlag(FLAGS_num_tflite_interpreters), &sess, &signatures));
882  session = sess.release();
883  }
884  Tensor input = test::AsTensor<int32>({1, 2, 3, 4, 5, 6}, TensorShape({6}));
885  Tensor input_shape = test::AsTensor<int32>({3, 2}, TensorShape({2}));
886  std::vector<Tensor> outputs;
887  for (auto _ : state) {
888  outputs.clear();
889  TF_ASSERT_OK(session->Run(
890  {{kTestModelInputList, input}, {kTestModelInputShape, input_shape}},
891  {kTestModelOutput}, {}, &outputs));
892  }
893 }
894 
895 static void BM_Reshape_Builtin(benchmark::State& state) {
896  BM_Reshape(state, /*use_flex_op=*/false);
897 }
898 BENCHMARK(BM_Reshape_Builtin)->UseRealTime()->ThreadRange(1, 64);
899 
900 static void BM_Reshape_Flex(benchmark::State& state) {
901  BM_Reshape(state, /*use_flex_op=*/true);
902 }
903 BENCHMARK(BM_Reshape_Flex)->UseRealTime()->ThreadRange(1, 64);
904 
905 void BM_HalfPlusTwo(benchmark::State& state) {
906  static TfLiteSession* session;
907  if (state.thread_index() == 0) {
908  string model_bytes;
909  TF_ASSERT_OK(ReadFileToString(
910  Env::Default(), test_util::TestSrcDirPath(kTestModel), &model_bytes));
911  ::google::protobuf::Map<string, SignatureDef> signatures;
912  std::unique_ptr<TfLiteSession> sess;
913  tensorflow::SessionOptions options;
914  TF_ASSERT_OK(TfLiteSession::Create(
915  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
916  absl::GetFlag(FLAGS_num_tflite_interpreters), &sess, &signatures));
917  session = sess.release();
918  }
919  Tensor input = test::AsTensor<float>({1.0, 2.0, 3.0}, TensorShape({3}));
920  std::vector<Tensor> outputs;
921  for (auto _ : state) {
922  outputs.clear();
923  TF_ASSERT_OK(session->Run({{"x", input}}, {"y"}, {}, &outputs));
924  }
925 }
926 BENCHMARK(BM_HalfPlusTwo)->UseRealTime()->ThreadRange(1, 64);
927 
928 void BM_MobileNet(benchmark::State& state) {
929  static TfLiteSession* session;
930  if (state.thread_index() == 0) {
931  string model_bytes;
932  TF_ASSERT_OK(ReadFileToString(Env::Default(),
933  test_util::TestSrcDirPath(kMobileNetModel),
934  &model_bytes));
935  ::google::protobuf::Map<string, SignatureDef> signatures;
936  std::unique_ptr<TfLiteSession> sess;
937  tensorflow::SessionOptions options;
938  TF_ASSERT_OK(TfLiteSession::Create(
939  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
940  absl::GetFlag(FLAGS_num_tflite_interpreters), &sess, &signatures));
941  session = sess.release();
942  }
943  std::vector<uint8> x_data(1 * 224 * 224 * 3, 1);
944  Tensor x = test::AsTensor<uint8>(x_data, TensorShape({1, 224, 224, 3}));
945  std::vector<Tensor> outputs;
946  for (auto _ : state) {
947  outputs.clear();
948  TF_ASSERT_OK(session->Run(
949  {{"input", x}}, {"MobilenetV1/Predictions/Reshape_1"}, {}, &outputs));
950  }
951 }
952 BENCHMARK(BM_MobileNet)->UseRealTime()->ThreadRange(1, 64);
953 
954 void BM_ParseExample(benchmark::State& state) {
955  static TfLiteSession* session;
956  if (state.thread_index() == 0) {
957  string model_bytes;
958  TF_ASSERT_OK(ReadFileToString(Env::Default(),
959  test_util::TestSrcDirPath(kParseExampleModel),
960  &model_bytes));
961  ::google::protobuf::Map<string, SignatureDef> signatures;
962  std::unique_ptr<TfLiteSession> sess;
963  tensorflow::SessionOptions options;
964  TF_ASSERT_OK(TfLiteSession::Create(
965  std::move(model_bytes), options, absl::GetFlag(FLAGS_num_pools),
966  absl::GetFlag(FLAGS_num_tflite_interpreters), &sess, &signatures));
967  session = sess.release();
968  }
969  const int kBatchSize = 500;
970  std::vector<tstring> example_list;
971  std::mt19937 random_engine;
972  auto random_func = [&]() {
973  return std::uniform_real_distribution<float>(-0.5, 0.5)(random_engine);
974  };
975  for (int i = 0; i < kBatchSize; i++) {
976  float val = random_func();
977  tensorflow::Example example;
978  std::string str;
979  auto* features = example.mutable_features();
980  (*features->mutable_feature())["x"].mutable_float_list()->add_value(val);
981  (*features->mutable_feature())["y"].mutable_bytes_list()->add_value("Test");
982  example.SerializeToString(&str);
983  example_list.push_back(str);
984  }
985 
986  Tensor input_batch =
987  test::AsTensor<tstring>(example_list, TensorShape({kBatchSize}));
988  std::vector<Tensor> outputs;
989  for (auto _ : state) {
990  outputs.clear();
991  TF_ASSERT_OK(session->Run(
992  {{"input", input_batch}},
993  {"ParseExample/ParseExampleV2", "ParseExample/ParseExampleV2:1"}, {},
994  &outputs));
995  }
996 }
997 BENCHMARK(BM_ParseExample)->UseRealTime()->ThreadRange(1, 64);
998 
999 #endif // PLATFORM_GOOGLE
1000 
1001 } // namespace
1002 } // namespace serving
1003 } // namespace tensorflow
Generated by doxygen 1.9.1