Program Listing for File queue.h¶
↰ Return to documentation for file (torch/csrc/api/include/torch/data/detail/queue.h
)
#pragma once
#include <torch/types.h>
#include <c10/util/Exception.h>
#include <chrono>
#include <condition_variable>
#include <cstddef>
#include <mutex>
#include <queue>
namespace torch {
namespace data {
namespace detail {
template <typename T>
class Queue {
public:
void push(T value) {
{
std::lock_guard<std::mutex> lock(mutex_);
queue_.push(std::move(value));
}
cv_.notify_one();
}
T pop(optional<std::chrono::milliseconds> timeout = nullopt) {
std::unique_lock<std::mutex> lock(mutex_);
if (timeout) {
if (!cv_.wait_for(
lock, *timeout, [this] { return !this->queue_.empty(); })) {
// clang-format off
AT_ERROR(
"Timeout in DataLoader queue while waiting for next batch"
" (timeout was ", timeout->count(), " ms)");
// clang-format on
}
} else {
cv_.wait(lock, [this] { return !this->queue_.empty(); });
}
AT_ASSERT(!queue_.empty());
T value = queue_.front();
queue_.pop();
lock.unlock();
return value;
}
size_t clear() {
std::lock_guard<std::mutex> lock(this->mutex_);
const auto size = queue_.size();
while (!queue_.empty()) {
queue_.pop();
}
return size;
}
private:
std::queue<T> queue_;
std::mutex mutex_;
std::condition_variable cv_;
};
} // namespace detail
} // namespace data
} // namespace torch