"Serialize LoDTensor, Save/Restore model" (#4602)

* "add model format design doc"

* "add restore function"

* "add parse protobuf"

* "move necessary information to saver.proto"

* "format code"

* "add gpu option"

* "add lod info"

* "add saveop python test wrapper"

* "checkpoint reuse save operator"

* "rewrite model format design doc"

* "async support needed"

* "fix run once"

* "fix doc based on comments"

* "refine based on comments"

* "fix based comments"

* "remove persistable flag from framework.proto"

* "add IndicateDataType to restore op"

* "add save test"

* "modify save restore code"

* "modified the restore logic"

* rm checkpoint_op.cc

* rm test_checkpoint_op.py

* "get inputs outputs name from execution context"

* Saving each variable to a independent file

* Fix bugs

* Rewrite save_restore_op_test with new Python framework

* Move `SaveOp` and `RestoreOp` from OpWithKernel to OpBase

* Refine unit test of SaveOp and RestoreOp

* fix compile errorwq

Author: dzhwinter

doc/design/model_format.md

@@ -0,0 +1,36 @@
+# Design Doc: Model Format
+
+## Motivation
+
+The model is the output of training process. One complete model consists of two parts, namely, the **topology** and the **parameters**. To support industrial deployment, we need to make the model format must be self-completed and do not expose any training source code.
+
+As a result, In PaddlePaddle, the **topology** represents as a  [ProgramDesc](https://github.com/PaddlePaddle/Paddle/blob/1c0a4c901c9fc881d120249c703b15d1c50dae7d/doc/design/program.md), which describes the model structure. The **parameters** contain all the trainable weights in the model, we must support large size parameter, and efficient serialization/deserialization. 
+
+## Implementation
+
+The topology is saved as a plain text, in detail, a self-contain protobuf file. 
+
+The parameters are saved as a binary file. As we all know, the protobuf message has the limits of [64M size](https://developers.google.com/protocol-buffers/docs/reference/cpp/google.protobuf.io.coded_stream#CodedInputStream.SetTotalBytesLimit.details). We do a (benchmark experiment)[https://github.com/PaddlePaddle/Paddle/pull/4610], its result shows protobuf is not fit in this scene.
+
+As a result, we design a particular format for tensor serialization. By default, arbitrary tensor in Paddle is a [LoDTensor](https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/framework/lod_tensor.md), and has a description information proto of (LoDTensorDesc)[https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/framework/framework.proto#L99]. We save the DescProto as the byte string header, it contains the necessary information, such as the `dims`, the `name` of the tensor, and the `LoD` information in [LoDTensor](https://github.com/PaddlePaddle/Paddle/blob/1c0a4c901c9fc881d120249c703b15d1c50dae7d/paddle/framework/lod_tensor.md). Tensor stores value in a continuous memory buffer, for speed we dump the raw memory to disk and save it as the byte string content. So, the binary format of one tensor is, 
+
+|HeaderLength|ContentLength|**LoDTensorDesc**|**TensorValue**|
+
+In detail, tensor's  byte view as the table shows. Note that all the signed value written in little-endian.
+
+```text
+[offset] [type]              [description] 
+0004     4 bytes integer      HeaderLength, the length of LoDTensorDesc
+0008     4 bytes integer      ContentLength, the length of LodTensor Buffer
+0009     1 bytes char         TensorDesc
+00010    1 bytes char         TensorDesc
+...
+00100    1 bytes char         TensorValue
+00101    1 bytes char         TensorValue
+00102    1 bytes char         TensorValue              ..
+...
+```
+
+## Summary
+
+We introduce the model format, the `ProgramDesc` describe the **topology**, and a bunch of particular format binary tensors describes the **parameters**.

paddle/framework/CMakeLists.txt

@@ -1,4 +1,7 @@
 # ddim lib
+proto_library(framework_proto SRCS framework.proto)
+proto_library(saver_proto SRCS framework.proto saver.proto)
+
 cc_library(ddim SRCS ddim.cc DEPS eigen3)
 cc_test(ddim_test SRCS ddim_test.cc DEPS ddim)
 nv_test(dim_test SRCS dim_test.cu DEPS ddim)
@@ -7,16 +10,15 @@ cc_library(tensor SRCS tensor.cc DEPS ddim place paddle_memory device_context)
 cc_test(tensor_test SRCS tensor_test.cc DEPS tensor)
 cc_test(eigen_test SRCS eigen_test.cc DEPS tensor)
 
-cc_library(lod_tensor SRCS lod_tensor.cc DEPS ddim place tensor)
-cc_test(lod_tensor_test SRCS lod_tensor_test.cc DEPS lod_tensor)
+cc_library(lod_tensor SRCS lod_tensor.cc DEPS ddim place tensor saver_proto framework_proto)
+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)
 
 cc_test(variable_test SRCS variable_test.cc)
 
 cc_library(scope SRCS scope.cc)
 cc_test(scope_test SRCS scope_test.cc DEPS scope)
 
-proto_library(framework_proto SRCS framework.proto)
 
 cc_library(attribute SRCS attribute.cc DEPS framework_proto)
 cc_test(program_desc_test SRCS program_desc_test.cc DEPS proto_desc)

paddle/framework/lod_tensor.cc

@@ -13,6 +13,15 @@
    limitations under the License. */
 
 #include "paddle/framework/lod_tensor.h"
+#include "paddle/framework/saver.pb.h"
+
+#include "paddle/memory/memcpy.h"
+#include "paddle/memory/memory.h"
+
+#include <stdint.h>
+#include <string.h>
+#include <algorithm>
+#include <iterator>
 
 #include <glog/logging.h>
 
@@ -112,5 +121,140 @@ void LoDTensor::ShrinkInLevel(size_t level, size_t elem_begin,
   lod_ = new_lod;
 }
 
+std::string LoDTensor::SerializeToString() const {
+  LoDTensorProto desc;
+
+  // set data_type
+  if (this->type() == typeid(int8_t)) desc.set_data_type(DataType::BOOL);
+  if (this->type() == typeid(int16_t)) desc.set_data_type(DataType::INT16);
+  if (this->type() == typeid(int32_t)) desc.set_data_type(DataType::INT32);
+  if (this->type() == typeid(int64_t)) desc.set_data_type(DataType::INT64);
+  // FIXME(dzh): there is no fp16 in standard c++
+
+  if (this->type() == typeid(float))  // NOLINT
+    desc.set_data_type(DataType::FP32);
+  if (this->type() == typeid(double))  // NOLINT
+    desc.set_data_type(DataType::FP64);
+
+  for (int i = 0; i < dims().size(); ++i) {
+    desc.add_dims(dims()[i]);
+  }
+
+  // set lod information
+  desc.set_lod_level(this->NumLevels());
+  for (size_t i = 0; i < this->NumLevels(); ++i) {
+    LoDInfo* lod = desc.add_levels();
+    for (size_t j = 0; j < lod_[i].size(); ++j) {
+      lod->add_level(lod_[i][j]);
+    }
+  }
+
+  desc.set_version(0);
+
+  std::string desc_bytes = desc.SerializeAsString();
+
+  // FIXME(dzh) : implement fix chunk size buffer.
+  size_t DESC_SIZE = desc_bytes.size();
+  size_t DATA_SIZE = holder_->size() - offset_;
+
+  const size_t BUFFER_SIZE = DESC_SIZE + DATA_SIZE + 2 * sizeof(size_t);
+  char* buffer =
+      static_cast<char*>(memory::Alloc(platform::CPUPlace(), BUFFER_SIZE));
+
+  // format: desc_size data_size, desc_bytes, data_bytes.
+  platform::CPUPlace src_place;
+  platform::CPUPlace dst_place;
+
+  memory::Copy(dst_place, buffer, src_place, &BUFFER_SIZE, sizeof(size_t));
+  memory::Copy(dst_place, buffer + sizeof(size_t), src_place, &DESC_SIZE,
+               sizeof(size_t));
+  memory::Copy(dst_place, buffer + sizeof(size_t) * 2, src_place,
+               desc_bytes.c_str(), desc_bytes.size());
+
+  PADDLE_ENFORCE(this->numel() != 0, "Serialize a empty Tensor!");
+
+  platform::Place place = holder_->place();
+  int element_width = holder_->size() / this->numel();
+
+  if (platform::is_cpu_place(place)) {
+    memory::Copy(dst_place, buffer + sizeof(size_t) * 2 + desc_bytes.size(),
+                 boost::get<platform::CPUPlace>(place),
+                 static_cast<char*>(holder_->ptr()) + offset_ / element_width,
+                 DATA_SIZE);
+  }
+#ifdef PADDLE_WITH_GPU
+  if (platform::is_gpu_place(place)) {
+    memory::Copy(dst_place, buffer + sizeof(size_t) * 2 + desc_bytes.size(),
+                 boost::get<platform::GPUPlace>(place),
+                 static_cast<char*>(holder_->ptr()) + offset_ / element_width,
+                 DATA_SIZE);
+  }
+#endif
+
+  std::string ret(buffer, BUFFER_SIZE);
+  memory::Free(platform::CPUPlace(), buffer);
+  return ret;
+}
+
+void LoDTensor::DeserializeFromString(const std::string& s,
+                                      const platform::Place& dst_place) {
+  size_t DESC_SIZE, BUFFER_SIZE;
+  platform::CPUPlace src_place;
+
+  memory::Copy(src_place, &BUFFER_SIZE, src_place, s.c_str(), sizeof(size_t));
+  memory::Copy(src_place, &DESC_SIZE, src_place, s.c_str() + sizeof(size_t),
+               sizeof(size_t));
+
+  const size_t DATA_SIZE = BUFFER_SIZE - DESC_SIZE - sizeof(size_t) * 2;
+
+  // parse LoDTensorDesc
+  LoDTensorProto desc;
+  desc.ParseFromArray(s.c_str() + sizeof(size_t) * 2, DESC_SIZE);
+
+  std::vector<int64_t> dims;
+  std::copy(desc.dims().begin(), desc.dims().end(), std::back_inserter(dims));
+  this->Resize(make_ddim(dims));
+
+  // parse data type
+  void* ptr = nullptr;
+  if (desc.data_type() == DataType::BOOL)
+    ptr = this->mutable_data<bool>(dst_place);
+  if (desc.data_type() == DataType::INT16)
+    ptr = this->mutable_data<int16_t>(dst_place);
+  if (desc.data_type() == DataType::INT32)
+    ptr = this->mutable_data<int32_t>(dst_place);
+  if (desc.data_type() == DataType::INT64)
+    ptr = this->mutable_data<int64_t>(dst_place);
+  // FIXME(dzh): there is no fp16 in standard c++
+
+  if (desc.data_type() == DataType::FP32)
+    ptr = this->mutable_data<float>(dst_place);
+  if (desc.data_type() == DataType::FP64)
+    ptr = this->mutable_data<double>(dst_place);
+
+  LoD lod;
+  std::vector<size_t> levels;
+  for (int i = 0; i < desc.levels().size(); ++i) {
+    auto current_level = desc.levels()[i].level();
+    std::copy(current_level.begin(), current_level.end(),
+              std::back_inserter(levels));
+    lod.emplace_back(levels);
+    levels.clear();
+  }
+
+  this->set_lod(lod);
+
+  if (platform::is_cpu_place(dst_place)) {
+    memory::Copy(boost::get<platform::CPUPlace>(dst_place), ptr, src_place,
+                 s.c_str() + sizeof(size_t) * 2 + DESC_SIZE, DATA_SIZE);
+  }
+#ifdef PADDLE_WITH_GPU
+  if (platform::is_gpu_place(dst_place)) {
+    memory::Copy(boost::get<platform::GPUPlace>(dst_place), ptr, src_place,
+                 s.c_str() + sizeof(size_t) * 2 + DESC_SIZE, DATA_SIZE);
+  }
+#endif
+}
+
 }  // namespace framework
 }  // namespace paddle

paddle/framework/lod_tensor.h

@@ -25,6 +25,7 @@
 #include "paddle/framework/ddim.h"
 #include "paddle/framework/tensor.h"
 #include "paddle/platform/enforce.h"
+#include "paddle/platform/place.h"
 
 namespace paddle {
 namespace framework {
@@ -132,6 +133,27 @@ class LoDTensor : public Tensor {
    */
   void ShrinkInLevel(size_t level, size_t elem_begin, size_t elem_end);
 
+  /**
+   *  @brief Serialize tensor to char bytes.
+   *  Please check model_format.md for the format detail.
+   *  NOTE: GPUTensor will copy data to cpu implicitly.
+   *  @return return string
+   */
+
+  // FIXME(dzh) : Currently, this interface should only be used in
+  // save/restore model and checkpoint. ParameterServer do not use shape
+  // information to do the optimization, as a result, when we serialize
+  // parameter/gradient to string, we should serialize the tensor
+  // to string in the ps trainer instead of LoDTensor.
+  std::string SerializeToString() const;
+
+  /**
+   *  @brief Deserialize char bytes to tensor.
+   *  @return return string
+   */
+  void DeserializeFromString(const std::string& s,
+                             const platform::Place& dst_place);
+
  private:
   LoD lod_;
 };

paddle/framework/lod_tensor_test.cc

@@ -17,10 +17,13 @@
 #include <gtest/gtest.h>
 #include <algorithm>
 #include <memory>
+#include <vector>
 
 namespace paddle {
 namespace framework {
 
+const int kLodTensorSize = 20 * 128;
+
 class LoDTensorTester : public ::testing::Test {
  public:
   virtual void SetUp() override {
@@ -38,7 +41,10 @@ class LoDTensorTester : public ::testing::Test {
 
     lod_tensor_.Resize({20 /*batch size*/, 128 /*dim*/});
     // malloc memory
-    lod_tensor_.mutable_data<float>(place);
+    float* dst_ptr = lod_tensor_.mutable_data<float>(place);
+    for (int i = 0; i < kLodTensorSize; ++i) {
+      dst_ptr[i] = i;
+    }
 
     lod_tensor_.set_lod(lod);
   }
@@ -101,5 +107,21 @@ TEST_F(LoDTensorTester, ShrinkInLevel) {
   ASSERT_EQ(new_lod_tensor.data<float>(), lod_tensor_.data<float>());
 }
 
+TEST_F(LoDTensorTester, SerializeDeserialize) {
+  LoDTensor new_lod_tensor = lod_tensor_;
+  float* src_ptr = lod_tensor_.data<float>();
+  std::string s = lod_tensor_.SerializeToString();
+  LoDTensor dst;
+  dst.DeserializeFromString(s, platform::CPUPlace());
+  float* dst_ptr = dst.data<float>();
+  for (int i = 0; i < kLodTensorSize; ++i) {
+    EXPECT_EQ(dst_ptr[i], src_ptr[i]);
+  }
+
+  ASSERT_EQ(dst.NumElements(0), 2UL);
+  ASSERT_EQ(dst.NumElements(1), 3UL);
+  ASSERT_EQ(dst.NumElements(2), 8UL);
+}
+
 }  // namespace framework
 }  // namespace paddle

paddle/framework/lod_tensor_test.cu

@@ -48,3 +48,30 @@ TEST(LoDTensor, LoDInGPU) {
     CHECK_EQ(lod[0].data()[i], src_lod[0].data()[i] * 2);
   }
 }
+
+TEST(LoDTensor, SerializeDeserialize) {
+  paddle::framework::LoDTensor lod_tensor;
+  paddle::platform::GPUPlace place(0);
+
+  paddle::framework::LoD src_lod;
+  src_lod.push_back(std::vector<size_t>{0, 2, 4, 6, 8, 10, 12, 14});
+
+  lod_tensor.Resize({14, 16});
+  lod_tensor.mutable_data<float>(place);
+
+  lod_tensor.set_lod(src_lod);
+  CHECK_EQ(lod_tensor.lod_element(0, 2).first, 4UL);
+  CHECK_EQ(lod_tensor.lod_element(0, 4).first, 8UL);
+
+  test<<<1, 8>>>(src_lod[0].data(), src_lod[0].size());
+  cudaDeviceSynchronize();
+
+  std::string s = lod_tensor.SerializeToString();
+  paddle::framework::LoDTensor dst;
+  dst.DeserializeFromString(s, place);
+  paddle::framework::LoD dst_lod = dst.lod();
+
+  for (size_t i = 0; i < dst_lod[0].size(); ++i) {
+    CHECK_EQ(src_lod[0].data()[i], dst_lod[0].data()[i] * 2);
+  }
+}

paddle/framework/saver.proto

@@ -0,0 +1,39 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
+
+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. */
+
+syntax = "proto2";
+option optimize_for = LITE_RUNTIME;
+package paddle.framework;
+
+import "framework.proto";
+
+/**
+ * This file contains necessary information for model, checkpoint.
+ * etc.
+ */
+
+message LoDInfo { repeated int64 level = 1; }
+
+/**
+ * Save the LoDTensorDesc information through LoDTensorProto, its data memory
+ * is copyed to c buffer immediately. See model_format.md for details.
+ */
+
+message LoDTensorProto {
+  optional DataType data_type = 1;
+  repeated int64 dims = 2; // [UNK, 640, 480] is saved as [-1, 640, 480]
+  repeated LoDInfo levels = 3;
+  optional int32 lod_level = 4 [ default = 0 ];
+  optional int32 version = 5;
+}

paddle/framework/scope.cc

@@ -65,6 +65,23 @@ void Scope::DropKids() {
   kids_.clear();
 }
 
+std::vector<std::string> Scope::GetAllNames(bool recursive) const {
+  std::vector<std::string> known_vars(vars_.size());
+
+  if (recursive) {
+    for (auto& kid : kids_) {
+      auto kid_vars = kid->GetAllNames();
+      for (auto& p : kid_vars) {
+        known_vars.emplace_back(p);
+      }
+    }
+  }
+  for (auto& p : vars_) {
+    known_vars.emplace_back(p.first);
+  }
+  return known_vars;
+}
+
 void Scope::DeleteScope(Scope* scope) {
   auto it = std::find(this->kids_.begin(), this->kids_.end(), scope);
   PADDLE_ENFORCE(it != this->kids_.end(), "Cannot find %p as kid scope", scope);