Merge pull request #8830 from reyoung/feature/recordio_file_reader

Feature/recordio file reader

Author: reyoung

paddle/fluid/framework/CMakeLists.txt

@@ -21,7 +21,7 @@ endif()
 cc_test(eigen_test SRCS eigen_test.cc DEPS tensor)
 
 nv_test(mixed_vector_test SRCS mixed_vector_test.cu DEPS place paddle_memory device_context init)
-cc_library(lod_tensor SRCS lod_tensor.cc DEPS ddim place tensor framework_proto)
+cc_library(lod_tensor SRCS lod_tensor.cc DEPS ddim place tensor framework_proto recordio)
 cc_test(lod_tensor_test SRCS lod_tensor_test.cc DEPS lod_tensor paddle_memory)
 nv_test(lod_tensor_gpu_test SRCS lod_tensor_test.cu DEPS lod_tensor init)
 

paddle/fluid/framework/lod_tensor.cc

@@ -19,6 +19,9 @@ limitations under the License. */
 #include "paddle/fluid/memory/memcpy.h"
 #include "paddle/fluid/memory/memory.h"
 
+#include "paddle/fluid/recordio/scanner.h"
+#include "paddle/fluid/recordio/writer.h"
+
 #include <stdint.h>
 #include <string.h>
 #include <algorithm>
@@ -291,6 +294,31 @@ void DeserializeFromStream(std::istream &is, LoDTensor *tensor,
   TensorFromStream(is, static_cast<Tensor *>(tensor), dev_ctx);
 }
 
+void WriteToRecordIO(recordio::Writer &writer,
+                     const std::vector<LoDTensor> &tensor,
+                     const platform::DeviceContext &dev_ctx) {
+  std::stringstream buffer;
+  size_t sz = tensor.size();
+  buffer.write(reinterpret_cast<const char *>(&sz), sizeof(uint32_t));
+  for (auto &each : tensor) {
+    SerializeToStream(buffer, each, dev_ctx);
+  }
+  writer.Write(buffer.str());
+}
+
+std::vector<LoDTensor> ReadFromRecordIO(
+    recordio::Scanner &scanner, const platform::DeviceContext &dev_ctx) {
+  std::istringstream sin(scanner.Next());
+  uint32_t sz;
+  sin.read(reinterpret_cast<char *>(&sz), sizeof(uint32_t));
+  std::vector<LoDTensor> result;
+  result.resize(sz);
+  for (uint32_t i = 0; i < sz; ++i) {
+    DeserializeFromStream(sin, &result[i], dev_ctx);
+  }
+  return result;
+}
+
 std::vector<LoDTensor> LoDTensor::SplitLoDTensor(
     const std::vector<platform::Place> places) const {
   check_memory_size();

paddle/fluid/framework/lod_tensor.h

@@ -29,6 +29,12 @@ limitations under the License. */
 #include "paddle/fluid/platform/place.h"
 
 namespace paddle {
+
+namespace recordio {
+class Writer;
+class Scanner;
+}
+
 namespace framework {
 
 /*
@@ -209,5 +215,12 @@ void SerializeToStream(std::ostream& os, const LoDTensor& tensor,
 void DeserializeFromStream(std::istream& is, LoDTensor* tensor,
                            const platform::DeviceContext& dev_ctx);
 
+extern void WriteToRecordIO(recordio::Writer& writer,
+                            const std::vector<LoDTensor>& tensor,
+                            const platform::DeviceContext& dev_ctx);
+
+extern std::vector<LoDTensor> ReadFromRecordIO(
+    recordio::Scanner& scanner, const platform::DeviceContext& dev_ctx);
+
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/framework/lod_tensor_test.cc

@@ -14,6 +14,9 @@
 
 #include "paddle/fluid/framework/lod_tensor.h"
 
+#include "paddle/fluid/recordio/scanner.h"
+#include "paddle/fluid/recordio/writer.h"
+
 #include <glog/logging.h>
 #include <gtest/gtest.h>
 #include <algorithm>
@@ -224,5 +227,43 @@ TEST(LoD, CheckAbsLoD) {
   abs_lod0.push_back(std::vector<size_t>({0}));
   ASSERT_FALSE(CheckAbsLoD(abs_lod0));
 }
+
+TEST(LoDTensor, RecordIO) {
+  LoDTensor tensor;
+  int* tmp = tensor.mutable_data<int>(make_ddim({4, 5}), platform::CPUPlace());
+  for (int i = 0; i < 20; ++i) {
+    tmp[i] = i;
+  }
+
+  std::stringstream* stream = new std::stringstream();
+  auto& ctx =
+      *platform::DeviceContextPool::Instance().Get(platform::CPUPlace());
+  {
+    recordio::Writer writer(stream, recordio::Compressor::kSnappy);
+    WriteToRecordIO(writer, {tensor, tensor}, ctx);
+    WriteToRecordIO(writer, {tensor, tensor}, ctx);
+    writer.Flush();
+  }
+
+  auto assert_tensor_ok = [](const LoDTensor& tensor) {
+    for (int i = 0; i < 20; ++i) {
+      ASSERT_EQ(tensor.data<int>()[i], i);
+    }
+  };
+
+  {
+    std::unique_ptr<std::istream> stream_ptr(stream);
+    recordio::Scanner scanner(std::move(stream_ptr));
+    auto tensors = ReadFromRecordIO(scanner, ctx);
+    ASSERT_EQ(tensors.size(), 2);
+    assert_tensor_ok(tensors[0]);
+    assert_tensor_ok(tensors[1]);
+    tensors = ReadFromRecordIO(scanner, ctx);
+    ASSERT_EQ(tensors.size(), 2);
+    assert_tensor_ok(tensors[0]);
+    assert_tensor_ok(tensors[1]);
+  }
+}
+
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/framework/reader.h

@@ -33,6 +33,8 @@ class ReaderBase {
   std::vector<DDim> shapes() const { return shapes_; }
   void set_shapes(const std::vector<DDim>& shapes) { shapes_ = shapes; }
 
+  virtual bool HasNext() const = 0;
+
   virtual ~ReaderBase() {}
 
  protected:
@@ -53,6 +55,8 @@ class DecoratedReader : public ReaderBase {
 
   void ReInit() override { reader_->ReInit(); }
 
+  bool HasNext() const override { return reader_->HasNext(); }
+
  protected:
   ReaderBase* reader_;
 };
@@ -87,6 +91,8 @@ class ReaderHolder {
     reader_->set_shapes(shapes);
   }
 
+  bool HasNext() const { return reader_->HasNext(); }
+
  private:
   std::unique_ptr<ReaderBase> reader_;
 };

paddle/fluid/operators/detail/safe_ref.h

@@ -14,6 +14,8 @@ limitations under the License. */
 
 #pragma once
 
+#include "paddle/fluid/platform/enforce.h"
+
 namespace paddle {
 namespace operators {
 namespace detail {

paddle/fluid/operators/reader/CMakeLists.txt

@@ -1,6 +1,24 @@
 cc_library(reader_op_registry SRCS reader_op_registry.cc DEPS operator op_registry reader)
-op_library(create_random_data_generator_op SRCS create_random_data_generator_op.cc DEPS reader_op_registry)
-op_library(create_shuffle_reader_op SRCS create_shuffle_reader_op.cc DEPS reader_op_registry)
-op_library(create_batch_reader_op SRCS create_batch_reader_op.cc DEPS reader_op_registry)
-op_library(create_double_buffer_reader_op SRCS create_double_buffer_reader_op.cc DEPS reader_op_registry)
-set(READER_LIBRARY create_random_data_generator_op create_shuffle_reader_op create_batch_reader_op create_double_buffer_reader_op PARENT_SCOPE)
+set(LOCAL_READER_LIBS)
+
+function(reader_library TARGET_NAME)
+    set(oneValueArgs "")
+    set(multiValueArgs SRCS DEPS)
+    set(options "")
+    set(common_deps reader_op_registry)
+    cmake_parse_arguments(reader_library "${options}" "${oneValueArgs}"
+            "${multiValueArgs}" ${ARGN})
+    op_library(${TARGET_NAME} SRCS ${reader_library_SRCS} DEPS ${common_deps} ${reader_library_DEPS})
+    set(LOCAL_READER_LIBS
+            ${TARGET_NAME}
+            ${LOCAL_READER_LIBS}
+        PARENT_SCOPE)
+endfunction()
+
+reader_library(create_random_data_generator_op SRCS create_random_data_generator_op.cc)
+reader_library(create_shuffle_reader_op SRCS create_shuffle_reader_op.cc)
+reader_library(create_batch_reader_op SRCS create_batch_reader_op.cc)
+reader_library(create_recordio_file_reader_op SRCS create_recordio_file_reader_op.cc)
+reader_library(create_double_buffer_reader_op SRCS create_double_buffer_reader_op.cc)
+# Export local libraries to parent
+set(READER_LIBRARY ${LOCAL_READER_LIBS} PARENT_SCOPE)

paddle/fluid/operators/reader/create_double_buffer_reader_op.cc

@@ -37,6 +37,8 @@ class DoubleBufferReader : public framework::DecoratedReader {
 
   ~DoubleBufferReader() { buffer_->Close(); }
 
+  bool HasNext() const override;
+
  private:
   void PrefetchThreadFunc();
 
@@ -106,6 +108,8 @@ void DoubleBufferReader::PrefetchThreadFunc() {
   }
 }
 
+bool DoubleBufferReader::HasNext() const { PADDLE_THROW("Not Implemented"); }
+
 }  // namespace reader
 }  // namespace operators
 }  // namespace paddle

paddle/fluid/operators/reader/create_random_data_generator_op.cc

@@ -52,6 +52,8 @@ class RandomDataGenerator : public framework::FileReader {
 
   void ReInit() override { return; }
 
+  bool HasNext() const override { return true; }
+
  private:
   float min_;
   float max_;

paddle/fluid/operators/reader/create_recordio_file_reader_op.cc

@@ -0,0 +1,87 @@
+//   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/reader_op_registry.h"
+#include "paddle/fluid/recordio/scanner.h"
+
+namespace paddle {
+namespace operators {
+namespace reader {
+class RecordIOFileReader : public framework::FileReader {
+ public:
+  RecordIOFileReader(const std::string& filename,
+                     const std::vector<framework::DDim>& shapes)
+      : FileReader(shapes),
+        scanner_(filename),
+        dev_ctx_(*platform::DeviceContextPool::Instance().Get(
+            platform::CPUPlace())) {}
+
+  void ReadNext(std::vector<framework::LoDTensor>* out) override {
+    *out = framework::ReadFromRecordIO(scanner_, dev_ctx_);
+  }
+
+  bool HasNext() const override { return scanner_.HasNext(); }
+
+  void ReInit() override { scanner_.Reset(); }
+
+ private:
+  recordio::Scanner scanner_;
+  const platform::DeviceContext& dev_ctx_;
+};
+
+class CreateRecordIOReaderOp : public framework::OperatorBase {
+ public:
+  using framework::OperatorBase::OperatorBase;
+
+ private:
+  void RunImpl(const framework::Scope& scope,
+               const platform::Place& dev_place) const override {
+    const auto& shape_concat = Attr<std::vector<int>>("shape_concat");
+    const auto& ranks = Attr<std::vector<int>>("ranks");
+    PADDLE_ENFORCE(!shape_concat.empty() && !ranks.empty());
+    PADDLE_ENFORCE_EQ(std::accumulate(ranks.begin(), ranks.end(), 0),
+                      int(shape_concat.size()),
+                      "The accumulate of all ranks should be equal to the "
+                      "shape concat's length.");
+    std::vector<framework::DDim> shapes = RestoreShapes(shape_concat, ranks);
+    std::string filename = Attr<std::string>("filename");
+
+    auto* out = scope.FindVar(Output("Out"))
+                    ->template GetMutable<framework::ReaderHolder>();
+    out->Reset(new RecordIOFileReader(filename, shapes));
+  }
+};
+
+class CreateRecordIOReaderOpMaker : public FileReaderMakerBase {
+ public:
+  CreateRecordIOReaderOpMaker(OpProto* op_proto, OpAttrChecker* op_checker)
+      : FileReaderMakerBase(op_proto, op_checker) {
+    AddAttr<std::string>("filename", "The filename of record io reader");
+    AddComment(R"DOC(
+      CreateRecordIOReader Operator
+
+      Create a reader from a record io file
+    )DOC");
+  }
+};
+
+}  // namespace reader
+}  // namespace operators
+}  // namespace paddle
+
+namespace reader = paddle::operators::reader;
+
+REGISTER_FILE_READER_OPERATOR(create_recordio_file_reader,
+                              reader::CreateRecordIOReaderOp,
+                              reader::CreateRecordIOReaderOpMaker);