Implementing Dynamic Data Sampling for BERT Language Model Training with PyTorch DataLoader

I'm currently in the process of building a BERT language model from scratch for educational purposes. While constructing the model itself was a smooth journey, I encountered challenges in creating the data processing pipeline, particularly with an issue that has me stuck.

Overview:

I am working with the IMDB dataset, treating each review as a document. Each document can be segmented into several sentences using punctuation marks (. ! ?). Each data sample consists of a sentence A, a sentence B, and an is_next label indicating whether the two sentences are consecutive. This implies that from each document (review), I can generate multiple training samples.

I am utilizing PyTorch and attempting to leverage the DataLoader for handling multiprocessing and parallelism.

The Problem:

The __getitem__ method in the Dataset class is designed to return a single training sample for each index. However, in my scenario, each index references a document (review), and an undefined number of training samples may be generated for each index.

The Question:

Is there a recommended way to handle such a situation? Alternatively, I am considering the following approach:

For each index, an undefined number of samples are returned to the DataLoader. The DataLoader would then assess whether the number of samples is sufficient to form a batch. Here are the three possible cases:

1. The number of samples returned for an index is less than the batch size. In this case, the DataLoader fetches additional samples from the next index (next document), and any excess is retained to form the next batch.
2. The number of samples returned for an index equals the batch size, and it passes it to the model.

I appreciate any guidance or insights into implementing this dynamic data sampling approach with PyTorch DataLoader.

I've got the answer here.

my_dataset = MyDataset(data)
data_loader = DataLoader(my_dataset, batch_size=None, batch_sampler=None)

batch_size = 32
current_batch = []

# The data_loader may return one or more samples...
for samples in data_loader:
    samples = process_samples(samples)

    # Accumulate samples until reaching the desired batch size
    # ...
    current_batch.append(samples)
    # ...

    if len(current_batch) == batch_size:
        processed_batch = process_batch(current_batch)

        # Forward pass, backward pass, and optimization steps...

        current_batch = []

By disabling automatic batching and manually forming batches in the training loop, I have the flexibility to handle an undefined number of training samples for each index directly within the __getitem__ method. This approach aligns with the dynamic nature of my data, where the number of samples generated for each index may vary.

While it might be more optimized to perform everything in __getitem__, I found this trade-off between simplicity and optimization to be reasonable for my specific dataset and processing requirements.

Feel free to adapt and refine this approach based on your own dataset characteristics and processing needs. If you have any further questions or insights, I'd be happy to discuss them.

Happy coding!