Program Listing for File dataloader.h

↰ Return to documentation for file (torch/csrc/api/include/torch/data/dataloader.h)

#pragma once

#include <torch/data/dataloader/stateful.h>
#include <torch/data/dataloader/stateless.h>

#include <torch/csrc/utils/variadic.h>

#include <c10/util/Exception.h>

#include <cstddef>
#include <memory>
#include <type_traits>
#include <utility>

namespace torch::data {

template <typename Dataset, typename Sampler>
std::enable_if_t<
    !Dataset::is_stateful,
    std::unique_ptr<StatelessDataLoader<Dataset, Sampler>>>
make_data_loader(Dataset dataset, Sampler sampler, DataLoaderOptions options) {
  return std::make_unique<StatelessDataLoader<Dataset, Sampler>>(
      std::move(dataset), std::move(sampler), options);
}

template <typename Sampler = samplers::RandomSampler, typename Dataset>
std::enable_if_t<
    !Dataset::is_stateful && std::is_constructible_v<Sampler, size_t>,
    std::unique_ptr<StatelessDataLoader<Dataset, Sampler>>>
make_data_loader(
    Dataset dataset,
    DataLoaderOptions options = DataLoaderOptions()) {
  const std::optional<size_t> size = dataset.size();
  TORCH_CHECK(
      size.has_value(),
      "Expected the dataset to be sized in "
      "order to construct the Sampler");
  return make_data_loader(std::move(dataset), Sampler(*size), options);
}

template <typename Dataset, typename = std::enable_if_t<Dataset::is_stateful>>
std::unique_ptr<StatefulDataLoader<Dataset>> make_data_loader(
    Dataset dataset,
    DataLoaderOptions options = DataLoaderOptions()) {
  return std::make_unique<StatefulDataLoader<Dataset>>(
      std::move(dataset), options);
}
} // namespace torch::data