data_loader.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <memory>
#include <optional>
 
#include "pulsatrix/collate.hpp"
#include "pulsatrix/dataset.hpp"
#include "pulsatrix/device_backend.hpp"
#include "pulsatrix/iterable_dataset.hpp"
#include "pulsatrix/sampler.hpp"
 
namespace pulsatrix {
 
struct DataLoaderOptions {
    int64_t batch_size = 1;
    bool shuffle = false;
    std::optional<unsigned> shuffle_seed;
    int num_workers = 0;
    int64_t prefetch_batches = 2;
    bool drop_last = false;
    CollateFn collate_fn = DefaultCollate;
};
 
class DataLoader {
public:
    DataLoader(std::shared_ptr<Dataset> dataset, DeviceBackend* backend, DataLoaderOptions options = {});
 
    DataLoader(std::shared_ptr<IterableDataset> dataset, DeviceBackend* backend, DataLoaderOptions options = {});
 
    void reset_epoch();
 
    [[nodiscard]] std::optional<Batch> next_batch();
 
    [[nodiscard]] int64_t num_batches() const;
 
private:
    [[nodiscard]] std::optional<Sample> fetch_next_sample();
 
    std::shared_ptr<Dataset> dataset_;                  // non-null iff constructed from a Dataset
    std::shared_ptr<IterableDataset> iterable_dataset_;  // non-null iff constructed from an IterableDataset
    DeviceBackend* backend_;
    DataLoaderOptions options_;
    std::unique_ptr<Sampler> sampler_;  // only used when dataset_ is set
};
 
}  // namespace pulsatrix