mxnet.gluon.data.dataloader

Dataset generator.

Functions

default_batchify_fn(data)

Collate data into batch.

default_mp_batchify_fn(data)

Collate data into batch.

fetcher_loop_v1(data_queue, data_buffer[, ...])

Fetcher loop for fetching data from queue and put in reorder dict.

rebuild_ndarray(pid, fd, shape, dtype)

Rebuild ndarray from pickled shared memory

rebuild_np_ndarray(pid, fd, shape, dtype)

Rebuild ndarray from pickled shared memory

reduce_ndarray(data)

Reduce ndarray to shared memory handle

reduce_np_ndarray(data)

Reduce ndarray to shared memory handle

worker_loop_v1(dataset, key_queue, ...)

Worker loop for multiprocessing DataLoader.

Classes

ConnectionWrapper(conn)

Connection wrapper for multiprocessing that supports sending NDArray via shared memory.

DataLoader(dataset[, batch_size, shuffle, ...])

Loads data from a dataset and returns mini-batches of data.

DataLoaderV1(dataset[, batch_size, shuffle, ...])

Loads data from a dataset and returns mini-batches of data.

Queue(*args, **kwargs)

Wrapper for multiprocessing queue that dumps NDArray with shared memory.

SimpleQueue(*args, **kwargs)

Wrapper for multiprocessing SimpleQueue that dumps NDArray with shared memory.
class mxnet.gluon.data.dataloader.DataLoader(dataset, batch_size=None, shuffle=False, sampler=None, last_batch=None, batch_sampler=None, batchify_fn=None, num_workers=0, pin_memory=False, pin_device_id=0, prefetch=None, thread_pool=False, timeout=120, try_nopython=None)[source]

Bases: object

Loads data from a dataset and returns mini-batches of data.

Parameters:

  • dataset (Dataset) – Source dataset. Note that numpy and mxnet arrays can be directly used as a Dataset.

  • batch_size (int) – Size of mini-batch.

  • shuffle (bool) – Whether to shuffle the samples.

  • sampler (Sampler) – The sampler to use. Either specify sampler or shuffle, not both.

  • last_batch ({'keep', 'discard', 'rollover'}) –

    How to handle the last batch if batch_size does not evenly divide len(dataset).

    keep - A batch with less samples than previous batches is returned. discard - The last batch is discarded if its incomplete. rollover - The remaining samples are rolled over to the next epoch.

  • batch_sampler (Sampler) – A sampler that returns mini-batches. Do not specify batch_size, shuffle, sampler, and last_batch if batch_sampler is specified.

  • batchify_fn (callable) –

    Callback function to allow users to specify how to merge samples into a batch. Defaults to gluon.data.batchify.Stack().

    def default_batchify_fn(data):
    if isinstance(data[0], nd.NDArray):
        return nd.stack(*data)
    elif isinstance(data[0], np.ndarray):
        return np.stack(*data)
    elif isinstance(data[0], tuple):
        data = zip(*data)
        return [default_batchify_fn(i) for i in data]
    else:
        data = np.asarray(data)
        return np.ndarray(data, dtype=data.dtype)

  • num_workers (int, default 0) – The number of multiprocessing workers to use for data preprocessing.

  • 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.

  • pin_device_id (int, default 0) – The device id to use for allocating pinned memory if pin_memory is True

  • prefetch (int, default is num_workers * 2) – The number of prefetching batches only works if num_workers > 0. If prefetch > 0, it allow worker process to prefetch certain batches before acquiring data from iterators. Note that using large prefetching batch will provide smoother bootstrapping performance, but will consume more shared_memory. Using smaller number may forfeit the purpose of using multiple worker processes, try reduce num_workers in this case. By default it defaults to num_workers * 2.

  • thread_pool (bool, default False) – If True, use threading pool instead of multiprocessing pool. Using threadpool can avoid shared memory usage. If DataLoader is more IO bounded or GIL is not a killing problem, threadpool version may achieve better performance than multiprocessing.

  • timeout (int, default is 120) – The timeout in seconds for each worker to fetch a batch data. Only modify this number unless you are experiencing timeout and you know it’s due to slow data loading. Sometimes full shared_memory will cause all workers to hang and causes timeout. In these cases please reduce num_workers or increase system shared_memory size instead.

  • try_nopython (bool or None, default is None) – Try compile python dataloading pipeline into pure MXNet c++ implementation. The benefit is potentially faster iteration, no shared_memory usage, and less processes managed by python. The compilation is not gauranteed to support all use cases, but it will fallback to python in case of failure. You can set try_nopython to False to disable auto-detection of the compilation feature or leave it to None to allow MXNet to determine it automatically. If you request try_nopython to True and the compilation fails, it will raise a RuntimeError with the failure reason.