FM / data

/usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/api/include/torch/data/

/usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/api/include/torch/data

Name	Size	Mode	Actions
dataloader/	-	0755	rm
datasets/	-	0755	rm
detail/	-	0755	rm
samplers/	-	0755	rm
transforms/	-	0755	rm
dataloader.h	2166	0644	edit dl rm
dataloader_options.h	2207	0644	edit dl rm
datasets.h	289	0644	edit dl rm
example.h	1308	0644	edit dl rm
iterator.h	5284	0644	edit dl rm
samplers.h	318	0644	edit dl rm
transforms.h	222	0644	edit dl rm
worker_exception.h	1140	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/api/include/torch/data/iterator.h (5284B)

#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 exhausted, 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

cmd: