Merge pull request #43538 from mansnils:scratch_tensors
PiperOrigin-RevId: 338767807 Change-Id: I2e1962f773872cb94abbe21881b57a1144cdf645
This commit is contained in:
commit
b35dceb805
tensorflow/lite/micro
@ -229,21 +229,6 @@ TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph,
|
||||
for (size_t n = 0; n < op->inputs()->size(); ++n) {
|
||||
const int tensor_index = op->inputs()->Get(n);
|
||||
AllocationInfo* current = &info_[tensor_index];
|
||||
|
||||
// TODO(b/166484865): Figure out a more general solution.
|
||||
// This workaround is needed to handle situations where subgraph input !=
|
||||
// operator input.
|
||||
// In case operator input(s) are not in subgraph inputs initialize them.
|
||||
if (current->first_created == 0) {
|
||||
for (size_t op_input = 0; op_input < op->inputs()->size(); ++op_input) {
|
||||
const int op_tensor_index = op->inputs()->Get(op_input);
|
||||
AllocationInfo* op_current = &info_[op_tensor_index];
|
||||
if (op_current->needs_allocating && op_current->first_created == -1) {
|
||||
op_current->first_created = i;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (((current->last_used == -1) || (current->last_used < i))) {
|
||||
current->last_used = i;
|
||||
}
|
||||
@ -257,16 +242,15 @@ TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph,
|
||||
}
|
||||
}
|
||||
|
||||
// Work out which tensors need to be allocated.
|
||||
// Sanity check for valid tensor lifetime.
|
||||
for (size_t i = 0; i < tensor_count_; ++i) {
|
||||
AllocationInfo* current = &info_[i];
|
||||
const bool is_read_only =
|
||||
// Even though tensor appears to be read only it may still need to be
|
||||
// allocated.
|
||||
const bool appears_read_only =
|
||||
(current->first_created == -1) && (current->last_used != -1);
|
||||
if (is_read_only) {
|
||||
current->needs_allocating = false;
|
||||
}
|
||||
const bool has_partial_lifetime =
|
||||
!is_read_only &&
|
||||
!appears_read_only &&
|
||||
((current->first_created == -1) || (current->last_used == -1));
|
||||
if (has_partial_lifetime && current->needs_allocating) {
|
||||
TF_LITE_REPORT_ERROR(
|
||||
|
@ -28,6 +28,12 @@ namespace testing {
|
||||
namespace {
|
||||
|
||||
constexpr int kExpectedAlignment = 4;
|
||||
constexpr int t0 = 0;
|
||||
constexpr int t1 = 1;
|
||||
constexpr int t2 = 2;
|
||||
constexpr int t3 = 3;
|
||||
constexpr int t4 = 4;
|
||||
constexpr int t5 = 5;
|
||||
|
||||
void VerifyMockTfLiteTensor(TfLiteTensor* tensor, bool is_variable = false) {
|
||||
TF_LITE_MICRO_EXPECT_EQ(kTfLiteInt32, tensor->type);
|
||||
@ -446,18 +452,18 @@ TF_LITE_MICRO_TEST(TestAllocationForComplexModelAllocation) {
|
||||
TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) {
|
||||
int version = 1;
|
||||
int subgraph = 0;
|
||||
constexpr int nbr_tensors = 4;
|
||||
constexpr int number_tensors = 4;
|
||||
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
|
||||
tflite::NodeAndRegistration* node_and_registration;
|
||||
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
|
||||
nbr_tensors] = {version, subgraph,
|
||||
nbr_tensors, // header
|
||||
// memory offsets:
|
||||
-1, -1, -1, -1};
|
||||
number_tensors] = {version, subgraph,
|
||||
number_tensors, // header
|
||||
// memory offsets:
|
||||
-1, -1, -1, -1};
|
||||
|
||||
// The structure is identical to the one in
|
||||
// TestAllocationForModelsWithBranches
|
||||
int num_conns = 3;
|
||||
int number_connections = 3;
|
||||
tflite::testing::NodeConnection node_list[3] = {{
|
||||
{0}, // input
|
||||
{1} // output
|
||||
@ -472,7 +478,7 @@ TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) {
|
||||
}};
|
||||
|
||||
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
|
||||
nbr_tensors, metadata_buffer, node_list, num_conns);
|
||||
number_tensors, metadata_buffer, node_list, number_connections);
|
||||
|
||||
TfLiteEvalTensor* eval_tensors = nullptr;
|
||||
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
|
||||
@ -506,37 +512,26 @@ TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) {
|
||||
}
|
||||
|
||||
TF_LITE_MICRO_TEST(OfflinePlannerBasic) {
|
||||
constexpr int nbr_tensors = 4;
|
||||
constexpr int number_tensors = 4;
|
||||
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
|
||||
tflite::NodeAndRegistration* node_and_registration;
|
||||
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
|
||||
nbr_tensors] = {1, 0, nbr_tensors,
|
||||
0, // t0
|
||||
48, // t1
|
||||
0, // t2
|
||||
48}; // t3
|
||||
|
||||
int t0 = 0;
|
||||
int t1 = 1;
|
||||
int t2 = 2;
|
||||
int t3 = 3;
|
||||
|
||||
int num_conns = 3;
|
||||
tflite::testing::NodeConnection node_list[3] = {{
|
||||
{t0}, // input
|
||||
{t1} // output
|
||||
},
|
||||
{
|
||||
{t1}, // input
|
||||
{t2} // output
|
||||
},
|
||||
{
|
||||
{t2}, // input
|
||||
{t3} // output
|
||||
}};
|
||||
number_tensors] = {1, 0, number_tensors,
|
||||
/*t0=*/0,
|
||||
/*t1=*/48,
|
||||
/*t2=*/0,
|
||||
/*t3=*/48};
|
||||
constexpr int number_connections = 3;
|
||||
tflite::testing::NodeConnection node_list[number_connections] = {
|
||||
{/*input=*/{tflite::testing::t0},
|
||||
/*output=*/{tflite::testing::t1}},
|
||||
{/*input=*/{tflite::testing::t1},
|
||||
/*output=*/{tflite::testing::t2}},
|
||||
{/*input=*/{tflite::testing::t2},
|
||||
/*output=*/{tflite::testing::t3}}};
|
||||
|
||||
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
|
||||
nbr_tensors, metadata_buffer, node_list, num_conns);
|
||||
number_tensors, metadata_buffer, node_list, number_connections);
|
||||
|
||||
TfLiteEvalTensor* eval_tensors = nullptr;
|
||||
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
|
||||
@ -561,37 +556,28 @@ TF_LITE_MICRO_TEST(OfflinePlannerBasic) {
|
||||
}
|
||||
|
||||
TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) {
|
||||
constexpr int nbr_tensors = 4;
|
||||
constexpr int number_tensors = 4;
|
||||
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
|
||||
tflite::NodeAndRegistration* node_and_registration;
|
||||
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
|
||||
nbr_tensors] = {
|
||||
1, 0, nbr_tensors, // header: version, subgraph, nbr tensors
|
||||
// memory offsets:
|
||||
0, // t0
|
||||
0, // t1
|
||||
48, // t2
|
||||
-1}; // t3
|
||||
number_tensors] = {/*version=*/1,
|
||||
/*subgraph=*/0,
|
||||
number_tensors,
|
||||
/*t0=*/0,
|
||||
/*t1=*/0,
|
||||
/*t2=*/48,
|
||||
/*t3=*/-1};
|
||||
|
||||
int t0 = 0;
|
||||
int t1 = 1;
|
||||
int t2 = 2;
|
||||
int t3 = 3;
|
||||
|
||||
int num_conns = 2;
|
||||
int number_connections = 2;
|
||||
tflite::testing::NodeConnection node_list[2] = {
|
||||
{
|
||||
{t0, t1}, // input, scratch
|
||||
{t2} // output
|
||||
},
|
||||
{
|
||||
{t2}, // input
|
||||
{t3} // output
|
||||
},
|
||||
{/*input, scratch=*/{tflite::testing::t0, tflite::testing::t1},
|
||||
/*output=*/{tflite::testing::t2}},
|
||||
{/*input=*/{tflite::testing::t2},
|
||||
/*output=*/{tflite::testing::t3}},
|
||||
};
|
||||
|
||||
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
|
||||
nbr_tensors, metadata_buffer, node_list, num_conns);
|
||||
number_tensors, metadata_buffer, node_list, number_connections);
|
||||
|
||||
TfLiteEvalTensor* eval_tensors = nullptr;
|
||||
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
|
||||
@ -617,39 +603,33 @@ TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) {
|
||||
}
|
||||
|
||||
TF_LITE_MICRO_TEST(OfflinePlannerOfflineOnline) {
|
||||
constexpr int nbr_tensors = 5;
|
||||
constexpr int number_tensors = 5;
|
||||
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
|
||||
tflite::NodeAndRegistration* node_and_registration;
|
||||
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
|
||||
nbr_tensors] = {
|
||||
1, 0, nbr_tensors, // header: version, subgraph, nbr tensors
|
||||
// memory offsets:
|
||||
0, // t0
|
||||
48, // t1
|
||||
-1, // t2
|
||||
0, // t3
|
||||
-1}; // t4
|
||||
number_tensors] = {/*version=*/1,
|
||||
/*subgraph=*/0,
|
||||
number_tensors,
|
||||
/*t0=*/0,
|
||||
/*t1=*/48,
|
||||
/*t2=*/-1,
|
||||
/*t3=*/0,
|
||||
/*t4=*/-1};
|
||||
|
||||
int t0 = 0;
|
||||
int t1 = 1;
|
||||
int t2 = 2;
|
||||
int t3 = 3;
|
||||
int t4 = 4;
|
||||
|
||||
int num_conns = 2;
|
||||
tflite::testing::NodeConnection node_list[2] = {
|
||||
constexpr int number_connections = 2;
|
||||
tflite::testing::NodeConnection node_list[number_connections] = {
|
||||
{
|
||||
{t0, t1}, // input, scratch
|
||||
{t2}, // output
|
||||
/*input, scratch=*/{tflite::testing::t0, tflite::testing::t1},
|
||||
/*output=*/{tflite::testing::t2},
|
||||
},
|
||||
{
|
||||
{t2}, // input
|
||||
{t3, t4}, // output1, output2
|
||||
/*input=*/{tflite::testing::t2},
|
||||
/*output1, output2=*/{tflite::testing::t3, tflite::testing::t4},
|
||||
},
|
||||
};
|
||||
|
||||
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
|
||||
nbr_tensors, metadata_buffer, node_list, num_conns);
|
||||
number_tensors, metadata_buffer, node_list, number_connections);
|
||||
|
||||
TfLiteEvalTensor* eval_tensors = nullptr;
|
||||
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
|
||||
@ -757,43 +737,34 @@ TF_LITE_MICRO_TEST(TestAllocateTfLiteTensorWithReset) {
|
||||
}
|
||||
|
||||
TF_LITE_MICRO_TEST(TestOperatorInputsNotInSubgraphInputs) {
|
||||
constexpr int nbr_tensors = 5;
|
||||
constexpr int number_tensors = 5;
|
||||
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
|
||||
tflite::NodeAndRegistration* node_and_registration;
|
||||
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
|
||||
nbr_tensors] = {
|
||||
1, 0, nbr_tensors, // header: version, subgraph, nbr tensors
|
||||
// memory offsets:
|
||||
0, // t0
|
||||
0, // t1
|
||||
0, // t2
|
||||
48, // t3
|
||||
-1}; // t4
|
||||
number_tensors] = {/*version=*/1,
|
||||
/*subgraph=*/0,
|
||||
number_tensors,
|
||||
/*t0=*/0,
|
||||
/*t1=*/0,
|
||||
/*t2=*/0,
|
||||
/*t3=*/48,
|
||||
/*t4=*/-1};
|
||||
|
||||
int t0 = 0;
|
||||
int t1 = 1;
|
||||
int t2 = 2;
|
||||
int t3 = 3;
|
||||
int t4 = 4;
|
||||
|
||||
int num_conns = 2;
|
||||
tflite::testing::NodeConnection node_list[2] = {
|
||||
{
|
||||
{t0, t1, t2}, // t0: input (actual input part of subgraph inputs as
|
||||
// well as operator inputs)
|
||||
// t1: scratch1 (only in operator inputs)
|
||||
// t2: scratch2 (only in operator inputs)
|
||||
{t3} // output
|
||||
},
|
||||
{
|
||||
{t3}, // input
|
||||
{t4} // output
|
||||
},
|
||||
constexpr int number_connections = 2;
|
||||
tflite::testing::NodeConnection node_list[number_connections] = {
|
||||
{// t0: input (actual input part of subgraph inputs as
|
||||
// well as operator inputs)
|
||||
// t1: scratch1 (only in operator inputs)
|
||||
// t2: scratch2 (only in operator inputs)
|
||||
{tflite::testing::t0, tflite::testing::t1, tflite::testing::t2},
|
||||
/*t3: output=*/{tflite::testing::t3}},
|
||||
{/*t3: input=*/{tflite::testing::t3},
|
||||
/*t4: output=*/{tflite::testing::t4}},
|
||||
};
|
||||
|
||||
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
|
||||
nbr_tensors, metadata_buffer, node_list, num_conns,
|
||||
1 /* only first tensor (t0) is in subgraph input list*/);
|
||||
number_tensors, metadata_buffer, node_list, number_connections,
|
||||
/*Only first tensor (t0) is in subgraph input list=*/1);
|
||||
|
||||
TfLiteEvalTensor* eval_tensors = nullptr;
|
||||
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
|
||||
@ -818,4 +789,61 @@ TF_LITE_MICRO_TEST(TestOperatorInputsNotInSubgraphInputs) {
|
||||
TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[4].data.uint8 - start);
|
||||
}
|
||||
|
||||
TF_LITE_MICRO_TEST(TestTypicalFirstOpAndSecondOpWithScratchTensors) {
|
||||
constexpr int number_tensors = 6;
|
||||
tflite::AllOpsResolver op_resolver = tflite::testing::GetOpResolver();
|
||||
tflite::NodeAndRegistration* node_and_registration;
|
||||
const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize +
|
||||
number_tensors] = {/*version=*/1,
|
||||
/*subgraph=*/0,
|
||||
number_tensors,
|
||||
/*t0=*/0,
|
||||
/*t1=*/0,
|
||||
/*t2=*/0,
|
||||
/*t3=*/0,
|
||||
/*t4=*/48,
|
||||
/*t5=*/-1};
|
||||
|
||||
constexpr int number_connections = 3;
|
||||
tflite::testing::NodeConnection node_list[number_connections] = {
|
||||
{/*t0: input (subgraph and operator input)=*/{tflite::testing::t0},
|
||||
/*t1: output=*/{tflite::testing::t1}},
|
||||
{// t1: input
|
||||
// t2: scratch1 (only in operator inputs)
|
||||
// t3: scratch2 (only in operator inputs)
|
||||
{tflite::testing::t1, tflite::testing::t2, tflite::testing::t3},
|
||||
|
||||
/*t4: output=*/{tflite::testing::t4}},
|
||||
{/*t4: input=*/{tflite::testing::t4},
|
||||
/*t5: output=*/{tflite::testing::t5}},
|
||||
};
|
||||
|
||||
const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning(
|
||||
number_tensors, metadata_buffer, node_list, number_connections,
|
||||
/*Only first tensor (t0) is in subgraph input list=*/1);
|
||||
|
||||
TfLiteEvalTensor* eval_tensors = nullptr;
|
||||
tflite::ScratchBufferHandle* scratch_buffer_handles = nullptr;
|
||||
constexpr size_t arena_size = 4096;
|
||||
uint8_t arena[arena_size];
|
||||
tflite::MicroAllocator* allocator =
|
||||
tflite::MicroAllocator::Create(arena, arena_size, micro_test::reporter);
|
||||
|
||||
TF_LITE_MICRO_EXPECT_EQ(
|
||||
kTfLiteOk,
|
||||
allocator->StartModelAllocation(model, op_resolver,
|
||||
&node_and_registration, &eval_tensors));
|
||||
TF_LITE_MICRO_EXPECT_EQ(
|
||||
kTfLiteOk, allocator->FinishModelAllocation(model, eval_tensors,
|
||||
&scratch_buffer_handles));
|
||||
|
||||
uint8_t* start = eval_tensors[0].data.uint8;
|
||||
TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[0].data.uint8 - start);
|
||||
TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[1].data.uint8 - start);
|
||||
TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[2].data.uint8 - start);
|
||||
TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[3].data.uint8 - start);
|
||||
TF_LITE_MICRO_EXPECT_EQ(48, eval_tensors[4].data.uint8 - start);
|
||||
TF_LITE_MICRO_EXPECT_EQ(0, eval_tensors[5].data.uint8 - start);
|
||||
}
|
||||
|
||||
TF_LITE_MICRO_TESTS_END
|
||||
|
@ -206,7 +206,7 @@ const Model* ModelBuilder::BuildModel(
|
||||
} else {
|
||||
// A non-zero value of num_subgraph_inputs means that some of
|
||||
// the operator input tensors are not subgraph inputs.
|
||||
TFLITE_DCHECK(num_subgraph_inputs < inputs.size());
|
||||
TFLITE_DCHECK(num_subgraph_inputs <= inputs.size());
|
||||
}
|
||||
|
||||
const flatbuffers::Offset<SubGraph> subgraphs[subgraphs_size] = {
|
||||
|
Loading…
Reference in New Issue
Block a user