Merge pull request #942 from yeliang2258/external_data_dev

External data support when model is bigger than 2G

README.md

@@ -57,6 +57,7 @@ paddle2onnx --model_dir saved_inference_model \
 |--deploy_backend |**[可选]** 量化模型部署的推理引擎，支持 onnxruntime、tensorrt 或 others，当选择 others 时，所有的量化信息存储于 max_range.txt 文件中，默认为 onnxruntime |
 |--save_calibration_file |**[可选]** TensorRT 8.X版本部署量化模型需要读取的 cache 文件的保存路径，默认为 calibration.cache |
 |--version |**[可选]** 查看 paddle2onnx 版本 |
+|--external_filename |**[可选]** 当导出的ONNX模型大于 2G 时，需要设置 external data 的存储路径，推荐设置为：external_data |
 
 - 使用 onnxruntime 验证转换模型, 请注意安装最新版本（最低要求 1.10.0）
 

paddle2onnx/__init__.py

@@ -42,19 +42,20 @@ def export(model_file,
            enable_optimize=True,
            custom_op_info=None,
            deploy_backend="onnxruntime",
-           calibration_file=""):
+           calibration_file="",
+           external_file=""):
     import paddle2onnx.paddle2onnx_cpp2py_export as c_p2o
     deploy_backend = deploy_backend.lower()
     if custom_op_info is None:
         onnx_model_str = c_p2o.export(
             model_file, params_file, opset_version, auto_upgrade_opset, verbose,
             enable_onnx_checker, enable_experimental_op, enable_optimize, {},
-            deploy_backend, calibration_file)
+            deploy_backend, calibration_file, external_file)
     else:
         onnx_model_str = c_p2o.export(
             model_file, params_file, opset_version, auto_upgrade_opset, verbose,
             enable_onnx_checker, enable_experimental_op, enable_optimize,
-            custom_op_info, deploy_backend, calibration_file)
+            custom_op_info, deploy_backend, calibration_file, external_file)
     if save_file is not None:
         with open(save_file, "wb") as f:
             f.write(onnx_model_str)

paddle2onnx/command.py

@@ -119,6 +119,11 @@ def arg_parser():
         type=ast.literal_eval,
         default=True,
         help="whether enable auto_update_opset, default is True")
+    parser.add_argument(
+        "--external_filename",
+        type=_text_type,
+        default=None,
+        help="The filename of external_data when the model is bigger than 2G.")
     return parser
 
 
@@ -132,12 +137,13 @@ def c_paddle_to_onnx(model_file,
                      enable_experimental_op=True,
                      enable_optimize=True,
                      deploy_backend="onnxruntime",
-                     calibration_file=""):
+                     calibration_file="",
+                     external_file=""):
     import paddle2onnx.paddle2onnx_cpp2py_export as c_p2o
     onnx_model_str = c_p2o.export(
         model_file, params_file, opset_version, auto_upgrade_opset, verbose,
         enable_onnx_checker, enable_experimental_op, enable_optimize, {},
-        deploy_backend, calibration_file)
+        deploy_backend, calibration_file, external_file)
     if save_file is not None:
         with open(save_file, "wb") as f:
             f.write(onnx_model_str)
@@ -214,6 +220,15 @@ def main():
             params_file = ""
         else:
             params_file = os.path.join(args.model_dir, args.params_filename)
+
+        if args.external_filename is None:
+            args.external_filename = "external_data"
+
+        base_path = os.path.dirname(args.save_file)
+        if base_path and not os.path.exists(base_path):
+            os.mkdir(base_path)
+        external_file = os.path.join(base_path, args.external_filename)
+
         calibration_file = args.save_calibration_file
         c_paddle_to_onnx(
             model_file=model_file,
@@ -226,7 +241,8 @@ def main():
             enable_experimental_op=True,
             enable_optimize=True,
             deploy_backend=args.deploy_backend,
-            calibration_file=calibration_file)
+            calibration_file=calibration_file,
+            external_file=external_file)
         logging.info("===============Make PaddlePaddle Better!================")
         logging.info("A little survey: https://iwenjuan.baidu.com/?code=r8hu2s")
         return

paddle2onnx/converter.cc

@@ -135,7 +135,8 @@ PADDLE2ONNX_DECL bool Export(const char* model, const char* params, char** out,
                              bool enable_experimental_op, bool enable_optimize,
                              CustomOp* ops, int op_count,
                              const char* deploy_backend,
-                             char** calibration_cache, int* calibration_size) {
+                             char** calibration_cache, int* calibration_size,
+                             const char* external_file, bool* save_external) {
   auto parser = PaddleParser();
   P2OLogger(verbose) << "Start to parsing Paddle model..." << std::endl;
   if (!parser.Init(model, params)) {
@@ -161,7 +162,7 @@ PADDLE2ONNX_DECL bool Export(const char* model, const char* params, char** out,
   std::string result =
       me.Run(parser, opset_version, auto_upgrade_opset, verbose,
              enable_onnx_checker, enable_experimental_op, enable_optimize,
-             deploy_backend, &calibration_str);
+             deploy_backend, &calibration_str, external_file, save_external);
   if (result.empty()) {
     P2OLogger(verbose) << "The exported ONNX model is invalid!" << std::endl;
     return false;
@@ -183,15 +184,13 @@ PADDLE2ONNX_DECL bool Export(const char* model, const char* params, char** out,
   return true;
 }
 
-PADDLE2ONNX_DECL bool Export(const void* model_buffer, int model_size,
-                             const void* params_buffer, int params_size,
-                             char** out, int* out_size, int32_t opset_version,
-                             bool auto_upgrade_opset, bool verbose,
-                             bool enable_onnx_checker,
-                             bool enable_experimental_op, bool enable_optimize,
-                             CustomOp* ops, int op_count,
-                             const char* deploy_backend,
-                             char** calibration_cache, int* calibration_size) {
+PADDLE2ONNX_DECL bool Export(
+    const void* model_buffer, int model_size, const void* params_buffer,
+    int params_size, char** out, int* out_size, int32_t opset_version,
+    bool auto_upgrade_opset, bool verbose, bool enable_onnx_checker,
+    bool enable_experimental_op, bool enable_optimize, CustomOp* ops,
+    int op_count, const char* deploy_backend, char** calibration_cache,
+    int* calibration_size, const char* external_file, bool* save_external) {
   auto parser = PaddleParser();
   P2OLogger(verbose) << "Start to parsing Paddle model..." << std::endl;
   if (!parser.Init(model_buffer, model_size, params_buffer, params_size)) {
@@ -216,7 +215,7 @@ PADDLE2ONNX_DECL bool Export(const void* model_buffer, int model_size,
   std::string result =
       me.Run(parser, opset_version, auto_upgrade_opset, verbose,
              enable_onnx_checker, enable_experimental_op, enable_optimize,
-             deploy_backend, &calibration_str);
+             deploy_backend, &calibration_str, external_file, save_external);
   if (result.empty()) {
     P2OLogger(verbose) << "The exported ONNX model is invalid!" << std::endl;
     return false;

paddle2onnx/converter.h

@@ -55,7 +55,8 @@ PADDLE2ONNX_DECL bool Export(
     bool enable_experimental_op = false, bool enable_optimize = true,
     CustomOp* ops = nullptr, int op_count = 0,
     const char* deploy_backend = "onnxruntime",
-    char** calibration_cache = nullptr, int* calibration_size = 0);
+    char** calibration_cache = nullptr, int* calibration_size = 0,
+    const char* external_file = "", bool* save_external = nullptr);
 
 PADDLE2ONNX_DECL bool Export(
     const void* model_buffer, int model_size, const void* params_buffer,
@@ -64,7 +65,8 @@ PADDLE2ONNX_DECL bool Export(
     bool enable_onnx_checker = true, bool enable_experimental_op = false,
     bool enable_optimize = true, CustomOp* ops = nullptr, int op_count = 0,
     const char* deploy_backend = "onnxruntime",
-    char** calibration_cache = nullptr, int* calibration_size = 0);
+    char** calibration_cache = nullptr, int* calibration_size = 0,
+    const char* external_file = "", bool* save_external = nullptr);
 
 // Following are inside usage, will remove it maybe
 struct PADDLE2ONNX_DECL ModelTensorInfo {

paddle2onnx/cpp2py_export.cc

@@ -35,7 +35,8 @@ PYBIND11_MODULE(paddle2onnx_cpp2py_export, m) {
                      bool enable_optimize = true,
                      const CustomOpInfo& info = CustomOpInfo(),
                      const std::string& deploy_backend = "onnxruntime",
-                     const std::string& calibration_file = "") {
+                     const std::string& calibration_file = "",
+                     const std::string& external_file = "") {
     P2OLogger(verbose) << "Start to parse PaddlePaddle model..." << std::endl;
     P2OLogger(verbose) << "Model file path: " << model_filename << std::endl;
     P2OLogger(verbose) << "Paramters file path: " << params_filename
@@ -49,7 +50,7 @@ PYBIND11_MODULE(paddle2onnx_cpp2py_export, m) {
                   opset_version, auto_upgrade_opset, verbose,
                   enable_onnx_checker, enable_experimental_op, enable_optimize,
                   nullptr, 0, deploy_backend.c_str(), &calibration_cache,
-                  &cache_size)) {
+                  &cache_size, external_file.c_str())) {
         P2OLogger(verbose) << "Paddle model convert failed." << std::endl;
         return pybind11::bytes("");
       }
@@ -86,7 +87,7 @@ PYBIND11_MODULE(paddle2onnx_cpp2py_export, m) {
                 opset_version, auto_upgrade_opset, verbose, enable_onnx_checker,
                 enable_experimental_op, enable_optimize, ops.data(),
                 info.size(), deploy_backend.c_str(), &calibration_cache,
-                &cache_size)) {
+                &cache_size, external_file.c_str())) {
       P2OLogger(verbose) << "Paddle model convert failed." << std::endl;
       return pybind11::bytes("");
     }

paddle2onnx/mapper/exporter.cc

@@ -14,6 +14,7 @@
 
 #include "paddle2onnx/mapper/exporter.h"
 
+#include <google/protobuf/message.h>
 #include <onnx/checker.h>
 
 #include <array>
@@ -171,13 +172,80 @@ void ModelExporter::ProcessGraphDumplicateNames(
   }
 }
 
-std::string ModelExporter::Run(const PaddleParser& parser, int opset_version,
-                               bool auto_upgrade_opset, bool verbose,
-                               bool enable_onnx_checker,
-                               bool enable_experimental_op,
-                               bool enable_optimize,
-                               const std::string& deploy_backend,
-                               std::string* calibration_cache) {
+void ModelExporter::SaveExternalData(::paddle2onnx::GraphProto* graph,
+                                     const std::string& external_file_path,
+                                     bool* save_external) {
+  P2OLogger() << "The exported ONNX model is bigger than 2G, external data "
+                 "will save to file: "
+              << external_file_path << std::endl;
+  std::string file_name = GetFilenameFromPath(external_file_path);
+  if (save_external) {
+    *save_external = true;
+  }
+  std::fstream f(external_file_path, std::ios::out);
+  Assert(f.is_open(), "Failed to open: " + external_file_path +
+                          " file to save external data");
+  for (auto index = 0; index < graph->node_size(); index++) {
+    auto node = graph->mutable_node(index);
+    if (node->op_type() != "Constant") {
+      continue;
+    }
+    for (auto i = 0; i < node->attribute_size(); i++) {
+      auto attr = node->mutable_attribute(i);
+      if (attr->name() != "value") {
+        continue;
+      }
+      auto tensor = attr->mutable_t();
+
+      if (tensor->raw_data().size() <= 128) {
+        continue;
+      }
+
+      tensor->set_data_location(TensorProto::EXTERNAL);
+      auto external_data = tensor->add_external_data();
+      external_data->set_key("location");
+      external_data->set_value(file_name);
+
+      external_data = tensor->add_external_data();
+      external_data->set_key("offset");
+      f.seekg(0, std::ios::end);
+      int64_t offset = f.tellg();
+      external_data->set_value(std::to_string(offset));
+      auto raw_data = tensor->raw_data();
+      f << raw_data;
+      external_data = tensor->add_external_data();
+      external_data->set_key("length");
+      int64_t raw_datas_size = raw_data.size();
+      external_data->set_value(std::to_string(raw_datas_size));
+      tensor->clear_raw_data();
+    }
+  }
+  f.close();
+}
+void ModelExporter::ONNXChecker(const ONNX_NAMESPACE::ModelProto& model,
+                                const bool& verbose) {
+  // TODO(jiangjiajun)
+  // If we need to integrate with framework
+  // this check will return a information
+  // to let framework know the conversion is
+  // pass or fail
+  try {
+    // ONNX_NAMESPACE::checker::check_model(*(model.get()));
+    ONNX_NAMESPACE::checker::check_model(model);
+  } catch (const std::exception& e) {
+    P2OLogger(verbose) << "The exported ONNX model is invalid." << std::endl;
+    P2OLogger(verbose) << "Model checker error log: " << e.what() << std::endl;
+  }
+  P2OLogger(verbose) << "PaddlePaddle model is exported as ONNX format now."
+                     << std::endl;
+}
+
+std::string ModelExporter::Run(
+    const PaddleParser& parser, int opset_version, bool auto_upgrade_opset,
+    bool verbose, bool enable_onnx_checker, bool enable_experimental_op,
+    bool enable_optimize, const std::string& deploy_backend,
+    std::string* calibration_cache, const std::string& external_file,
+    bool* save_external) {
   _deploy_backend = deploy_backend;
   _helper.SetOpsetVersion(opset_version);
   _total_ops_num = 0;
@@ -274,9 +342,6 @@ std::string ModelExporter::Run(const PaddleParser& parser, int opset_version,
         deploy_backend, parser, calibration_cache);
     // Update int8 weights in quantized OP to float32
     UpdateParameters(_helper.updated_params);
-    // std::ofstream cache_file;
-    // cache_file.open(calibration_file, std::ios::out);
-    // cache_file << calibration_cache;
   }
   // RemoveIsolatedNodes(&parameters, &inputs, &outputs, &_helper.nodes);
 
@@ -296,32 +361,38 @@ std::string ModelExporter::Run(const PaddleParser& parser, int opset_version,
     *(graph->add_value_info()) = (*item.get());
   }
 
-  // TODO(jiangjiajun)
-  // If we need to integrate with framework
-  // this check will return a information
-  // to let framework know the conversion is
-  // pass or fail
-  if (enable_onnx_checker) {
-    try {
-      ONNX_NAMESPACE::checker::check_model(*(model.get()));
-    } catch (...) {
-      P2OLogger(verbose) << "The exported ONNX model is invalid." << std::endl;
-      return "";
+  std::string external_data_file;
+  if (model->ByteSizeLong() > INT_MAX) {
+    if (external_file.empty()) {
+      external_data_file = "external_data";
+    } else {
+      external_data_file = external_file;
     }
-    P2OLogger(verbose) << "PaddlePaddle model is exported as ONNX format now."
-                       << std::endl;
   }
 
   std::string out;
   if (enable_optimize) {
-    auto const opt_model = Optimize(*(model.get()));
+    auto opt_model = Optimize(*(model.get()));
+    if (external_data_file.size()) {
+      SaveExternalData(opt_model.mutable_graph(), external_data_file,
+                       save_external);
+    }
+    if (enable_onnx_checker) {
+      ONNXChecker(opt_model, verbose);
+    }
     if (!opt_model.SerializeToString(&out)) {
       P2OLogger(verbose)
           << "Error happenedd while optimizing the exported ONNX model."
           << std::endl;
       return "";
     }
   } else {
+    if (external_data_file.size()) {
+      SaveExternalData(graph, external_data_file, save_external);
+    }
+    if (enable_onnx_checker) {
+      ONNXChecker(*(model.get()), verbose);
+    }
     if (!model->SerializeToString(&out)) {
       P2OLogger(verbose)
           << "Error happened while optimizing the exported ONNX model."

paddle2onnx/mapper/exporter.h

@@ -22,6 +22,20 @@
 #include "paddle2onnx/mapper/quantize_helper.h"
 #include "paddle2onnx/parser/parser.h"
 
+#ifdef _MSC_VER
+#define PATH_SEP "\\"
+#else
+#define PATH_SEP "/"
+#endif
+
+inline std::string GetFilenameFromPath(const std::string& path) {
+  auto pos = path.find_last_of(PATH_SEP);
+  if (pos == std::string::npos) {
+    return path;
+  }
+  return path.substr(pos + 1);
+}
+
 namespace paddle2onnx {
 
 struct ModelExporter {
@@ -85,13 +99,22 @@ struct ModelExporter {
                           std::set<std::string>* unsupported_ops,
                           bool enable_experimental_op);
 
+  void SaveExternalData(::paddle2onnx::GraphProto* graph,
+                        const std::string& external_file_path,
+                        bool* save_external = nullptr);
+
+  void ONNXChecker(const ONNX_NAMESPACE::ModelProto& model,
+                   const bool& verbose);
+
   std::string Run(const PaddleParser& parser, int opset_version = 9,
                   bool auto_upgrade_opset = true, bool verbose = false,
                   bool enable_onnx_checker = true,
                   bool enable_experimental_op = false,
                   bool enable_optimize = true,
                   const std::string& deploy_backend = "onnxruntime",
-                  std::string* calibration_cache = nullptr);
+                  std::string* calibration_cache = nullptr,
+                  const std::string& external_file = "",
+                  bool* save_external = nullptr);
 };
 
 }  // namespace paddle2onnx

paddle2onnx/mapper/onnx_helper.cc

@@ -14,6 +14,8 @@
 
 #include "paddle2onnx/mapper/onnx_helper.h"
 
+#include <fstream>
+
 namespace paddle2onnx {
 
 void AddAttribute(std::shared_ptr<ONNX_NAMESPACE::NodeProto> node,