Improve EnsureTensorsVectorCapacity performance.

PiperOrigin-RevId: 307325783
Change-Id: I60fa334f04da194dd31e2d5c82f7f1d68c8d3a7a
This commit is contained in:
Yu-Cheng Ling 2020-04-19 20:15:10 -07:00 committed by TensorFlower Gardener
parent 871f12a7b6
commit 62a6012a0c
2 changed files with 12 additions and 3 deletions

View File

@ -558,7 +558,13 @@ class Subgraph {
void EnsureTensorsVectorCapacity() {
const size_t required_capacity = tensors_.size() + kTensorsCapacityHeadroom;
if (required_capacity > tensors_.capacity()) {
tensors_.reserve(required_capacity);
// Whenever it's required to increase the vector capacity, make it at
// least twice bigger. The behavior is consistent with the default
// behavior of GCC STL's `std::vector::resize()`. This avoids frequently
// allocating and copying the underlying buffer.
size_t reserved_capacity =
std::max(required_capacity, tensors_.capacity() * 2);
tensors_.reserve(reserved_capacity);
context_.tensors = tensors_.data();
}
}

View File

@ -1079,8 +1079,11 @@ TEST(InterpreterTensorsCapacityTest, TestExceedHeadroom) {
TfLiteTensor* first_tensor = context->tensors;
int new_tensor_index;
context->AddTensors(context, Interpreter::kTensorsCapacityHeadroom + 1,
&new_tensor_index);
// Add enough tensors to trigger buffer re-allocation.
context->AddTensors(
context,
(context->tensors_size + Interpreter::kTensorsCapacityHeadroom + 1) * 2,
&new_tensor_index);
EXPECT_NE(first_tensor, context->tensors);
return kTfLiteOk;
};