szha commented on a change in pull request #13209: [WIP] Allow dataloader
iterator to be reused
URL: https://github.com/apache/incubator-mxnet/pull/13209#discussion_r232466296
##########
File path: python/mxnet/gluon/data/dataloader.py
##########
@@ -281,6 +335,58 @@ def shutdown(self):
self._shutdown = True
+class _SameProcessIter(object):
+ """Same Process Iterator, which allow reset().
+
+ Parameters
+ ----------
+ dataset : Dataset
+ Source dataset. Note that numpy and mxnet arrays can be directly used
+ as a Dataset.
+ batchify_fn : callable
+ Callback function to allow users to specify how to merge samples
+ into a batch.
+ batch_sampler : Sampler
+ A sampler that returns mini-batches. Do not specify batch_size,
+ shuffle, sampler, and last_batch if batch_sampler is specified.
+ pin_memory : boolean, default False
+ If ``True``, the dataloader will copy NDArrays into pinned memory
+ before returning them. Copying from CPU pinned memory to GPU is faster
+ than from normal CPU memory.
+
+ """
+ def __init__(self, dataset, batchify_fn, batch_sampler, pin_memory=False):
+ self._dataset = dataset
+ self._batchify_fn = batchify_fn
+ self._batch_sampler = batch_sampler
+ self._iter = iter(self._batch_sampler)
Review comment:
call reset here in case more states need to be maintained in reset()
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
