fix process hang when training is over.

This commit is contained in:
liyong 2020-12-20 14:33:04 +08:00
parent 33faba7100
commit 2f783a6e25
3 changed files with 10 additions and 1 deletions

View File

@ -205,9 +205,13 @@ class DatasetHelper:
return self.iter.get_sink_size()
def stop_send(self):
"""Free up resources about data sink."""
"""stop send data about data sink."""
self.iter.stop_send()
def release(self):
"""Free up resources about data sink."""
self.iter.release()
def continue_send(self):
"""continue send data to device at the beginning of epoch."""
self.iter.continue_send()
@ -241,6 +245,7 @@ class _DatasetIter:
_send_data_no_flag(dataset, epoch_num)
self.stop_send = dataset.__transfer_dataset__.stop_send
self.release = dataset.__transfer_dataset__.release
self.continue_send = dataset.__transfer_dataset__.continue_send
self.get_data_info = dataset.__transfer_dataset__.get_data_info
self.dataset_types, self.dataset_shapes = _get_types_and_shapes(dataset)

View File

@ -465,6 +465,7 @@ class Model:
if should_stop:
break
dataset_helper.stop_send()
dataset_helper.release()
list_callback.end(run_context)

View File

@ -64,6 +64,9 @@ class MindData:
def stop_send(self):
pass
def release(self):
pass
def continue_send(self):
pass