Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
ChunkDataset API proposal
Problem to be solved
A typical data loading in PyTorch assumes all the data is accessible to every participating process. Randomization is performed by the sampler with the knowledge of total length of the dataset. While this approach is simpler and natural to scenarios such as a directory full of images, it does not map well to situations where a large dataset with unknown size is available in collection of files or a single large file. The global randomization incurs many disk seeks and the user needs to carefully partition data to support distributed training. Manually splitting the data, distribute amongst computing units without duplicates and performing efficient shuffling are not strictly related to training models, but are still important. We often implement similar boiler plate code in different projects, leading to increase in development time.
Proposed solution
The proposed
ChunkDataset
is a stateful dataset that supports hierarchical sampling and efficient reading through chunks. Achunk
, in this context, could be a file, such as audio or image, section of a file in the case of a large text-file, a folder, a URL, or any other abstraction that allows data to be segmented roughly the same size.Unlike regular datasets,
ChunkDataset
implements two levels of sampling, i.e. hierarchical sampling, to operate. In the first level, achunk
is selected based on a sampling strategy and second, a sample is selected from thechunk
using another or similar sampling strategy. The hierarchical sampling approach adopted here provides satisfactory randomness and is inspired by the following paper.By using ChunkDataset API, tasks such as splitting data between computing units with proper randomization become trivial. All user has to do is to provide a
ChunkDataReader
implementation that reads a chunk, instantiate aDistributedChunkSampler
with the desired shuffling strategy and finally putting all together in aChunkDataSet
instance. Once this dataset is passed to PyTorchDataLoader
, every worker will learn its correct rank, reads their pieces of data and continue on the regularDataloader
flow.Brief discussion on API
ChunkDataReader class
In order to perform reading of a particular chunk chosen by
DistributedChunkSampler
, the user has to implement a reader class that extendsChunkDataReader
:DistributedChunkSampler class
DistributedChunkSampler
is already implemented and the user only needs to instantiate it and inject intoChunkDataset
.Similarly to
DistributedSampler
,DistributedChunkSampler
takes :attr:num_replicas, :attr:rank
and :attr:shuffle
on its constructor to specify the number of processes participating in the distributed training, the current rank of a process and the shuffling strategy. One main difference between two samplers is that becauseDistributedChunkSampler
operates onIterableDataset
with unknown size, it takes :attr:num_chunks
as input to draw indices as opposed toDistributedSampler
:attr:dataset
parameter. Another important difference between both samplers is thatDistributedSampler
performs padding on its generated indices, which can't be done for chunks to prevent duplicate reading on different workers.The
DistributedChunkSampler
public API is:ChunkDataset class
ChunkDataset
is already implemented and the user only needs to instantiate it and inject into PyTorchDataLoader
.As mentioned before,
ChunkDataset
is anIterableDataset
implementation, which focus on representing a dataset with unknown size. Once it is passed in to PyTorchDataLoader
, it iterates over the dataset until it is exhausted. At this point, an exception is raised and reading is gracefully finished.ChunkDataset
must bereset
after each epoch to reset the internal state of the sampler and to optionally improve shuffling by injectingepoch
.The
ChunkDataset
public API is:ps: This PR builds on IterableDataset and the original C++ implementation for ChunkDataset API