distributed.h

#pragma once

#include <torch/csrc/WindowsTorchApiMacro.h>
#include <torch/data/samplers/base.h>

#include <cstddef>
#include <vector>

namespace torch {
namespace serialize {
class OutputArchive;
class InputArchive;
} // namespace serialize
} // namespace torch

namespace torch {
namespace data {
namespace samplers {

template <typename BatchRequest = std::vector<size_t>>
class DistributedSampler : public Sampler<BatchRequest> {
 public:
  TORCH_API DistributedSampler(
      size_t size,
      size_t num_replicas = 1,
      size_t rank = 0,
      bool allow_duplicates = true)
      : size_(size),
        num_replicas_(num_replicas),
        rank_(rank),
        epoch_(0),
        allow_duplicates_(allow_duplicates) {}

  void set_epoch(size_t epoch) {
    epoch_ = epoch;
  }

  size_t epoch() const {
    return epoch_;
  }

 protected:
  size_t local_sample_count() {
    if (allow_duplicates_) {
      return (size_ + num_replicas_ - 1) / num_replicas_;
    } else {
      return size_ / num_replicas_;
    }
  }

  size_t size_;
  size_t num_replicas_;
  size_t rank_;
  size_t epoch_;
  bool allow_duplicates_;
};

class DistributedRandomSampler : public DistributedSampler<> {
 public:
  TORCH_API DistributedRandomSampler(
      size_t size,
      size_t num_replicas = 1,
      size_t rank = 0,
      bool allow_duplicates = true);

  TORCH_API void reset(optional<size_t> new_size = nullopt) override;

  TORCH_API optional<std::vector<size_t>> next(size_t batch_size) override;

  TORCH_API void save(serialize::OutputArchive& archive) const override;

  TORCH_API void load(serialize::InputArchive& archive) override;

  TORCH_API size_t index() const noexcept;

 private:
  void populate_indices();

  size_t begin_index_;
  size_t end_index_;
  size_t sample_index_;
  std::vector<size_t> all_indices_;
};

class DistributedSequentialSampler : public DistributedSampler<> {
 public:
  TORCH_API DistributedSequentialSampler(
      size_t size,
      size_t num_replicas = 1,
      size_t rank = 0,
      bool allow_duplicates = true);

  TORCH_API void reset(optional<size_t> new_size = nullopt) override;

  TORCH_API optional<std::vector<size_t>> next(size_t batch_size) override;

  TORCH_API void save(serialize::OutputArchive& archive) const override;

  TORCH_API void load(serialize::InputArchive& archive) override;

  TORCH_API size_t index() const noexcept;

 private:
  void populate_indices();

  size_t begin_index_;
  size_t end_index_;
  size_t sample_index_;
  std::vector<size_t> all_indices_;
};

} // namespace samplers
} // namespace data
} // namespace torch