fix model pool memory leak

mindspore/lite/src/cxx_api/model_pool/model_pool.cc

@@ -220,44 +220,30 @@ Status ModelPool::SplitInputTensorByBatch(const std::vector<MSTensor> &inputs,
           MS_LOG(ERROR) << "malloc size is wrong.";
           return kLiteError;
         }
-        void *data = malloc(input_size * sizeof(float));
+        auto data =
-        if (data == nullptr) {
+          reinterpret_cast<float *>(const_cast<MSTensor &>(inputs[i]).MutableData()) + input_data_split_size[i];
-          MS_LOG(ERROR) << "malloc failed.";
+        auto new_tensor = MSTensor(inputs[i].Name(), static_cast<enum DataType>(kNumberTypeFloat32), shape, data,
-          return kLiteError;
+                                   input_size * sizeof(float));
-        }
-        memcpy(reinterpret_cast<float *>(data),
-               reinterpret_cast<float *>(const_cast<MSTensor &>(inputs[i]).MutableData()) + input_data_split_size[i],
-               input_size * sizeof(float));
-        auto new_tensor = mindspore::MSTensor::CreateTensor(
-          inputs[i].Name(), static_cast<enum DataType>(kNumberTypeFloat32), shape, data, input_size * sizeof(float));
         if (new_tensor == nullptr) {
           MS_LOG(ERROR) << "create tensor failed.";
           return kLiteError;
         }
-        new_inputs_tensor.push_back(*new_tensor);
+        new_inputs_tensor.push_back(new_tensor);
-        free(data);
         input_data_split_size[i] += input_size;
       } else if (inputs[i].DataType() == static_cast<enum DataType>(kNumberTypeInt32)) {
         if (input_size * sizeof(int32_t) > MAX_MALLOC_SIZE) {
           MS_LOG(ERROR) << "malloc size is wrong.";
           return kLiteError;
         }
-        void *data = malloc(input_size * sizeof(int32_t));
+        auto data =
-        if (data == nullptr) {
+          reinterpret_cast<int32_t *>(const_cast<MSTensor &>(inputs[i]).MutableData()) + input_data_split_size[i];
-          MS_LOG(ERROR) << "malloc failed.";
+        auto new_tensor = MSTensor(inputs[i].Name(), static_cast<enum DataType>(kNumberTypeInt32), shape, data,
-          return kLiteError;
+                                   input_size * sizeof(int32_t));
-        }
-        memcpy(reinterpret_cast<int32_t *>(data),
-               reinterpret_cast<int32_t *>(const_cast<MSTensor &>(inputs[i]).MutableData()) + input_data_split_size[i],
-               input_size * sizeof(int32_t));
-        auto new_tensor = mindspore::MSTensor::CreateTensor(
-          inputs[i].Name(), static_cast<enum DataType>(kNumberTypeInt32), shape, data, input_size * sizeof(int32_t));
         if (new_tensor == nullptr) {
           MS_LOG(ERROR) << "create tensor failed.";
           return kLiteError;
         }
-        new_inputs_tensor.push_back(*new_tensor);
+        new_inputs_tensor.push_back(new_tensor);
-        free(data);
         input_data_split_size[i] += input_size;
       } else {
         MS_LOG(ERROR) << "not support data type in split batch.";
@@ -323,9 +309,10 @@ Status ModelPool::ConcatPredictOutput(std::vector<std::vector<MSTensor>> *output
       return kLiteError;
     }
     for (size_t j = 0; j < outputs->size(); j++) {
-      void *out_data = outputs->at(j)[i].MutableData();
+      void *out_data = outputs->at(j).at(i).MutableData();
       if (out_data == nullptr) {
         free(all_out_data);
+        all_out_data = nullptr;
         MS_LOG(ERROR) << "output data is nullptr.";
         return kLiteError;
       }
@@ -343,6 +330,26 @@ Status ModelPool::ConcatPredictOutput(std::vector<std::vector<MSTensor>> *output
       all_out_data = nullptr;
     }
     new_outputs->push_back(*new_tensor);
+    delete new_tensor;
+  }
+  return kSuccess;
+}
+
+Status ModelPool::FreeSplitTensor(std::vector<std::vector<MSTensor>> *new_inputs,
+                                  std::vector<std::vector<MSTensor>> *new_outputs) {
+  for (size_t i = 0; i < new_inputs->size(); i++) {
+    for (size_t j = 0; j < new_inputs->at(i).size(); j++) {
+      new_inputs->at(i).at(j).SetData(nullptr);
+    }
+  }
+  new_inputs->clear();
+  if (is_user_data_) {
+    for (size_t i = 0; i < new_outputs->size(); i++) {
+      for (size_t j = 0; j < new_outputs->at(i).size(); j++) {
+        new_outputs->at(i).at(j).SetData(nullptr);
+      }
+    }
+    new_outputs->clear();
   }
   return kSuccess;
 }
@@ -383,12 +390,10 @@ Status ModelPool::Predict(const std::vector<MSTensor> &inputs, std::vector<MSTen
       MS_LOG(ERROR) << "ConcatPredictOutput failed.";
       return kLiteError;
     }
-    if (is_user_data_) {
+    status = FreeSplitTensor(&new_inputs, &new_outputs);
-      for (size_t i = 0; i < batch_split_num; i++) {
+    if (status != kSuccess) {
-        for (size_t j = 0; j < new_outputs.size(); j++) {
+      MS_LOG(ERROR) << "free split tensor failed.";
-          new_outputs.at(i).at(j).SetData(nullptr);
+      return kLiteError;
-        }
-      }
     }
   } else {
     if (wait_model_num == 1) {

mindspore/lite/src/cxx_api/model_pool/model_pool.h

@@ -52,6 +52,8 @@ class ModelPool {
   Status SplitOutputTensorByBatch(std::vector<std::vector<MSTensor>> *outputs, std::vector<MSTensor> *new_outputs,
                                   size_t batch_split_num);
   Status ConcatPredictOutput(std::vector<std::vector<MSTensor>> *outputs, std::vector<MSTensor> *new_outputs);
+  Status FreeSplitTensor(std::vector<std::vector<MSTensor>> *new_inputs,
+                         std::vector<std::vector<MSTensor>> *new_outputs);
 
   std::vector<std::thread> model_thread_vec_;
   std::vector<MSTensor> model_inputs_;

mindspore/lite/src/cxx_api/model_pool/model_thread.cc

@@ -35,6 +35,7 @@ void ModelThread::Run() {
       continue;
     }
     if (is_copy_output_) {
+      std::vector<MSTensor> new_outputs;
       auto output_size = outputs->size();
       for (size_t i = 0; i < output_size; i++) {
         auto copy_tensor =
@@ -46,9 +47,11 @@ void ModelThread::Run() {
           PredictTaskQueue::GetInstance()->ActiveTask();
           continue;
         }
-        outputs->erase(outputs->begin());
+        new_outputs.push_back(*copy_tensor);
-        outputs->push_back(*copy_tensor);
+        delete copy_tensor;
       }
+      outputs->clear();
+      outputs->insert(outputs->end(), new_outputs.begin(), new_outputs.end());
     }
     task->ready = true;
     PredictTaskQueue::GetInstance()->ActiveTask();

mindspore/lite/tools/benchmark/benchmark_unified_api.cc

@@ -247,6 +247,7 @@ int BenchmarkUnifiedApi::GenerateInputData() {
       auto new_tensor =
         mindspore::MSTensor::CreateTensor(tensor_name, ms_inputs_for_api_[i].DataType(), shape, input_data, size);
       inputs.push_back(*new_tensor);
+      delete new_tensor;
     }
     all_inputs_.push_back(inputs);
     return RET_OK;
@@ -324,6 +325,7 @@ int BenchmarkUnifiedApi::ReadInputFile() {
       auto new_tensor =
         mindspore::MSTensor::CreateTensor(tensor_name, ms_inputs_for_api_[i].DataType(), shape, input_data, size);
       inputs.push_back(*new_tensor);
+      delete new_tensor;
     }
     all_inputs_.push_back(inputs);
     return RET_OK;