未验证 提交 c03092b1 编写于 作者: K Kaipeng Deng 提交者: GitHub

fix dataloader performace decrease & unittest hanging. test=develop (#26739)

上级 1ec30cb1
...@@ -608,22 +608,24 @@ class _DataLoaderIterMultiProcess(_DataLoaderIterBase): ...@@ -608,22 +608,24 @@ class _DataLoaderIterMultiProcess(_DataLoaderIterBase):
# in _send_idx but will not increase _rcvd_idx, so we check # in _send_idx but will not increase _rcvd_idx, so we check
# whether the worker is still alive here to skip the discarded # whether the worker is still alive here to skip the discarded
# batch indices and increase _rcvd_idx # batch indices and increase _rcvd_idx
while self._rcvd_idx < self._send_idx: if self._dataset_kind == _DatasetKind.ITER:
info = self._task_infos[self._rcvd_idx] while self._rcvd_idx < self._send_idx:
if len(info) == 2 or self._worker_status[info[0]]: info = self._task_infos[self._rcvd_idx]
break if len(info) == 2 or self._worker_status[info[0]]:
del self._task_infos[self._rcvd_idx] break
self._rcvd_idx += 1 del self._task_infos[self._rcvd_idx]
self._batches_outstanding -= 1 self._rcvd_idx += 1
else: self._batches_outstanding -= 1
# NOTE: _rcvd_idx and _send_idx only record batches among else:
# workers, if batches among workers drained, there # NOTE: _rcvd_idx and _send_idx only record batches among
# may also be data in blocking queue # workers, if batches among workers drained, there
if self._batches_outstanding < len(self._places): # may also be data in blocking queue
return None if self._batches_outstanding < len(self._places):
continue return None
continue
if len(self._task_infos[self._rcvd_idx]) == 2:
if self._rcvd_idx in self._task_infos and \
len(self._task_infos[self._rcvd_idx]) == 2:
return self._task_infos.pop(self._rcvd_idx)[1] return self._task_infos.pop(self._rcvd_idx)[1]
try: try:
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册