zhreshold commented on a change in pull request #17841:
URL: https://github.com/apache/incubator-mxnet/pull/17841#discussion_r417660437
##########
File path: python/mxnet/gluon/data/dataloader.py
##########
@@ -607,28 +609,51 @@ def __init__(self, dataset, batch_size=None,
shuffle=False, sampler=None,
self._num_workers = num_workers if num_workers >= 0 else 0
self._worker_pool = None
self._prefetch = max(0, int(prefetch) if prefetch is not None else 2 *
self._num_workers)
- if self._num_workers > 0:
- if self._thread_pool:
- self._worker_pool = ThreadPool(self._num_workers,
-
initializer=_thread_worker_initializer,
- initargs=(is_np_shape(),
is_np_array()))
- else:
- # set ignore keyboard interupt signal before forking processes
- original_sigint_handler = signal.signal(signal.SIGINT,
signal.SIG_IGN)
- self._worker_pool = multiprocessing.Pool(
- self._num_workers, initializer=_worker_initializer,
- initargs=[self._dataset, is_np_shape(), is_np_array()])
- # resume keyboard interupt signal in main process
- signal.signal(signal.SIGINT, original_sigint_handler)
if batchify_fn is None:
if num_workers > 0:
- self._batchify_fn = default_mp_batchify_fn
+ self._batchify_fn = _batchify.Stack(use_shared_mem=True)
else:
- self._batchify_fn = default_batchify_fn
+ self._batchify_fn = _batchify.Stack()
else:
self._batchify_fn = batchify_fn
+ if num_workers > 0 and (try_nopython or try_nopython is None):
+ # check for capability to use mx backend threadedLoader
+ use_mx_iter, mx_iter_args = _check_mx_loader_capability(
+ self._dataset, self._batch_sampler, self._batchify_fn)
+ if not use_mx_iter:
+ if try_nopython:
+ warnings.warn(mx_iter_args)
+ else:
+ use_mx_iter = False
+
+ if use_mx_iter:
+ logging.info("Using MXNet backend ThreadedDataLoader with %s
workers "
+ "instead of python dataloader.", self._num_workers)
+ self._mx_iter = MXThreadedDataLoader(
+ num_workers=self._num_workers,
+ pin_memory=self._pin_memory,
+ pin_device_id=self._pin_device_id,
+ prefetch=self._prefetch, **mx_iter_args)
+ else:
+ if self._num_workers > 0:
Review comment:
num_workers == 0 is handled in `__iter__` and there's nothing needs to
be done in `__init__`
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
