p47068915
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!20584 [MS][LITE] fix lite train demo memory increase along training 

Merge pull request !20584 from zhengjun10/master
This commit is contained in:
i-robot 2021-07-21 08:07:51 +00:00 committed by Gitee
parent 937bfeaf00 f094444e63
commit 2f65ad73ea
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mindspore/lite/src/train/train_loop.cc
View File

@ -42,11 +42,11 @@ int TrainLoop::Train(int epochs, Dataset *ds, std::vector<session::TrainLoopCall
for (auto cb : cbs) cb->Begin(cb_data);
std::shared_ptr<Iterator> iter = ds->CreateIterator();
for (int i = 0; i < epochs; i++) {
cb_data.epoch_ = epoch_++;
for (auto cb : cbs) cb->EpochBegin(cb_data);
std::shared_ptr<Iterator> iter = ds->CreateIterator();
MSTensorVec row_vec;
int s = 0;
@ -62,7 +62,6 @@ int TrainLoop::Train(int epochs, Dataset *ds, std::vector<session::TrainLoopCall
for (auto cb : cbs) cb->StepEnd(cb_data);
iter->GetNextRow(&row_vec);
}
iter->Stop();
int break_loop = false;
for (auto cb : cbs) {
int ret = cb->EpochEnd(cb_data);
@ -80,7 +79,7 @@ int TrainLoop::Train(int epochs, Dataset *ds, std::vector<session::TrainLoopCall
break;
}
}
iter->Stop();
for (auto cb : cbs) cb->End(cb_data);
return RET_OK;
}