Repository: pytorch
Issue: [RFC, Tracker] DataLoader improvements
Labels: module: dataloader,triaged

This issue created to track all current problems of `DataLoader` (and related components such as `DataSet`, `Sampler`, `Transforms`). It is focused on **what** we want to archive. Implementation details will follow after prioritizing.

### Known problems

- Issues interpretability. DataLoader periodically ends in a situation when the process hangs for an unknown to the user reason. Root cause analysis of such cases is overcomplicated due to cryptic Exception traces and/or inability to reproduce quickly. Examples #39570 #36375 #33296 #31758 #30147 #25522. We plan to:

	- [ ]  Implement proper timeouts handling and timeouts control.
	- [ ]  Make DataLoader return readable and actionable exceptions.
	- [ ]  Make DataLoader return usable traces in the case of Ctrl+C and similar OS signals #22924.

 
- Issues with CPU Utilization. Usage of DataLoader frequently ends with oversubscribing to CPU threads or CPU underutilization. We plan to:

	- [ ] Make it easier to manage and understand the number of used threads/processes. 

 
- GPU underutilization. Multiple users reported problems with the CUDA context and multiprocessing. Examples #40403 #35759. We plan to:

	- [ ] Review the process of spawning processes to make sure it is easy/understandable to use. We plan to:
	- [ ] Unlock the ability to use GPU `Transforms`  easily.

 
- Recreation of the processes. Spawning processes is a costly procedure, especially in the case of CUDA usage. We plan to:
          
	- [ ] Reuse processes between training epochs.

- Memory usage issues. In case of unbalanced loading/training speed DataLoader behavior leads to OOM.  Examples #31101 #25691. We plan to:
          
	- [ ] Add ability to control amount of pre-loaded data.


### Documentation

- Clear Initialization process doc. Some users getting confused at what time DataSet / DataLoader getting initialized. The order must be clear and understandable.

- Clear Multiprocessing / Multithreading docs. With recipes (torchaudio is a good example). Especially covering cases when threading meets forking and when cuda meets forking.

### Benchmarking

- To make sure that performance improvements are trackable and we are not introducing slow downs we need to write and document benchmark scripts / methodology. 

- Add documentation how to measure performance of data loading, so users can identify bottlenecks.

- Ability to benchmark CPU vs GPU Transforms.
 
### Improvements

- Controllable size of batch on every step of pipeline. Various features (such as WebData, Batch Loading in general, Minibatces, batch level Transforms) requires various sizes of batch (1-N) on every step of the process. We need to make it possible to specify it on every level of loading graph.
     - WebData #38419. As datasets become larger and larger, storing training samples as individual files becomes impractical and inefficient. This can be addressed using sequential storage formats and sharding.
     - Related issues: #12672 #24985 #26957 #42405

- Improved Sampling. Sampling is important component of data loading pipeline. Next requirements should be considered in case of architecture changes.
     - #25162
     - Cascading Sampling
     - Seq2Seq requirements #25743
     - #28743

- Introduce ability to choose between multi-threading and multi-processing. Right now we provide only multi-processing solution, making multi-threading possible will potentially unlock more performant pipelines.

- Preprocessing/Networking. We currently support Distributed training but it is impossible to do distributed data preprocessing at this moment. We should consider this option.

- Infinite DataSet/DataLoader. There is significant number of requests to provide DataSet/DataLoader which never stops to produce batches. Ex: https://discuss.pytorch.org/t/implementing-an-infinite-loop-dataset-dataloader-combo/35567 https://www.google.com/search?q=pytorch+Infinite+DataSet

- Extend ConcatDataSet functionality. There are more than one way to combine two DataSets. 

- Caching support. In some cases caching is the easiest way to store entire DataSet in the memory. We can unify caching API to avoid multiple non optimal solutions. #35642 #39274

- State saving / restoration for DataSet / DataLoader / Sampler. Restarting training from specific checkpoint is problematic when size of single epoch is too large. We can introduce method to save/restore data pipeline state. #36650.

- Improve collate_fn experience
    - #33181
    - #27617

- Unify Transforms Interface. 
    - Unlock ability to make JIT-able transforms
    - Make it painless to make GPU transforms
    - Make it possible to execute transforms on batch

### Considering Improvements

- C++ implementation of DataLoader components for performance/c++ parity

cc @SsnL
cc @vadimkantorov