fix model pool memory leak

This commit is contained in:
yefeng 2022-02-16 09:56:00 +08:00
parent 470bb27fd9
commit a4f796b0c1
4 changed files with 45 additions and 33 deletions

View File

@ -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));
if (data == nullptr) {
MS_LOG(ERROR) << "malloc failed.";
return kLiteError;
}
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));
auto data =
reinterpret_cast<float *>(const_cast<MSTensor &>(inputs[i]).MutableData()) + input_data_split_size[i];
auto new_tensor = MSTensor(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);
free(data);
new_inputs_tensor.push_back(new_tensor);
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));
if (data == nullptr) {
MS_LOG(ERROR) << "malloc failed.";
return kLiteError;
}
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));
auto data =
reinterpret_cast<int32_t *>(const_cast<MSTensor &>(inputs[i]).MutableData()) + input_data_split_size[i];
auto new_tensor = MSTensor(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);
free(data);
new_inputs_tensor.push_back(new_tensor);
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_) {
for (size_t i = 0; i < batch_split_num; i++) {
for (size_t j = 0; j < new_outputs.size(); j++) {
new_outputs.at(i).at(j).SetData(nullptr);
}
}
status = FreeSplitTensor(&new_inputs, &new_outputs);
if (status != kSuccess) {
MS_LOG(ERROR) << "free split tensor failed.";
return kLiteError;
}
} else {
if (wait_model_num == 1) {

View File

@ -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_;

View File

@ -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());
outputs->push_back(*copy_tensor);
new_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();

View File

@ -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;