From d0f9158c1ee6abd5ec0b6dfd60d1cdd8ea57b8c3 Mon Sep 17 00:00:00 2001 From: Emir Haleva Date: Sun, 1 Aug 2021 16:59:32 +0300 Subject: [PATCH] Fixed RefCount of Tensor also in TrainSession Eval Mode --- mindspore/lite/src/lite_kernel.cc | 13 +++++---- mindspore/lite/src/lite_kernel.h | 4 +-- mindspore/lite/src/lite_kernel_util.cc | 2 +- mindspore/lite/src/sub_graph_kernel.cc | 4 +-- mindspore/lite/src/sub_graph_kernel.h | 8 +++--- mindspore/lite/src/train/train_session.cc | 35 ++++++++++++++++------- 6 files changed, 42 insertions(+), 24 deletions(-) diff --git a/mindspore/lite/src/lite_kernel.cc b/mindspore/lite/src/lite_kernel.cc index 926a94f3bfd..db1ad97e1d0 100644 --- a/mindspore/lite/src/lite_kernel.cc +++ b/mindspore/lite/src/lite_kernel.cc @@ -38,15 +38,18 @@ bool LiteKernel::IsReady(const std::vector &scope_tensors) { }); } -void LiteKernel::InitOutTensorInitRefCount() { +void LiteKernel::InitOutTensorInitRefCount(const std::vector *mask_kernels) { for (auto *tensor : this->out_tensors()) { MS_ASSERT(tensor != nullptr); size_t init_ref_count = 0; for (auto *post_kernel : this->out_kernels_) { - auto &post_in_tensors = post_kernel->in_tensors(); - init_ref_count += - std::count_if(post_in_tensors.begin(), post_in_tensors.end(), - [&tensor](const lite::Tensor *post_kernel_in_tensor) { return post_kernel_in_tensor == tensor; }); + if ((mask_kernels == nullptr) || + std::find(mask_kernels->begin(), mask_kernels->end(), post_kernel) != mask_kernels->end()) { + auto &post_in_tensors = post_kernel->in_tensors(); + init_ref_count += std::count_if( + post_in_tensors.begin(), post_in_tensors.end(), + [&tensor](const lite::Tensor *post_kernel_in_tensor) { return post_kernel_in_tensor == tensor; }); + } } tensor->set_init_ref_count(init_ref_count); } diff --git a/mindspore/lite/src/lite_kernel.h b/mindspore/lite/src/lite_kernel.h index 55456c46c05..a499d95e9a9 100644 --- a/mindspore/lite/src/lite_kernel.h +++ b/mindspore/lite/src/lite_kernel.h @@ -327,7 +327,7 @@ class LiteKernel { virtual bool IsReady(const std::vector &in_tensor); - virtual void InitOutTensorInitRefCount(); + virtual void InitOutTensorInitRefCount(const std::vector *mask_kernels = nullptr); KernelKey desc() const { return desc_; } @@ -378,4 +378,4 @@ kernel::InnerKernel *LiteKernelCreator(const std::vector &inputs } } // namespace mindspore::kernel -#endif // MINDSPORE_LITE_SRC_INNER_KERNEL_H_ +#endif // MINDSPORE_LITE_SRC_LITE_KERNEL_H_ diff --git a/mindspore/lite/src/lite_kernel_util.cc b/mindspore/lite/src/lite_kernel_util.cc index d3b2df08187..e56f8860506 100644 --- a/mindspore/lite/src/lite_kernel_util.cc +++ b/mindspore/lite/src/lite_kernel_util.cc @@ -190,7 +190,7 @@ int LiteKernelUtil::TopologicalSortKernels(std::vector *ke void LiteKernelUtil::InitTensorInitRefCount(const std::vector &kernels) { for (auto *kernel : kernels) { - kernel->InitOutTensorInitRefCount(); + kernel->InitOutTensorInitRefCount(&kernels); } } diff --git a/mindspore/lite/src/sub_graph_kernel.cc b/mindspore/lite/src/sub_graph_kernel.cc index c75b955fea9..ce85591207e 100644 --- a/mindspore/lite/src/sub_graph_kernel.cc +++ b/mindspore/lite/src/sub_graph_kernel.cc @@ -144,9 +144,9 @@ void SubGraphKernel::InitInputTensorInitRefCount() { } } -void SubGraphKernel::InitOutTensorInitRefCount() { +void SubGraphKernel::InitOutTensorInitRefCount(const std::vector *mask_kernels) { for (auto *node : nodes_) { - node->InitOutTensorInitRefCount(); + node->InitOutTensorInitRefCount(mask_kernels); } } diff --git a/mindspore/lite/src/sub_graph_kernel.h b/mindspore/lite/src/sub_graph_kernel.h index 0200b2ebd8b..1671fbdb414 100644 --- a/mindspore/lite/src/sub_graph_kernel.h +++ b/mindspore/lite/src/sub_graph_kernel.h @@ -14,8 +14,8 @@ * limitations under the License. */ -#ifndef MINDSPORE_LITE_SRC_SUB_GRAPH_H -#define MINDSPORE_LITE_SRC_SUB_GRAPH_H +#ifndef MINDSPORE_LITE_SRC_SUB_GRAPH_KERNEL_H_ +#define MINDSPORE_LITE_SRC_SUB_GRAPH_KERNEL_H_ #include #include @@ -101,7 +101,7 @@ class SubGraphKernel : public LiteKernel { // called after Run int ReSize() override; - void InitOutTensorInitRefCount() override; + void InitOutTensorInitRefCount(const std::vector *mask_kernels) override; void InitInputTensorInitRefCount(); @@ -226,4 +226,4 @@ class CustomSubGraph : public SubGraphKernel { int Execute(const KernelCallBack &before, const KernelCallBack &after) override; }; } // namespace mindspore::kernel -#endif // MINDSPORE_LITE_SRC_SUB_GRAPH_H +#endif // MINDSPORE_LITE_SRC_SUB_GRAPH_KERNEL_H_ diff --git a/mindspore/lite/src/train/train_session.cc b/mindspore/lite/src/train/train_session.cc index 8a5514be5a7..2cf97f2718d 100644 --- a/mindspore/lite/src/train/train_session.cc +++ b/mindspore/lite/src/train/train_session.cc @@ -24,22 +24,23 @@ #include #include #include "include/errorcode.h" -#include "src/common/utils.h" -#include "src/tensor.h" -#include "src/lite_model.h" -#include "src/train/loss_kernel.h" -#include "src/train/optimizer_kernel.h" -#include "src/sub_graph_kernel.h" -#include "src/train/train_populate_parameter.h" -#include "src/train/train_populate_parameter_v0.h" #include "src/executor.h" +#include "src/lite_model.h" +#include "src/lite_kernel_util.h" +#include "src/sub_graph_kernel.h" +#include "src/tensor.h" #include "src/kernel_registry.h" +#include "src/common/prim_util.h" +#include "src/common/tensor_util.h" +#include "src/common/utils.h" #include "src/runtime/kernel/arm/fp32_grad/convolution.h" #include "src/runtime/kernel/arm/fp32/batchnorm_fp32.h" -#include "src/common/tensor_util.h" +#include "src/train/loss_kernel.h" +#include "src/train/optimizer_kernel.h" #include "src/train/train_utils.h" #include "src/train/train_export.h" -#include "src/common/prim_util.h" +#include "src/train/train_populate_parameter.h" +#include "src/train/train_populate_parameter_v0.h" namespace mindspore { namespace lite { @@ -412,6 +413,13 @@ int TrainSession::Train() { output_node_map_ = train_output_node_map_; output_tensor_map_ = train_output_tensor_map_; output_tensor_names_ = train_output_tensor_names_; + kernel::LiteKernelUtil::InitTensorInitRefCount(train_kernels_); + for (auto &ms_tensors : eval_output_node_map_) { // Allow to look at prediction also during training + for (auto &ms_tensor : ms_tensors.second) { + lite::Tensor *lite_tensor = static_cast(ms_tensor); + lite_tensor->set_init_ref_count(lite_tensor->init_ref_count() + 1); + } + } return RET_OK; } @@ -431,6 +439,13 @@ int TrainSession::Eval() { output_node_map_ = eval_output_node_map_; output_tensor_map_ = eval_output_tensor_map_; output_tensor_names_ = eval_output_tensor_names_; + kernel::LiteKernelUtil::InitTensorInitRefCount(inference_kernels_); + for (auto &ms_tensors : eval_output_node_map_) { + for (auto &ms_tensor : ms_tensors.second) { + lite::Tensor *lite_tensor = static_cast(ms_tensor); + lite_tensor->set_init_ref_count(lite_tensor->init_ref_count() + 1); + } + } return RET_OK; }