diff options
author | A. Unique TensorFlower <gardener@tensorflow.org> | 2018-03-20 13:02:36 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-03-20 13:04:59 -0700 |
commit | 278ead7d06e427df09f910031cb9195c8a4da559 (patch) | |
tree | 4df3f24f67e2a7e8fbdf1b83606cb4525edb9215 /tensorflow/contrib/lite | |
parent | e4313551d184932c9a135d4edacf42711e5b3483 (diff) |
In allocate_transient_arrays.cc, was not handling the case where the same
array occurs more than once in the list of inputs or outputs of a node.
PiperOrigin-RevId: 189794090
Diffstat (limited to 'tensorflow/contrib/lite')
-rw-r--r-- | tensorflow/contrib/lite/toco/allocate_transient_arrays.cc | 32 |
1 files changed, 26 insertions, 6 deletions
diff --git a/tensorflow/contrib/lite/toco/allocate_transient_arrays.cc b/tensorflow/contrib/lite/toco/allocate_transient_arrays.cc index 49cc1fc2aa..621fbcb98d 100644 --- a/tensorflow/contrib/lite/toco/allocate_transient_arrays.cc +++ b/tensorflow/contrib/lite/toco/allocate_transient_arrays.cc @@ -248,29 +248,49 @@ void AllocateTransientArrays(Model* model, op_index++) { const auto& op = model->operators[op_index]; // Allocate those arrays whose lifespan starts exactly here. + std::vector<string> arrays_to_allocate; for (const auto& input : op->inputs) { if (StartsAt(array_lifespans[input], op_index)) { - AllocateTransientArray(*model, input, &allocator, - transient_data_alignment); + if (std::find(arrays_to_allocate.begin(), arrays_to_allocate.end(), + input) == arrays_to_allocate.end()) { + arrays_to_allocate.push_back(input); + } } } for (const auto& output : op->outputs) { if (StartsAt(array_lifespans[output], op_index)) { - AllocateTransientArray(*model, output, &allocator, - transient_data_alignment); + if (std::find(arrays_to_allocate.begin(), arrays_to_allocate.end(), + output) == arrays_to_allocate.end()) { + arrays_to_allocate.push_back(output); + } } } + for (const string& array : arrays_to_allocate) { + AllocateTransientArray(*model, array, &allocator, + transient_data_alignment); + } + // Deallocate those arrays whose lifespan ends exactly here. + std::vector<string> arrays_to_deallocate; for (const auto& input : op->inputs) { if (EndsAt(array_lifespans[input], op_index)) { - DeallocateTransientArray(*model, input, &allocator); + if (std::find(arrays_to_deallocate.begin(), arrays_to_deallocate.end(), + input) == arrays_to_deallocate.end()) { + arrays_to_deallocate.push_back(input); + } } } for (const auto& output : op->outputs) { if (EndsAt(array_lifespans[output], op_index)) { - DeallocateTransientArray(*model, output, &allocator); + if (std::find(arrays_to_deallocate.begin(), arrays_to_deallocate.end(), + output) == arrays_to_deallocate.end()) { + arrays_to_deallocate.push_back(output); + } } } + for (const string& array : arrays_to_deallocate) { + DeallocateTransientArray(*model, array, &allocator); + } } // Just out of curiosity (not used in the actual allocation process) |