#pragma once
|
|
#include <torch/csrc/utils/variadic.h>
|
#include <torch/types.h>
|
|
#include <c10/util/Exception.h>
|
|
#include <functional>
|
#include <iterator>
|
#include <memory>
|
#include <type_traits>
|
#include <utility>
|
|
namespace torch {
|
namespace data {
|
namespace detail {
|
// For increased safety and more separated logic, this implementation of
|
// `Iterator` consists of a `ValidIterator` and a `SentinelIterator`. A
|
// `ValidIterator` yields new batches until the `DataLoader` is exhausted. While
|
// the `DataLoader` is not exhausted, `ValidIterator`s compare equal if they are
|
// the same object. When the `ValidIterator` becomes exhauted, it compares equal
|
// to the `SentinelIterator`, but not before. Half the code here is to implement
|
// double dispatch for the comparison. Got damnit, C++.
|
|
template <typename Batch>
|
struct ValidIterator;
|
|
template <typename Batch>
|
struct SentinelIterator;
|
|
/// Base class for the `ValidIterator` and `SentinelIterator`
|
template <typename Batch>
|
struct IteratorImpl {
|
virtual ~IteratorImpl() = default;
|
virtual void next() = 0;
|
virtual Batch& get() = 0;
|
virtual bool operator==(const IteratorImpl& other) const = 0;
|
virtual bool operator==(const ValidIterator<Batch>& other) const = 0;
|
virtual bool operator==(const SentinelIterator<Batch>& other) const = 0;
|
};
|
|
template <typename Batch>
|
struct ValidIterator : public IteratorImpl<Batch> {
|
using BatchProducer = std::function<optional<Batch>()>;
|
|
explicit ValidIterator(BatchProducer next_batch)
|
: next_batch_(std::move(next_batch)) {}
|
|
/// Fetches the next batch.
|
void next() override {
|
// If we didn't get the very first batch yet, get it now.
|
lazy_initialize();
|
TORCH_CHECK(
|
batch_.has_value(), "Attempted to increment iterator past the end");
|
// Increment to the next batch.
|
batch_ = next_batch_();
|
}
|
|
/// Returns the current batch. The precondition for this operation to not
|
/// throw an exception is that it has been compared to the `SentinelIterator`
|
/// and did not compare equal.
|
Batch& get() override {
|
// If we didn't get the very first batch yet, get it now.
|
lazy_initialize();
|
TORCH_CHECK(
|
batch_.has_value(),
|
"Attempted to dereference iterator that was past the end");
|
return batch_.value();
|
}
|
|
/// Does double dispatch.
|
bool operator==(const IteratorImpl<Batch>& other) const override {
|
return other == *this;
|
}
|
|
/// A `ValidIterator` is equal to the `SentinelIterator` iff. the
|
/// `ValidIterator` has reached the end of the dataloader.
|
bool operator==(const SentinelIterator<Batch>& /* unused */) const override {
|
lazy_initialize();
|
return !batch_;
|
}
|
|
/// Returns true if the memory address of `other` equals that of `this`.
|
bool operator==(const ValidIterator<Batch>& other) const override {
|
return &other == this;
|
}
|
|
/// Gets the very first batch if it has not yet been fetched.
|
void lazy_initialize() const {
|
if (!initialized_) {
|
batch_ = next_batch_();
|
initialized_ = true;
|
}
|
}
|
|
BatchProducer next_batch_;
|
mutable optional<Batch> batch_;
|
mutable bool initialized_ = false;
|
};
|
|
template <typename Batch>
|
struct SentinelIterator : public IteratorImpl<Batch> {
|
void next() override {
|
AT_ERROR(
|
"Incrementing the DataLoader's past-the-end iterator is not allowed");
|
}
|
|
Batch& get() override {
|
AT_ERROR(
|
"Dereferencing the DataLoader's past-the-end iterator is not allowed");
|
}
|
|
/// Does double dispatch.
|
bool operator==(const IteratorImpl<Batch>& other) const override {
|
return other == *this;
|
}
|
|
/// Calls the comparison operator between `ValidIterator` and
|
/// `SentinelIterator`.
|
bool operator==(const ValidIterator<Batch>& other) const override {
|
return other == *this;
|
}
|
|
/// Sentinel iterators always compare equal.
|
bool operator==(const SentinelIterator<Batch>& other) const override {
|
return true;
|
}
|
};
|
} // namespace detail
|
|
template <typename Batch>
|
class Iterator {
|
public:
|
// Type aliases to make the class recognized as a proper iterator.
|
using difference_type = std::ptrdiff_t;
|
using value_type = Batch;
|
using pointer = Batch*;
|
using reference = Batch&;
|
using iterator_category = std::input_iterator_tag;
|
|
explicit Iterator(std::unique_ptr<detail::IteratorImpl<Batch>> impl)
|
: impl_(std::move(impl)) {}
|
|
/// Increments the iterator.
|
/// Only permitted for valid iterators (not past the end).
|
Iterator& operator++() {
|
impl_->next();
|
return *this;
|
}
|
|
/// Returns the current batch.
|
/// Only permitted for valid iterators (not past the end).
|
Batch& operator*() {
|
return impl_->get();
|
}
|
|
/// Returns a pointer to the current batch.
|
/// Only permitted for valid iterators (not past the end).
|
Batch* operator->() {
|
return &impl_->get();
|
}
|
|
/// Compares two iterators for equality.
|
bool operator==(const Iterator& other) const {
|
return *impl_ == *other.impl_;
|
}
|
|
/// Compares two iterators for inequality.
|
bool operator!=(const Iterator& other) const {
|
return !(*this == other);
|
}
|
|
private:
|
/// Points either to a `ValidIterator` or to a `SentinelIterator`.
|
std::shared_ptr<detail::IteratorImpl<Batch>> impl_;
|
};
|
} // namespace data
|
} // namespace torch
|
