commit
01fbcb0bbb
@ -0,0 +1,84 @@
|
|||||||
|
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include "paddle/fluid/operators/reader/lod_tensor_blocking_queue.h"
|
||||||
|
#include "paddle/fluid/operators/reader/reader_op_registry.h"
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace operators {
|
||||||
|
namespace reader {
|
||||||
|
|
||||||
|
class PyReader : public framework::ReaderBase {
|
||||||
|
public:
|
||||||
|
explicit PyReader(const std::shared_ptr<LoDTensorBlockingQueue>& queue) {
|
||||||
|
PADDLE_ENFORCE(queue != nullptr, "LoDTensorBlockingQueue must not be null");
|
||||||
|
queue_ = queue;
|
||||||
|
}
|
||||||
|
|
||||||
|
void ReadNext(std::vector<framework::LoDTensor>* out) override {
|
||||||
|
bool success;
|
||||||
|
*out = queue_->Pop(&success);
|
||||||
|
if (!success) out->clear();
|
||||||
|
}
|
||||||
|
|
||||||
|
void ReInit() override {}
|
||||||
|
|
||||||
|
private:
|
||||||
|
std::shared_ptr<LoDTensorBlockingQueue> queue_;
|
||||||
|
};
|
||||||
|
|
||||||
|
class CreatePyReaderOp : public framework::OperatorBase {
|
||||||
|
public:
|
||||||
|
using framework::OperatorBase::OperatorBase;
|
||||||
|
|
||||||
|
private:
|
||||||
|
void RunImpl(const framework::Scope& scope,
|
||||||
|
const platform::Place& dev_place) const override {
|
||||||
|
auto* out = scope.FindVar(Output("Out"))
|
||||||
|
->template GetMutable<framework::ReaderHolder>();
|
||||||
|
if (out->Get() != nullptr) return;
|
||||||
|
|
||||||
|
const std::string& queue_name = Input("blocking_queue");
|
||||||
|
auto* queue_holder_var = scope.FindVar(queue_name);
|
||||||
|
PADDLE_ENFORCE(
|
||||||
|
queue_holder_var != nullptr,
|
||||||
|
"No LoDTensorBlockingQueueHolder variable with name %s found",
|
||||||
|
queue_name);
|
||||||
|
auto* queue_holder =
|
||||||
|
queue_holder_var->template GetMutable<LoDTensorBlockingQueueHolder>();
|
||||||
|
|
||||||
|
out->Reset(new PyReader(queue_holder->GetQueue()));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
class CreatePyReaderOpMaker : public FileReaderMakerBase {
|
||||||
|
protected:
|
||||||
|
void Apply() override {
|
||||||
|
AddInput("blocking_queue",
|
||||||
|
"Name of the `LoDTensorBlockingQueueHolder` variable");
|
||||||
|
|
||||||
|
AddComment(R"DOC(
|
||||||
|
Create PyReader to support LoDTensor data feeding in Python side.
|
||||||
|
)DOC");
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
} // namespace reader
|
||||||
|
} // namespace operators
|
||||||
|
} // namespace paddle
|
||||||
|
|
||||||
|
namespace reader = ::paddle::operators::reader;
|
||||||
|
|
||||||
|
REGISTER_FILE_READER_OPERATOR(create_py_reader, reader::CreatePyReaderOp,
|
||||||
|
reader::CreatePyReaderOpMaker);
|
@ -0,0 +1,103 @@
|
|||||||
|
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#pragma once
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "paddle/fluid/framework/ddim.h"
|
||||||
|
#include "paddle/fluid/framework/lod_tensor.h"
|
||||||
|
#include "paddle/fluid/operators/reader/blocking_queue.h"
|
||||||
|
#include "paddle/fluid/platform/place.h"
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace operators {
|
||||||
|
namespace reader {
|
||||||
|
|
||||||
|
class LoDTensorBlockingQueueHolder;
|
||||||
|
|
||||||
|
class LoDTensorBlockingQueue {
|
||||||
|
friend class LoDTensorBlockingQueueHolder;
|
||||||
|
|
||||||
|
private:
|
||||||
|
LoDTensorBlockingQueue(size_t capacity,
|
||||||
|
const std::vector<framework::DDim>& dims)
|
||||||
|
: queue_(capacity), dims_(dims) {}
|
||||||
|
|
||||||
|
public:
|
||||||
|
bool Push(const std::vector<framework::LoDTensor>& lod_tensor_vec) {
|
||||||
|
CheckDims(lod_tensor_vec);
|
||||||
|
return queue_.Send(lod_tensor_vec);
|
||||||
|
}
|
||||||
|
|
||||||
|
bool Push(std::vector<framework::LoDTensor>&& lod_tensor_vec) {
|
||||||
|
CheckDims(lod_tensor_vec);
|
||||||
|
return queue_.Send(std::move(lod_tensor_vec));
|
||||||
|
}
|
||||||
|
|
||||||
|
std::vector<framework::LoDTensor> Pop(bool* ok = nullptr) {
|
||||||
|
std::vector<framework::LoDTensor> lod_tensor_vec;
|
||||||
|
bool success = queue_.Receive(&lod_tensor_vec);
|
||||||
|
if (ok != nullptr) *ok = success;
|
||||||
|
return lod_tensor_vec;
|
||||||
|
}
|
||||||
|
|
||||||
|
inline size_t Cap() const { return queue_.Cap(); }
|
||||||
|
|
||||||
|
inline size_t Size() const { return queue_.Size(); }
|
||||||
|
|
||||||
|
inline void Close() { return queue_.Close(); }
|
||||||
|
|
||||||
|
inline bool IsClosed() const { return queue_.IsClosed(); }
|
||||||
|
|
||||||
|
private:
|
||||||
|
void CheckDims(const std::vector<framework::LoDTensor>& lod_tensor_vec) {
|
||||||
|
PADDLE_ENFORCE(dims_.size() == lod_tensor_vec.size(),
|
||||||
|
"Expect input size is %d but found %s", dims_.size(),
|
||||||
|
lod_tensor_vec.size());
|
||||||
|
for (size_t i = 0; i < dims_.size(); ++i) {
|
||||||
|
const auto& in_dims = framework::slice_ddim(
|
||||||
|
lod_tensor_vec[i].dims(), 1, lod_tensor_vec[i].dims().size());
|
||||||
|
const auto& expect_dims =
|
||||||
|
framework::slice_ddim(dims_[i], 1, dims_[i].size());
|
||||||
|
PADDLE_ENFORCE(in_dims == expect_dims,
|
||||||
|
"Dims of the %d-th input tensor do not match", i);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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) {
|
||||||
|
PADDLE_ENFORCE(
|
||||||
|
queue_ == nullptr,
|
||||||
|
"LoDTensorBlockingQueueHolder::InitOnce() can only be called once");
|
||||||
|
queue_.reset(new LoDTensorBlockingQueue(capacity, dims));
|
||||||
|
}
|
||||||
|
|
||||||
|
inline const std::shared_ptr<LoDTensorBlockingQueue>& GetQueue() const {
|
||||||
|
return queue_;
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
std::shared_ptr<LoDTensorBlockingQueue> queue_;
|
||||||
|
};
|
||||||
|
|
||||||
|
} // namespace reader
|
||||||
|
} // namespace operators
|
||||||
|
} // namespace paddle
|
Loading…
Reference in new issue