tensorflow/lite/tools/serialization/writer_lib_test.cc
// Tests writing only a portion of the subgraph.
Tests writing only a portion of the subgraph.
Tests writing only a portion of the subgraph .
TEST_P(SingleSubgraphTest, CustomInputOutputTest) {
  Interpreter interpreter;
  interpreter.AddTensors(4);
  constexpr float kFoo[] = {1, 2, 3};
  interpreter.SetTensorParametersReadWrite(0, kTfLiteFloat32, "a", {3},
                                           TfLiteQuantization());
  interpreter.SetTensorParametersReadOnly(
      1, kTfLiteFloat32, "b", {3}, TfLiteQuantization(),
      reinterpret_cast<const char*>(kFoo), sizeof(kFoo));
  interpreter.SetTensorParametersReadWrite(2, kTfLiteFloat32, "c", {3},
                                           TfLiteQuantization());
  interpreter.SetTensorParametersReadWrite(3, kTfLiteFloat32, "d", {3},
                                           TfLiteQuantization());
  interpreter.SetInputs({0, 1});
  interpreter.SetOutputs({3});
  const char* initial_data = "";
  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver;
  TfLiteAddParams* builtin_data =
      reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
  builtin_data->activation = kTfLiteActNone;
  builtin_data->pot_scale_int16 = false;
  const TfLiteRegistration* reg = resolver.FindOp(BuiltinOperator_ADD, 1);
  interpreter.AddNodeWithParameters({0, 1}, {2}, initial_data, 0,
                                    reinterpret_cast<void*>(builtin_data), reg);
  const TfLiteRegistration* reg2 = resolver.FindOp(BuiltinOperator_RELU, 1);
  interpreter.AddNodeWithParameters({2}, {3}, nullptr, 0, nullptr, reg2);
  const std::string test_file = CreateFilePath("test_custom.tflite");
  SubgraphWriter writer(&interpreter.primary_subgraph());
  EXPECT_EQ(writer.SetCustomInputOutput({2}, {3},
                                        {1}),
            kTfLiteOk);
  writer.SetUnusedTensors({0, 1});
  writer.Write(test_file);
  std::unique_ptr<FlatBufferModel> model =
      FlatBufferModel::BuildFromFile(test_file.c_str());
  InterpreterBuilder builder(*model, resolver);
  std::unique_ptr<Interpreter> new_interpreter;
  builder(&new_interpreter);
  ASSERT_EQ(new_interpreter->AllocateTensors(), kTfLiteOk);
}
TEST_P ( SingleSubgraphTest , CustomInputOutputTest ) { Interpreter interpreter ; interpreter . AddTensors ( 4 ) ; constexpr float kFoo [ ] = { 1 , 2 , 3 } ; interpreter . SetTensorParametersReadWrite ( 0 , kTfLiteFloat32 , " " , { 3 } , TfLiteQuantization ( ) ) ; interpreter . SetTensorParametersReadOnly ( 1 , kTfLiteFloat32 , " " , { 3 } , TfLiteQuantization ( ) , reinterpret_cast < const char * > ( kFoo ) , sizeof ( kFoo ) ) ; interpreter . SetTensorParametersReadWrite ( 2 , kTfLiteFloat32 , " " , { 3 } , TfLiteQuantization ( ) ) ; interpreter . SetTensorParametersReadWrite ( 3 , kTfLiteFloat32 , " " , { 3 } , TfLiteQuantization ( ) ) ; interpreter . SetInputs ( { 0 , 1 } ) ; interpreter . SetOutputs ( { 3 } ) ; const char * initial_data = " " ; tflite :: ops :: builtin :: BuiltinOpResolverWithoutDefaultDelegates resolver ; TfLiteAddParams * builtin_data = reinterpret_cast < TfLiteAddParams * > ( malloc ( sizeof ( TfLiteAddParams ) ) ) ; builtin_data -> activation = kTfLiteActNone ; builtin_data -> pot_scale_int16 = false ; const TfLiteRegistration * reg = resolver . FindOp ( BuiltinOperator_ADD , 1 ) ; interpreter . AddNodeWithParameters ( { 0 , 1 } , { 2 } , initial_data , 0 , reinterpret_cast < void * > ( builtin_data ) , reg ) ; const TfLiteRegistration * reg2 = resolver . FindOp ( BuiltinOperator_RELU , 1 ) ; interpreter . AddNodeWithParameters ( { 2 } , { 3 } , nullptr , 0 , nullptr , reg2 ) ; const std :: string test_file = CreateFilePath ( " " ) ; SubgraphWriter writer ( & interpreter . primary_subgraph ( ) ) ; EXPECT_EQ ( writer . SetCustomInputOutput ( { 2 } , { 3 } , { 1 } ) , kTfLiteOk ) ; writer . SetUnusedTensors ( { 0 , 1 } ) ; writer . Write ( test_file ) ; std :: unique_ptr < FlatBufferModel > model = FlatBufferModel :: BuildFromFile ( test_file . c_str ( ) ) ; InterpreterBuilder builder ( * model , resolver ) ; std :: unique_ptr < Interpreter > new_interpreter ; builder ( & new_interpreter ) ; ASSERT_EQ ( new_interpreter -> AllocateTensors ( ) , kTfLiteOk ) ; }
Tests writing only a portion of the subgraph.
Tests writing only a portion of the subgraph .
// Add two ops: Add and Relu // Only write the second op. /*inputs=*/ /*outputs=*/ /*execution_plan=*/