supports loading model from memory, test=develop (#24098)

paddle/fluid/operators/save_combine_op.cc

@@ -71,6 +71,23 @@ to a file on disk.
         "The \"file_path\" where the LoDTensor variables will be saved.")
         .AddCustomChecker(
             [](const std::string& path) { return !path.empty(); });
+    AddAttr<bool>("save_to_memory",
+                  "(boolean, default false)"
+                  "If true, the variables will be saved to binary strings.")
+        .SetDefault(false);
+    AddOutput("Y",
+              "(RAW, default empty)."
+              "This output is used when saving variables to binary strings.")
+        .AsDispensable();
+  }
+};
+
+class SaveCombineOpInferVarType : public framework::VarTypeInference {
+ public:
+  void operator()(framework::InferVarTypeContext* ctx) const override {
+    for (auto& o : ctx->Output("Y")) {
+      ctx->SetType(o, framework::proto::VarType::RAW);
+    }
   }
 };
 
@@ -80,7 +97,7 @@ to a file on disk.
 namespace ops = paddle::operators;
 
 REGISTER_OPERATOR(save_combine, ops::SaveCombineOp,
-                  ops::SaveCombineOpProtoMaker);
+                  ops::SaveCombineOpProtoMaker, ops::SaveCombineOpInferVarType);
 
 REGISTER_OP_CPU_KERNEL(
     save_combine,

paddle/fluid/operators/save_combine_op.h

@@ -38,6 +38,8 @@ class SaveCombineOpKernel : public framework::OpKernel<T> {
     auto filename = ctx.Attr<std::string>("file_path");
     auto overwrite = ctx.Attr<bool>("overwrite");
     auto save_as_fp16 = ctx.Attr<bool>("save_as_fp16");
+    auto save_to_memory = ctx.Attr<bool>("save_to_memory");
+    auto output = ctx.Output<std::string>("Y");
 
     bool is_present = FileExists(filename);
     if (is_present && !overwrite) {
@@ -47,12 +49,7 @@ class SaveCombineOpKernel : public framework::OpKernel<T> {
           filename, overwrite));
     }
 
-    MkDirRecursively(DirName(filename).c_str());
-    std::ofstream fout(filename, std::ios::binary);
-    PADDLE_ENFORCE_EQ(static_cast<bool>(fout), true,
-                      platform::errors::Unavailable(
-                          "Cannot open %s to save variables.", filename));
-
+    std::ostringstream ss;
     auto inp_var_names = ctx.InputNames("X");
     auto &inp_vars = ctx.MultiInputVar("X");
     PADDLE_ENFORCE_GT(inp_var_names.size(), 0UL,
@@ -91,12 +88,25 @@ class SaveCombineOpKernel : public framework::OpKernel<T> {
         // copy LoD info to the new tensor
         out.set_lod(tensor.lod());
         framework::TransDataType(in_kernel_type, out_kernel_type, tensor, &out);
-        framework::SerializeToStream(fout, out, dev_ctx);
+        framework::SerializeToStream(ss, out, dev_ctx);
       } else {
-        framework::SerializeToStream(fout, tensor, dev_ctx);
+        framework::SerializeToStream(ss, tensor, dev_ctx);
       }
     }
-    fout.close();
+    if (save_to_memory) {
+      PADDLE_ENFORCE_NE(output, nullptr,
+                        platform::errors::InvalidArgument(
+                            "Cannot find variable Y for save_combine_op"));
+      *output = ss.str();
+    } else {
+      MkDirRecursively(DirName(filename).c_str());
+      std::ofstream fout(filename, std::ios::binary);
+      PADDLE_ENFORCE_EQ(static_cast<bool>(fout), true,
+                        platform::errors::Unavailable(
+                            "Cannot open %s to save variables.", filename));
+      fout << ss.str();
+      fout.close();
+    }
   }
 };
 

paddle/fluid/pybind/pybind.cc

@@ -957,6 +957,10 @@ All parameter, weight, gradient are variables in Paddle.
              return self.GetMutable<LoDTensor>();
            },
            py::return_value_policy::reference)
+      .def("get_bytes",
+           [](Variable &self) {
+             return py::bytes(*self.GetMutable<std::string>());
+           })
       .def("get_lod_rank_table",
            [](Variable &self) { return self.GetMutable<LoDRankTable>(); },
            py::return_value_policy::reference)

python/paddle/fluid/tests/unittests/test_inference_model_io.py

@@ -16,6 +16,7 @@ from __future__ import print_function
 
 import unittest
 
+import os
 import six
 import numpy as np
 import paddle.fluid.core as core
@@ -27,13 +28,20 @@ import paddle.fluid.layers as layers
 import paddle.fluid.optimizer as optimizer
 from paddle.fluid.compiler import CompiledProgram
 from paddle.fluid.framework import Program, program_guard
-from paddle.fluid.io import save_inference_model, load_inference_model
+from paddle.fluid.io import save_inference_model, load_inference_model, save_persistables
 from paddle.fluid.transpiler import memory_optimize
 
 
 class TestBook(unittest.TestCase):
+    class InferModel(object):
+        def __init__(self, list):
+            self.program = list[0]
+            self.feed_var_names = list[1]
+            self.fetch_vars = list[2]
+
     def test_fit_line_inference_model(self):
         MODEL_DIR = "./tmp/inference_model"
+        UNI_MODEL_DIR = "./tmp/inference_model1"
 
         init_program = Program()
         program = Program()
@@ -65,30 +73,43 @@ class TestBook(unittest.TestCase):
                           'y': tensor_y},
                     fetch_list=[avg_cost])
 
+        # Separated model and unified model
         save_inference_model(MODEL_DIR, ["x", "y"], [avg_cost], exe, program)
+        save_inference_model(UNI_MODEL_DIR, ["x", "y"], [avg_cost], exe,
+                             program, 'model', 'params')
+        main_program = program.clone()._prune_with_input(
+            feeded_var_names=["x", "y"], targets=[avg_cost])
+        params_str = save_persistables(exe, None, main_program, None)
+
         expected = exe.run(program,
                            feed={'x': tensor_x,
                                  'y': tensor_y},
                            fetch_list=[avg_cost])[0]
 
         six.moves.reload_module(executor)  # reload to build a new scope
-        exe = executor.Executor(place)
 
-        [infer_prog, feed_var_names, fetch_vars] = load_inference_model(
-            MODEL_DIR, exe)
-
-        outs = exe.run(
-            infer_prog,
-            feed={feed_var_names[0]: tensor_x,
-                  feed_var_names[1]: tensor_y},
-            fetch_list=fetch_vars)
-        actual = outs[0]
-
-        self.assertEqual(feed_var_names, ["x", "y"])
-        self.assertEqual(len(fetch_vars), 1)
-        print("fetch %s" % str(fetch_vars[0]))
-        self.assertTrue("scale" in str(fetch_vars[0]))
-        self.assertEqual(expected, actual)
+        model_0 = self.InferModel(load_inference_model(MODEL_DIR, exe))
+        with open(os.path.join(UNI_MODEL_DIR, 'model'), "rb") as f:
+            model_str = f.read()
+        model_1 = self.InferModel(
+            load_inference_model(None, exe, model_str, params_str))
+
+        for model in [model_0, model_1]:
+            outs = exe.run(model.program,
+                           feed={
+                               model.feed_var_names[0]: tensor_x,
+                               model.feed_var_names[1]: tensor_y
+                           },
+                           fetch_list=model.fetch_vars)
+            actual = outs[0]
+
+            self.assertEqual(model.feed_var_names, ["x", "y"])
+            self.assertEqual(len(model.fetch_vars), 1)
+            print("fetch %s" % str(model.fetch_vars[0]))
+            self.assertEqual(expected, actual)
+
+        self.assertRaises(ValueError, fluid.io.load_inference_model, None, exe,
+                          model_str, None)
 
 
 class TestSaveInferenceModel(unittest.TestCase):