Fixed RefCount of Tensor also in TrainSession Eval Mode

mindspore/lite/src/lite_kernel.cc

@@ -38,15 +38,18 @@ bool LiteKernel::IsReady(const std::vector<lite::Tensor *> &scope_tensors) {
   });
 }
 
-void LiteKernel::InitOutTensorInitRefCount() {
+void LiteKernel::InitOutTensorInitRefCount(const std::vector<LiteKernel *> *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);
   }

mindspore/lite/src/lite_kernel.h

@@ -327,7 +327,7 @@ class LiteKernel {
 
   virtual bool IsReady(const std::vector<lite::Tensor *> &in_tensor);
 
-  virtual void InitOutTensorInitRefCount();
+  virtual void InitOutTensorInitRefCount(const std::vector<LiteKernel *> *mask_kernels = nullptr);
 
   KernelKey desc() const { return desc_; }
 
@@ -378,4 +378,4 @@ kernel::InnerKernel *LiteKernelCreator(const std::vector<lite::Tensor *> &inputs
 }
 }  // namespace mindspore::kernel
 
-#endif  // MINDSPORE_LITE_SRC_INNER_KERNEL_H_
+#endif  // MINDSPORE_LITE_SRC_LITE_KERNEL_H_

mindspore/lite/src/lite_kernel_util.cc

@@ -190,7 +190,7 @@ int LiteKernelUtil::TopologicalSortKernels(std::vector<kernel::LiteKernel *> *ke
 
 void LiteKernelUtil::InitTensorInitRefCount(const std::vector<kernel::LiteKernel *> &kernels) {
   for (auto *kernel : kernels) {
-    kernel->InitOutTensorInitRefCount();
+    kernel->InitOutTensorInitRefCount(&kernels);
   }
 }
 

mindspore/lite/src/sub_graph_kernel.cc

@@ -144,9 +144,9 @@ void SubGraphKernel::InitInputTensorInitRefCount() {
   }
 }
 
-void SubGraphKernel::InitOutTensorInitRefCount() {
+void SubGraphKernel::InitOutTensorInitRefCount(const std::vector<LiteKernel *> *mask_kernels) {
   for (auto *node : nodes_) {
-    node->InitOutTensorInitRefCount();
+    node->InitOutTensorInitRefCount(mask_kernels);
   }
 }
 

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 <atomic>
 #include <utility>
@@ -101,7 +101,7 @@ class SubGraphKernel : public LiteKernel {
   // called after Run
   int ReSize() override;
 
-  void InitOutTensorInitRefCount() override;
+  void InitOutTensorInitRefCount(const std::vector<LiteKernel *> *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_

mindspore/lite/src/train/train_session.cc

@@ -24,22 +24,23 @@
 #include <memory>
 #include <map>
 #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<lite::Tensor *>(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<lite::Tensor *>(ms_tensor);
+      lite_tensor->set_init_ref_count(lite_tensor->init_ref_count() + 1);
+    }
+  }
   return RET_OK;
 }