sneaxiy
7 years ago
committed by
GitHub
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