Template Class BatchDataBuffer

Nested Relationships

Class Documentation

template<typename UnwrappedBatch, typename ExampleSampler = samplers::RandomSampler>
class torch::data::datasets::detail::BatchDataBuffer

BatchDataBuffer manages a queue of UnwrappedBatchData.

After a new chunk is loaded, BatchDataBuffer splits it into small batches and push them into the queue. When get_batch is called from data loader, it pops cached batches and return. If the cache is empty, it either waits to load more chunks or return null if all chunks are loaded.

Public Types

using UnwrappedBatchType = UnwrappedBatch
using BatchType = torch::optional<UnwrappedBatchType>
using BatchRequestType = typename ExampleSampler::BatchRequestType

Public Functions

BatchDataBuffer(size_t batch_size, ExampleSampler &example_sampler, size_t queue_capacity)
BatchType get_batch()

Return batch data from the queue.

Called from the ChunkDataset main thread.

void add_chunk_data(UnwrappedBatchType data)

Push preloaded chunks to batch queue.

Called from the ChunkDataset worker threads.

void add_chunk_data(std::exception_ptr e_ptr)

Push exceptions thrown during preloading into batch queue.

Called from the ChunkDataset worker threads.

void stop()

Public Members

size_t batch_size_ = 0

The batch size is needed to create batches from the chunk data.

Similar to regular dataloader where the batches are created with prefetches, BatchDataBuffer perform the batch creation using the provided batch size.

size_t total_example_count_in_queue_ = 0

count of total example stored in the queue

std::queue<UnwrappedBatchData> batch_queue_

local cache to store example batches from loaded chunk

std::mutex queue_mutex_
std::condition_variable cv_read_
std::condition_variable cv_write_
ExampleSampler &example_sampler_
size_t queue_capacity_
bool stop_ = false
struct UnwrappedBatchData

struct that contains a raw unwrapped batch unit.

An unwrapped batch unit is the raw data without ‘optional’ wrapper. It can be a collection of images, utterances, e.t.c.

Public Functions

UnwrappedBatchData(UnwrappedBatchType data)
UnwrappedBatchData(std::exception_ptr e)

Public Members

UnwrappedBatchType batch_data

batch data to return

std::exception_ptr exception

exception pointer which captures any abnormal exceptions while creating the batch.


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources