dataloader_options.h

#pragma once

#include <torch/arg.h>
#include <torch/types.h>

#include <chrono>
#include <cstddef>

namespace torch {
namespace data {

struct DataLoaderOptions {
  DataLoaderOptions() = default;
  /* implicit */ DataLoaderOptions(size_t batch_size)
      : batch_size_(batch_size) {}

  TORCH_ARG(size_t, batch_size) = 1;

  TORCH_ARG(size_t, workers) = 0;

  TORCH_ARG(optional<size_t>, max_jobs);

  TORCH_ARG(optional<std::chrono::milliseconds>, timeout);

  TORCH_ARG(bool, enforce_ordering) = true;

  TORCH_ARG(bool, drop_last) = false;
};

struct FullDataLoaderOptions {
  explicit FullDataLoaderOptions(DataLoaderOptions options)
      : batch_size(options.batch_size_),
        workers(options.workers_),
        max_jobs(options.max_jobs_.value_or(2 * workers)),
        timeout(options.timeout_),
        enforce_ordering(options.enforce_ordering_),
        drop_last(options.drop_last_) {}

  size_t batch_size;
  size_t workers;
  size_t max_jobs;
  optional<std::chrono::milliseconds> timeout;
  bool enforce_ordering;
  bool drop_last;
};
} // namespace data
} // namespace torch