Add new load and save ops for storing model params in a single file

paddle/operators/CMakeLists.txt

@@ -191,3 +191,4 @@ if(WITH_GPU)
     cc_test(nccl_op_test SRCS nccl_op_test.cu.cc DEPS nccl_op gpu_info device_context)
 endif()
 cc_test(save_load_op_test SRCS save_load_op_test.cc DEPS save_op load_op)
+cc_test(save_load_combine_op_test SRCS save_load_combine_op_test.cc DEPS save_combine_op load_combine_op)

paddle/operators/load_combine_op.cc

@@ -0,0 +1,122 @@
+/* Copyright (c) 2018 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. */
+#include <fstream>
+
+#include "paddle/framework/op_registry.h"
+#include "paddle/platform/device_context.h"
+
+namespace paddle {
+namespace operators {
+
+class LoadCombineOp : public framework::OperatorBase {
+ public:
+  LoadCombineOp(const std::string &type,
+                const framework::VariableNameMap &inputs,
+                const framework::VariableNameMap &outputs,
+                const framework::AttributeMap &attrs)
+      : OperatorBase(type, inputs, outputs, attrs) {}
+  void Run(const framework::Scope &scope,
+           const platform::Place &place) const override {
+    auto filename = Attr<std::string>("file_path");
+    auto position_counter = Attr<int>("position_counter");
+
+    std::ifstream fin(filename);
+    PADDLE_ENFORCE(static_cast<bool>(fin),
+                   "Cannot open file %s for load_combine op", filename);
+
+    auto out_var_name = Output("Out");
+    auto *out_var = scope.FindVar(out_var_name);
+    PADDLE_ENFORCE(out_var != nullptr, "Output variable %s cannot be found",
+                   out_var_name);
+
+    auto *tensor = out_var->GetMutable<framework::LoDTensor>();
+
+    platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
+    auto &dev_ctx = *pool.Get(place);
+
+    uint64_t data_length;
+    char *buffer = NULL;
+    for (int i = 0; i <= position_counter; i++) {
+      if (!buffer) delete[] buffer;
+
+      // Error checking
+      PADDLE_ENFORCE(static_cast<bool>(fin), "Cannot read more from file %s",
+                     filename);
+
+      // Read a fixed-width int, to get the number of bytes
+      // for the serialized data.
+      fin.read(reinterpret_cast<char *>(&data_length), sizeof(data_length));
+
+      // Error checking
+      PADDLE_ENFORCE(static_cast<bool>(fin), "Cannot read more from file %s",
+                     filename);
+
+      buffer = new char[data_length];
+
+      // Read the serialized data into the buffer
+      fin.read(buffer, data_length);
+    }
+
+    std::string current_serialized_data;
+    current_serialized_data.assign(buffer, data_length);
+
+    // Create an input string stream
+    std::istringstream ist(current_serialized_data);
+    DeserializeFromStream(ist, tensor, dev_ctx);
+
+    if (!buffer) delete[] buffer;  // delete the last allocated memory
+
+    if (platform::is_gpu_place(place)) {
+      // copy CPU to GPU
+      framework::LoDTensor cpu_tensor;
+      cpu_tensor.ShareDataWith(*tensor);
+      cpu_tensor.set_lod(tensor->lod());
+
+      // reset tensor
+      out_var->Clear();
+      tensor = out_var->GetMutable<framework::LoDTensor>();
+      tensor->set_lod(cpu_tensor.lod());
+      Copy(cpu_tensor, place, dev_ctx, tensor);
+    }
+  }
+};
+
+class LoadCombineOpProtoMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  LoadCombineOpProtoMaker(OpProto *proto, OpAttrChecker *op_checker)
+      : OpProtoAndCheckerMaker(proto, op_checker) {
+    AddOutput("Out", "(Tensor) The tensor need to be load_combineed");
+    AddAttr<int>("position_counter",
+                 "(int) "
+                 "It specifies the relative ordering of different parameters.")
+        .AddCustomChecker([](const int &counter) { return counter >= 0; });
+    AddAttr<std::string>("file_path",
+                         "(string) "
+                         "Variable will be load_combined from \"file_path\".")
+        .AddCustomChecker(
+            [](const std::string &path) { return !path.empty(); });
+    AddComment(R"DOC(
+LoadCombine Operator.
+
+LoadCombine operator combines together various tensor variable into a file.
+
+)DOC");
+  }
+};
+}  // namespace operators
+}  // namespace paddle
+namespace ops = paddle::operators;
+
+REGISTER_OPERATOR(load_combine, ops::LoadCombineOp,
+                  ops::LoadCombineOpProtoMaker);

paddle/operators/save_combine_op.cc

@@ -0,0 +1,159 @@
+/* Copyright (c) 2018 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. */
+
+#include <stdint.h>
+#include <sys/stat.h>
+#include <fstream>
+#include <numeric>
+#include <sstream>
+#include "paddle/framework/data_type.h"
+#include "paddle/framework/framework.pb.h"
+#include "paddle/framework/lod_tensor.h"
+#include "paddle/framework/op_registry.h"
+#include "paddle/platform/device_context.h"
+
+namespace paddle {
+namespace operators {
+
+// TODO(sidgoyal78): These function are needed by other files (save_op), move
+// them to paddle::filesystem namespace. (as noted by yuyang18 in save_op).
+constexpr char kSEP = '/';
+static bool FileExists(const std::string &filepath) {
+  struct stat buffer;
+  return (stat(filepath.c_str(), &buffer) == 0);
+}
+
+static std::string DirName(const std::string &filepath) {
+  auto pos = filepath.rfind(kSEP);
+  if (pos == std::string::npos) {
+    return "";
+  }
+  return filepath.substr(0, pos);
+}
+
+static void MkDir(const char *path) {
+  if (mkdir(path, 0755)) {
+    PADDLE_ENFORCE_EQ(errno, EEXIST, "%s mkdir failed!", path);
+  }
+}
+
+static void MkDirRecursively(const char *fullpath) {
+  if (*fullpath == '\0') return;  // empty string
+  if (FileExists(fullpath)) return;
+
+  MkDirRecursively(DirName(fullpath).c_str());
+  MkDir(fullpath);
+}
+
+class SaveCombineOp : public framework::OperatorBase {
+ public:
+  SaveCombineOp(const std::string &type,
+                const framework::VariableNameMap &inputs,
+                const framework::VariableNameMap &outputs,
+                const framework::AttributeMap &attrs)
+      : OperatorBase(type, inputs, outputs, attrs) {}
+  void Run(const framework::Scope &scope,
+           const platform::Place &place) const override {
+    auto filename = Attr<std::string>("file_path");
+    auto overwrite = Attr<bool>("overwrite");
+    auto position_counter = Attr<int>("position_counter");
+
+    bool is_present = FileExists(filename);
+    if (is_present && !overwrite && position_counter == 0) {
+      PADDLE_THROW(
+          "%s is existed, cannot save_combine to it when overwrite=false",
+          filename, overwrite);
+    }
+
+    MkDirRecursively(DirName(filename).c_str());
+
+    std::ofstream fout;
+
+    // if position_counter is 0, we open the file in write mode,
+    // otherwise, we open in append mode.
+    if (position_counter == 0) {
+      fout.open(filename);
+    } else {
+      fout.open(filename, std::ios_base::app);
+    }
+    PADDLE_ENFORCE(static_cast<bool>(fout), "Cannot open %s to write",
+                   filename);
+
+    auto iname = Input("X");
+
+    auto *var = scope.FindVar(iname);
+    PADDLE_ENFORCE(var != nullptr,
+                   "Cannot find variable %s for save_combine_op", iname);
+
+    PADDLE_ENFORCE(var->IsType<framework::LoDTensor>(),
+                   "SaveCombineOp only support LoDTensor, %s has wrong type",
+                   iname);
+
+    auto &tensor = var->Get<framework::LoDTensor>();
+
+    // get device context from pool
+    platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
+    auto &dev_ctx = *pool.Get(place);
+
+    // Create "output string stream" to get the serialized LodTensor
+    std::ostringstream str_stream;
+    framework::SerializeToStream(str_stream, tensor, dev_ctx);
+    std::string current_serialized_data = str_stream.str();
+
+    // Save 'current_size' information as a fixed width integer, and
+    // further save the serialized data using 'current_size' bytes
+    uint64_t current_size = current_serialized_data.size();
+    fout.write(reinterpret_cast<const char *>(&current_size),
+               sizeof(current_size));
+    fout.write(reinterpret_cast<const char *>(current_serialized_data.c_str()),
+               static_cast<std::streamsize>(current_size));
+    fout.close();
+  }
+};
+
+class SaveCombineOpProtoMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  SaveCombineOpProtoMaker(OpProto *proto, OpAttrChecker *op_checker)
+      : OpProtoAndCheckerMaker(proto, op_checker) {
+    AddInput("X", "(Tensor ) Input tensor to be save_combined");
+    AddComment(R"DOC(
+Save_combine operator
+
+This operator will serialize and write tensor variables to file on disk in a 
+combined fashion.
+)DOC");
+    AddAttr<bool>("overwrite",
+                  "(boolean, default true)"
+                  "Overwrite the output file if exist")
+        .SetDefault(true);
+    AddAttr<int>("position_counter",
+                 "(int) "
+                 "It specifies the relative ordering of different parameters.")
+        .AddCustomChecker([](const int &counter) { return counter >= 0; });
+    AddAttr<std::string>(
+        "file_path",
+        "(string)"
+        "The \"file_path\" where the variable will be save_combined.")
+        .AddCustomChecker(
+            [](const std::string &path) { return !path.empty(); });
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+
+REGISTER_OPERATOR(save_combine, ops::SaveCombineOp,
+                  ops::SaveCombineOpProtoMaker);