Caffe2 - C++ API
A deep learning, cross platform ML framework
Data Structures | Public Types | Public Member Functions | Data Fields
torch::data::datasets::detail::BatchDataBuffer< UnwrappedBatch, ExampleSampler > Class Template Reference

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< UnwrappedBatchDatabatch_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
 

Detailed Description

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

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.

Definition at line 39 of file chunk.h.

Member Function Documentation

template<typename UnwrappedBatch = std::vector<Example<>>, typename ExampleSampler = samplers::RandomSampler>
void torch::data::datasets::detail::BatchDataBuffer< UnwrappedBatch, ExampleSampler >::add_chunk_data ( UnwrappedBatchType  data)
inline

Push preloaded chunks to batch queue.

Called from the ChunkDataset worker threads.

Definition at line 85 of file chunk.h.

template<typename UnwrappedBatch = std::vector<Example<>>, typename ExampleSampler = samplers::RandomSampler>
void torch::data::datasets::detail::BatchDataBuffer< UnwrappedBatch, ExampleSampler >::add_chunk_data ( std::exception_ptr  e_ptr)
inline

Push exceptions thrown during preloading into batch queue.

Called from the ChunkDataset worker threads.

Definition at line 146 of file chunk.h.

template<typename UnwrappedBatch = std::vector<Example<>>, typename ExampleSampler = samplers::RandomSampler>
BatchType torch::data::datasets::detail::BatchDataBuffer< UnwrappedBatch, ExampleSampler >::get_batch ( )
inline

Return batch data from the queue.

Called from the ChunkDataset main thread.

Definition at line 55 of file chunk.h.

Field Documentation

template<typename UnwrappedBatch = std::vector<Example<>>, typename ExampleSampler = samplers::RandomSampler>
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.

Definition at line 191 of file chunk.h.


The documentation for this class was generated from the following file: