save_combine_op.cc 4.9 KB
Newer Older
1
/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18

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 <fstream>
#include <numeric>
#include <sstream>
Y
Yi Wang 已提交
19
#include "paddle/fluid/framework/data_type.h"
20
#include "paddle/fluid/framework/data_type_transform.h"
Y
Yi Wang 已提交
21 22 23 24
#include "paddle/fluid/framework/framework.pb.h"
#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/platform/device_context.h"
D
dzhwinter 已提交
25
#include "paddle/fluid/platform/port.h"
26 27 28 29 30 31 32 33 34 35 36

namespace paddle {
namespace operators {

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) {}
37 38 39 40

 private:
  void RunImpl(const framework::Scope &scope,
               const platform::Place &place) const override {
41 42
    auto filename = Attr<std::string>("file_path");
    auto overwrite = Attr<bool>("overwrite");
43
    auto save_as_fp16 = Attr<bool>("save_as_fp16");
44 45 46 47 48 49 50 51

    bool is_present = FileExists(filename);
    if (is_present && !overwrite) {
      PADDLE_THROW("%s exists!, cannot save_combine to it when overwrite=false",
                   filename, overwrite);
    }

    MkDirRecursively(DirName(filename).c_str());
52
    std::ofstream fout(filename, std::ios::binary);
53
    PADDLE_ENFORCE(static_cast<bool>(fout), "Cannot open %s to write",
54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74
                   filename);

    auto inp_var_names = Inputs("X");
    PADDLE_ENFORCE_GT(static_cast<int>(inp_var_names.size()), 0,
                      "The number of input variables should be greater than 0");

    // get device context from pool
    platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
    auto &dev_ctx = *pool.Get(place);

    for (size_t i = 0; i < inp_var_names.size(); i++) {
      auto *var = scope.FindVar(inp_var_names[i]);

      PADDLE_ENFORCE(var != nullptr,
                     "Cannot find variable %s for save_combine_op",
                     inp_var_names[i]);
      PADDLE_ENFORCE(var->IsType<framework::LoDTensor>(),
                     "SaveCombineOp only supports LoDTensor, %s has wrong type",
                     inp_var_names[i]);

      auto &tensor = var->Get<framework::LoDTensor>();
75 76 77
      // Serialize tensors one by one

      // Check types to see if a fp16 transformation is required
Y
Yu Yang 已提交
78
      auto in_dtype = tensor.type();
79 80 81 82 83 84 85 86 87 88
      auto out_dtype =
          save_as_fp16 ? framework::proto::VarType::FP16 : in_dtype;

      if (in_dtype != out_dtype) {
        auto in_kernel_type = framework::OpKernelType(in_dtype, place);
        auto out_kernel_type = framework::OpKernelType(out_dtype, place);
        framework::LoDTensor out;
        // copy LoD info to the new tensor
        out.set_lod(tensor.lod());
        framework::TransDataType(in_kernel_type, out_kernel_type, tensor, &out);
89
        framework::SerializeToStream(fout, out, dev_ctx);
90
      } else {
91
        framework::SerializeToStream(fout, tensor, dev_ctx);
92
      }
93
    }
94
    fout.close();
95 96 97 98 99
  }
};

class SaveCombineOpProtoMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
100
  void Make() override {
101 102 103 104 105 106 107 108 109 110 111 112 113 114
    AddInput(
        "X",
        "(vector) Input LoDTensors that need to be saved together in a file.")
        .AsDuplicable();
    AddComment(R"DOC(
SaveCombine operator

This operator will serialize and write a list of input LoDTensor variables 
to a file on disk.
)DOC");
    AddAttr<bool>("overwrite",
                  "(boolean, default true)"
                  "Overwrite the output file if it exists.")
        .SetDefault(true);
115 116 117 118 119 120
    AddAttr<bool>("save_as_fp16",
                  "(boolean, default false)"
                  "If true, the tensor will be converted to float16 data "
                  "type and then saved. Otherwise, the tensor will be "
                  "directly saved without data type conversion.")
        .SetDefault(false);
121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136
    AddAttr<std::string>(
        "file_path",
        "(string)"
        "The \"file_path\" where the LoDTensor variables will be saved.")
        .AddCustomChecker(
            [](const std::string &path) { return !path.empty(); });
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

REGISTER_OPERATOR(save_combine, ops::SaveCombineOp,
                  ops::SaveCombineOpProtoMaker);