experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Improve code quality in ArenaPlanner's PlanAllocations method.

Browse Source 
* Merge ref-counting / allocation of variable tensors in one go.
* Input seems to be allocated twice, remove the second allocation loop.

PiperOrigin-RevId: 350264826
Change-Id: Icbc6283a6f2357f49dad3992040e286e5424dd92
This commit is contained in:
Haoliang Zhang 2021-01-05 19:05:18 -08:00 committed by TensorFlower Gardener
parent 2d94d904df
commit 173603e851
1 changed files with 7 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
tensorflow/lite/arena_planner.cc
View File

@ -116,7 +116,14 @@ TfLiteStatus ArenaPlanner::PlanAllocations() {
// Variable tensors also should be ensured to be never overwritten and need to // Variable tensors also should be ensured to be never overwritten and need to
// be alive all the time. // be alive all the time.
for (int tensor_index : graph_info_->variables()) { for (int tensor_index : graph_info_->variables()) {
// Increase the reference count for variable tensors by one, so it will
// never be deallocated.
refcounts[tensor_index]++; refcounts[tensor_index]++;
// `variables` is a subgraph-level list and it should never be
// kTfLiteOptionalTensor.
TF_LITE_ENSURE(context_, tensor_index != kTfLiteOptionalTensor);
// Variable tensor should be allocated at the very beginning.
TF_LITE_ENSURE_STATUS(allocate(0, tensor_index));
} }
// Queue all graph inputs for allocation. If preserve_inputs_ is true, make // Queue all graph inputs for allocation. If preserve_inputs_ is true, make
@ -130,15 +137,6 @@ TfLiteStatus ArenaPlanner::PlanAllocations() {
} }
} }
// Queue all graph variable tensors for allocation.
for (int tensor_index : graph_info_->variables()) {
if (tensor_index != kTfLiteOptionalTensor) {
// Increase the reference count for input tensors by one, so it will
// never be deallocated.
TF_LITE_ENSURE_STATUS(allocate(0, tensor_index));
}
}
// Count references to node input tensors. // Count references to node input tensors.
for (size_t i = 0; i < graph_info_->num_execution_nodes(); ++i) { for (size_t i = 0; i < graph_info_->num_execution_nodes(); ++i) {
const TfLiteNode& node = graph_info_->node(i); const TfLiteNode& node = graph_info_->node(i);
@ -151,12 +149,6 @@ TfLiteStatus ArenaPlanner::PlanAllocations() {
} }
} }
// Queue all graph inputs for allocation.
for (int tensor_index : graph_info_->inputs()) {
if (tensor_index != kTfLiteOptionalTensor) {
TF_LITE_ENSURE_STATUS(allocate(0, tensor_index));
}
}
// Go through the graph in execution order. // Go through the graph in execution order.
for (size_t i = 0; i < graph_info_->num_execution_nodes(); ++i) { for (size_t i = 0; i < graph_info_->num_execution_nodes(); ++i) {
const TfLiteNode& node = graph_info_->node(i); const TfLiteNode& node = graph_info_->node(i);