BatchDataBuffer manages a queue of UnwrappedBatchData. More...
#include <chunk.h>
Data Structures | |
struct | UnwrappedBatchData |
struct that contains a raw unwrapped batch unit. More... | |
Public Types | |
using | UnwrappedBatchType = UnwrappedBatch |
using | BatchType = torch::optional< UnwrappedBatchType > |
using | BatchRequestType = typename ExampleSampler::BatchRequestType |
Public Member Functions | |
BatchDataBuffer (size_t batch_size, ExampleSampler &example_sampler, size_t queue_capacity) | |
BatchType | get_batch () |
Return batch data from the queue. More... | |
void | add_chunk_data (UnwrappedBatchType data) |
Push preloaded chunks to batch queue. More... | |
void | add_chunk_data (std::exception_ptr e_ptr) |
Push exceptions thrown during preloading into batch queue. More... | |
void | stop () |
Data Fields | |
size_t | batch_size_ = 0 |
The batch size is needed to create batches from the chunk data. More... | |
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 |
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.
|
inline |
Push preloaded chunks to batch queue.
Called from the ChunkDataset worker threads.
|
inline |
Push exceptions thrown during preloading into batch queue.
Called from the ChunkDataset worker threads.
|
inline |
Return batch data from the queue.
Called from the ChunkDataset main thread.
size_t torch::data::datasets::detail::BatchDataBuffer< UnwrappedBatch, ExampleSampler >::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.