|
|
|
@ -32,10 +32,8 @@ class LoDTensorBlockingQueue {
|
|
|
|
|
friend class LoDTensorBlockingQueueHolder;
|
|
|
|
|
|
|
|
|
|
private:
|
|
|
|
|
LoDTensorBlockingQueue(size_t capacity,
|
|
|
|
|
const std::vector<framework::DDim>& dims,
|
|
|
|
|
bool speed_test_mode = false)
|
|
|
|
|
: queue_(capacity, speed_test_mode), dims_(dims) {}
|
|
|
|
|
explicit LoDTensorBlockingQueue(size_t capacity, bool speed_test_mode = false)
|
|
|
|
|
: queue_(capacity, speed_test_mode) {}
|
|
|
|
|
|
|
|
|
|
public:
|
|
|
|
|
bool Push(const std::vector<framework::LoDTensor>& lod_tensor_vec) {
|
|
|
|
@ -65,17 +63,15 @@ class LoDTensorBlockingQueue {
|
|
|
|
|
|
|
|
|
|
private:
|
|
|
|
|
BlockingQueue<std::vector<framework::LoDTensor>> queue_;
|
|
|
|
|
std::vector<framework::DDim> dims_;
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
class LoDTensorBlockingQueueHolder {
|
|
|
|
|
public:
|
|
|
|
|
void InitOnce(size_t capacity, const std::vector<framework::DDim>& dims,
|
|
|
|
|
bool speed_test_mode = false) {
|
|
|
|
|
void InitOnce(size_t capacity, bool speed_test_mode = false) {
|
|
|
|
|
PADDLE_ENFORCE(
|
|
|
|
|
queue_ == nullptr,
|
|
|
|
|
"LoDTensorBlockingQueueHolder::InitOnce() can only be called once");
|
|
|
|
|
queue_.reset(new LoDTensorBlockingQueue(capacity, dims, speed_test_mode));
|
|
|
|
|
queue_.reset(new LoDTensorBlockingQueue(capacity, speed_test_mode));
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
inline const std::shared_ptr<LoDTensorBlockingQueue>& GetQueue() const {
|
|
|
|
|