提交 27e1a9ee 编写于 作者: S sneaxiy

fix hang bug

test=develop
上级 7160cb0f
...@@ -69,29 +69,32 @@ class PyReader(object): ...@@ -69,29 +69,32 @@ class PyReader(object):
class Iterator(object): class Iterator(object):
def __init__(self, reader): def __init__(self, reader):
self._reader = reader self._reader = reader._reader
self._reset = reader._reset
def __iter__(self): def __iter__(self):
return self return self
def next(self): def next(self):
ret = self._reader._reader.read_next() ret = self._reader.read_next()
if len(ret): if len(ret):
return ret return ret
else: else:
self._reader._restart_reader() self._reset()
self._reader._reader.reset()
raise StopIteration raise StopIteration
return Iterator(self) return Iterator(self)
def _restart_reader(self): def _reset(self):
if not self._exited: if not self._exited:
for q in self._queues: for q in self._queues:
q.close() q.close()
if self._thread:
self._thread.join() self._thread.join()
self._reader.reset()
def __thread_main__(): def __thread_main__():
queue_num = len(self._queues) queue_num = len(self._queues)
idx = 0 idx = 0
...@@ -138,4 +141,4 @@ class PyReader(object): ...@@ -138,4 +141,4 @@ class PyReader(object):
assert self._tensor_reader is None, \ assert self._tensor_reader is None, \
"Cannot reset the data source of PyReader" "Cannot reset the data source of PyReader"
self._tensor_reader = reader self._tensor_reader = reader
self._restart_reader() self._reset()
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册