Rewrite dataloader, improves responsiveness and reliability

[zhreshold]

Description

Rewrite dataloader with multiprocess.pool, the logic of dataloader with num_worker>0 is cleaner and more robust.

Changes with this PR:

• Secure MXIndexRecordIO reader handler with process id, note that to perseve perf, no lock is introduced, separate processes will obtain new handlers to the record file.
• Replaced manual queue handling with built-in multiprocessing.Pool, no matter how many epochs and iterators are launched, they share same process pool. The benefit is two fold: 1) It's faster, since new iterators will no longer trying to fork new processes 2) it's more reliable, the process pool is managed along with the entire life of dataloader, not iterators.
• Benchmarked with ResNet 50 on ImageNet, p3-x16 instance with 8 GPUs and 60 CPU workers, images/sec bumped from (1500-2400) to (3000-3100), it's faster and more stable
• Added an option prefetch=None:

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

The rest functionality and APIs are intact with this PR. No action is required for exisiting users.

Thanks @szha @leezu for the ideas.
@piiswrong @eric-haibin-lin @hetong007 @YutingZhang for review and questions

This PR may affect gluoncv and gluonnlp dataloader variants, I will address them separately.

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [MXNET-$JIRA_ID], where $JIRA_ID refers to the relevant JIRA issue created (except PRs with tiny changes)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
• Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
• Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
• Code is well-documented:
• For user-facing API changes, API doc string has been updated.
• For new C++ functions in header files, their functionalities and arguments are documented.
• For new examples, README.md is added to explain the what the example does, the source of the dataset, expected performance on test set and reference to the original paper if applicable
• Check the API doc at http://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
• To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• Feature1, tests, (and when applicable, API doc)
• Feature2, tests, (and when applicable, API doc)

Comments

• If this change is a backward incompatible change, why must this change be made.
• Interesting edge cases to note here

[ThomasDelteil]

Thanks that's a massive improvement, especially for datasets that gets iterated in a small number of iterations (small datasets or big batch sizes). The killing and respawning of workers was one reason why I was questioning more and more the fact that Gluon is privileging the concept of epochs vs number of iterations in measuring progress through training for checkpointing, iterating, etc.

The next improvements I would love to see for the data loader would be the cold start issue at beginning of epoch:

• prefetching the first batch of the next epoch before the previous epoch is completed.
• Or when the queue is empty, first batch should be prefetched in a distributed manner across workers to avoid cold start problem on large batch-sizes. But that would become secondary if the first point is addressed.

[vandanavk]

@mxnet-label-bot add [Data-loading, pr-awaiting-review]

@sandeep-krishnamurthy @anirudhacharya @kalyc

[leezu]

Won't this break when using multiple DataLoader with different datasets?

[zhreshold]

No, there is one global dataset per worker process, different data loaders have independent process pools

[leezu]

@ThomasDelteil regarding point 1, this is currently available in GluonNLP with PrefetchingStream. The current API is as follows: Consider a some Stream (generalization of Dataset that does not support random access but only sequential) of batches, we can wrap it in a PrefetchingStream as data = PrefetchingStream(streamofbatches). When calling iter(data) will begin prefetching. It is currently used as follows

    prefetched_iters = []
    for _ in range(min(args.num_prefetch_epoch, args.epochs)):
        prefetched_iters.append(iter(data))
    for epoch in range(args.epochs):
        if epoch + len(prefetched_iters) < args.epochs:
            prefetched_iters.append(iter(data))
        data_iter = prefetched_iters.pop(0)

The usage-pattern may be made more user-friendly, but the basic idea may also be applicable to Dataset API?

[ThomasDelteil]

@leezu, the problem with the stream is that the shuffling if it happens is to be done ahead of time so that the sequential access remains random.

Currently one solution that works and that I prefer, with the current Dataset and DataLoader API, is to use a continuous batch sampler and effectively iterate through your batches rather than using the concept of epoch.
You iterate your dataloader until you break out of the loop after N iterations.

You can design your continuous batch sampler to return an iterator that is effectively creating epochs times sequences of indices without replacement in a single iterator.

[eric-haibin-lin]

Looks good! Nice improvement

[eric-haibin-lin]

remove unused code

[leezu]

@ThomasDelteil I'm not sure if I can follow the argument. You can view the current DataLoader as exposing a stream of batches. The sampling (shuffling) still happens lazily. In that case the data pipeline is composed of both Dataset (with random access for the lazy shuffling) and the notion of Stream over batches that can be prefetched

[leezu]

It seems the new DataLoader preserves the API, so why not remove DataLoaderV1?

[zhreshold]

If some specific implementations were relying on the old queue based methods, I think leaving the older version is still preferable for some user, what do you think?

[leezu]

I don't have a strong opinion about it. If you deem the reliance on queue as reliance on an undocumented implementation detail of V1, then it may be better to remove from Gluon and copy the V1 code over to the specific implementation? It may be reduce the number of APIs to maintain in the future. But I'm also OK with keeping it here if you prefer

[ThomasDelteil]

@leezu ok, I believe I hadn't well understood the PrefetchingStream API you described above. Makes sense now after second read.

[eric-haibin-lin]

I am merging this now so that it can be cherry-picked to v1.4 branch quickly

[eric-haibin-lin]

I've added this PR to 1.4.x branch.

[kohillyang]

Is it possible to implement another Dataloader with threadpool? According to my test, Dataloader with threadpool is faster than Dataloader with multiprocess(about 2~5x), especially when the dataset is small or most of codes are written with C++, for example, rotating the image is written by Opencv library. And it has less bugs, less requirements. Because some people are running there program with docker and maybe their shared memory is not enough. And Dataloader with threadpool start more quickly as its overhead is smaller.

I am wonder if mxnet can give us a choise, using threadpool or multiprocess.pool.

[leezu]

You may just add an option to use multiprocessing.pool.ThreadPool instead of the process pool that is currently always used. Probably the worker_initializer function needs to be changed too due to the global variable. Do you want to contribute a PR?

[anirudhacharya]

Dataloader with threadpool is faster than Dataloader with multiprocess(about 2~5x), especially when the dataset is small or most of codes are written with C++, for example, rotating the image is written by Opencv library

what is the basis for this comment. If you have multiple CPU cores isn't multiprocessing more efficient, has easier memory management and bypasses GIL?

[leezu]

@anirudh2290 multiprocessing is often extremely inefficient. The reason is that any object that is passed between the processes must first be pickled and unpickled. Pickling is very slow. Using multi-threading of course requires extra care to bypass the GIL, but can make sense if the pickling overhead is too much. I assume @kohillyang has such case.

[ThomasDelteil]

Yes I think the problem with multi-threading is to make sure the user understands that to get a real parallel processing boost, it needs to be done in the underlying libraries like opencv to bypass the GIL. I think there's a case for it but it shouldn't be the default and have appropriate warning in the documentation. I believe different pickling protocols can be chosen to trade off compute for memory @leezu ? Maybe we could also expose that to the users

[leezu]

@ThomasDelteil yes, threading should certainly not be the default. Unfortunately Python does not allow choosing the Pickle protocol version in the multiprocessing context. "Multiprocessing uses default protocol and there is no simple (without hacking the stdlib) way to pass the right protocol manually." https://bugs.python.org/issue26507 https://bugs.python.org/issue28053 https://bugs.python.org/issue23403

Thus we are stuck with a high pickling overhead when using multiprocessing for the foreseeable future

[zhreshold]

Regarding threadpool version, I think it's quite easy to swap to. Performance-wise, it's really case-by-case. And @kohillyang 's suggestion is good to have, we can have both version ready without a lot of change.