sampler.hpp

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <cstddef>
#include <numeric>
#include <random>
#include <vector>
 
#include <cpp-toolbox/cpp-toolbox_export.hpp>
 
namespace toolbox::io
{
 
struct sequential_policy_t
{
  void operator()(std::vector<std::size_t>& /*indices*/) const noexcept {}
};
 
class shuffle_policy_t
{
public:
  explicit shuffle_policy_t(unsigned seed = std::random_device {}())
      : m_rng(seed)
  {
  }
 
  void operator()(std::vector<std::size_t>& indices)
  {
    std::shuffle(indices.begin(), indices.end(), m_rng);
  }
 
  void set_seed(unsigned seed) { m_rng.seed(seed); }
 
private:
  std::mt19937 m_rng;
};
 
template<typename PolicyT = sequential_policy_t>
class CPP_TOOLBOX_EXPORT sampler_t
{
public:
  using index_type = std::size_t;
 
  sampler_t(const sampler_t& other)
      : m_dataset_size(other.m_dataset_size)
      , m_policy(other.m_policy)
      , m_indices(other.m_indices)
  {
    // 计算迭代器偏移/Calculate iterator offset
    auto offset = other.m_iter - other.m_indices.begin();
    m_iter = m_indices.begin() + offset;
  }
 
  sampler_t& operator=(const sampler_t& other)
  {
    if (this != &other) {
      m_dataset_size = other.m_dataset_size;
      m_policy = other.m_policy;
      m_indices = other.m_indices;
      // 计算迭代器偏移/Calculate iterator offset
      auto offset = other.m_iter - other.m_indices.begin();
      m_iter = m_indices.begin() + offset;
    }
    return *this;
  }
 
  explicit sampler_t(std::size_t dataset_size, PolicyT policy = PolicyT {})
      : m_dataset_size(dataset_size)
      , m_policy(std::move(policy))
  {
    m_indices.resize(m_dataset_size);
    std::iota(m_indices.begin(), m_indices.end(), index_type {0});
    reset();
  }
 
  void reset()
  {
    std::iota(m_indices.begin(), m_indices.end(), index_type {0});
    m_policy(m_indices);
    m_iter = m_indices.begin();
  }
 
  [[nodiscard]] bool has_next() const noexcept
  {
    return m_iter != m_indices.end();
  }
 
  index_type next() { return *m_iter++; }
 
  std::vector<index_type> next_batch(std::size_t batch_size)
  {
    std::vector<index_type> batch;
    batch.reserve(batch_size);
    for (std::size_t i = 0; i < batch_size && has_next(); ++i) {
      batch.push_back(next());
    }
    return batch;
  }
 
private:
  std::size_t m_dataset_size;
  PolicyT m_policy;
  std::vector<index_type> m_indices;
  typename std::vector<index_type>::iterator m_iter;
};
 
}  // namespace toolbox::io