Support load state_dict from save_params/persistables (#27298)

* support load state_dict from save_params/persistables

* remove failed unittest

* add load eof check & unittest

* remove eof check

python/paddle/fluid/dygraph/checkpoint.py

@@ -195,13 +195,39 @@ def load_dygraph(model_path, config=None):
     params_file_path = model_prefix + ".pdparams"
     opti_file_path = model_prefix + ".pdopt"
 
-    # deal with argument `configs`
-    configs = config
-    if configs is None:
-        configs = SaveLoadConfig()
+    # deal with argument `config`
+    if config is None:
+        config = SaveLoadConfig()
+
+    if os.path.exists(params_file_path) or os.path.exists(opti_file_path):
+        # Load state dict by `save_dygraph` save format
+        para_dict = {}
+        if os.path.exists(params_file_path):
+            with open(params_file_path, 'rb') as f:
+                para_dict = pickle.load(f) if six.PY2 else pickle.load(
+                    f, encoding='latin1')
+
+        if not config.keep_name_table and "StructuredToParameterName@@" in para_dict:
+            del para_dict["StructuredToParameterName@@"]
+
+        if os.path.exists(opti_file_path):
+            with open(opti_file_path, 'rb') as f:
+                opti_dict = pickle.load(f) if six.PY2 else pickle.load(
+                    f, encoding='latin1')
+    else:
+        # check model path
+        if not os.path.isdir(model_prefix):
+            raise ValueError("Model saved directory '%s' is not exists." %
+                             model_prefix)
+
+        # check whether model file exists
+        if config.model_filename is None:
+            model_filename = '__model__'
+        else:
+            model_filename = config.model_filename
+        model_file_path = os.path.join(model_path, model_filename)
 
-    if not os.path.exists(params_file_path) and not os.path.exists(
-            opti_file_path):
+        if os.path.exists(model_file_path):
             # Load state dict by `jit.save/io.save_inference_model` save format
             # NOTE(chenweihang): [ Compatibility of save_inference_model save format ]
             # The model saved by `save_inference_model` does not completely correspond to 
@@ -210,32 +236,28 @@ def load_dygraph(model_path, config=None):
             # the user to configure the `use_structured_name` argument when `set_state_dict`
             # NOTE(chenweihang): `jit.save` doesn't save optimizer state 
 
-        # 1. check model path
-        if not os.path.isdir(model_prefix):
-            raise ValueError("Model saved directory '%s' is not exists." %
-                             model_prefix)
-
-        # 2. load program desc & construct _ProgramHolder
+            # 1. load program desc & construct _ProgramHolder
             programs = _construct_program_holders(model_path,
-                                              configs.model_filename)
+                                                  config.model_filename)
 
-        # 3. load layer parameters & buffers
+            # 2. load layer parameters & buffers
             # NOTE: using fluid.dygraph.guard() here will cause import error in py2
             with guard():
                 persistable_var_dict = _construct_params_and_buffers(
                     model_prefix,
                     programs,
-                configs.separate_params,
-                configs.params_filename,
+                    config.separate_params,
+                    config.params_filename,
                     append_suffix=False)
 
-            # 4. construct state_dict
+                # 3. construct state_dict
                 para_dict = dict()
                 for var_name in persistable_var_dict:
                     para_dict[var_name] = persistable_var_dict[var_name].numpy()
 
                 # if __variables.info__ exists, we can recover structured_name
-            var_info_path = os.path.join(model_prefix, EXTRA_VAR_INFO_FILENAME)
+                var_info_path = os.path.join(model_prefix,
+                                             EXTRA_VAR_INFO_FILENAME)
                 if os.path.exists(var_info_path):
                     with open(var_info_path, 'rb') as f:
                         extra_var_info = pickle.load(f)
@@ -244,22 +266,44 @@ def load_dygraph(model_path, config=None):
                         structured_name = extra_var_info[var_name].get(
                             'structured_name', None)
                         assert structured_name is not None, "Cannot find saved variable (%s)'s structured name in saved model." % var_name
-                    structured_para_dict[structured_name] = para_dict[var_name]
+                        structured_para_dict[structured_name] = para_dict[
+                            var_name]
                     para_dict = structured_para_dict
         else:
-        # Load state dict by `save_dygraph` save format
-        para_dict = {}
-        if os.path.exists(params_file_path):
-            with open(params_file_path, 'rb') as f:
-                para_dict = pickle.load(f) if six.PY2 else pickle.load(
-                    f, encoding='latin1')
-
-        if not configs.keep_name_table and "StructuredToParameterName@@" in para_dict:
-            del para_dict["StructuredToParameterName@@"]
-
-        if os.path.exists(opti_file_path):
-            with open(opti_file_path, 'rb') as f:
-                opti_dict = pickle.load(f) if six.PY2 else pickle.load(
-                    f, encoding='latin1')
+            # load state dict by `io.save_params/persistables` save format
+            # TODO(chenweihang): [ Now only supports loading parameters seperately ]
+            # If users save all parameters as one file, the [ variable.name -> variable ]
+            # mapping info will lost, so users need to give variable list, but users build 
+            # variable list in dygraph mode is difficult, we recommend users to use
+            # paddle.io.load_program_state in this case
+
+            # Try to load all the files in the directory in VarBase format, 
+            # the file name is used as the name of VarBase
+            load_var_list = []
+
+            # 1. load file names
+            var_name_list = []
+            for root, _, files in os.walk(model_path):
+                for filename in files:
+                    file_path = os.path.join(root, filename)
+                    tmp_var_name = os.path.relpath(file_path, model_path)
+                    var_name = tmp_var_name.replace("\\", "/")
+                    var_name_list.append(var_name)
+
+            # 2. create and load VarBase
+            with guard():
+                for name in var_name_list:
+                    new_var = _varbase_creator(name=name, persistable=True)
+                    _dygraph_tracer().trace_op(
+                        type='load',
+                        inputs={},
+                        outputs={'Out': new_var},
+                        attrs={'file_path': os.path.join(model_path, name)})
+                    load_var_list.append(new_var)
+
+            # 3. construct state_dict
+            para_dict = dict()
+            for var in load_var_list:
+                para_dict[var.name] = var.numpy()
 
     return para_dict, opti_dict

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

@@ -64,7 +64,7 @@ class TestLoadStateDictFromSaveInferenceModel(unittest.TestCase):
         self.batch_size = 128
         self.batch_num = 10
 
-    def train_and_save_model(self):
+    def train_and_save_model(self, only_params=False):
         with new_program_scope():
             startup_program = fluid.default_startup_program()
             main_program = fluid.default_main_program()
@@ -102,6 +102,10 @@ class TestLoadStateDictFromSaveInferenceModel(unittest.TestCase):
                 static_param_dict[param.name] = fluid.executor._fetch_var(
                     param.name)
 
+            if only_params:
+                fluid.io.save_params(
+                    exe, self.save_dirname, filename=self.params_filename)
+            else:
                 fluid.io.save_inference_model(
                     self.save_dirname, ["img"], [prediction],
                     exe,
@@ -120,9 +124,7 @@ class TestLoadStateDictFromSaveInferenceModel(unittest.TestCase):
         self.params_filename = None
         orig_param_dict = self.train_and_save_model()
 
-        configs = paddle.SaveLoadConfig()
-        configs.separate_params = True
-        load_param_dict, _ = paddle.load(self.save_dirname, configs)
+        load_param_dict, _ = paddle.load(self.save_dirname)
         self.check_load_state_dict(orig_param_dict, load_param_dict)
 
     def test_load_with_model_filename(self):
@@ -160,6 +162,14 @@ class TestLoadStateDictFromSaveInferenceModel(unittest.TestCase):
         load_param_dict, _ = paddle.load(self.save_dirname, configs)
         self.check_load_state_dict(orig_param_dict, load_param_dict)
 
+    def test_load_state_dict_from_save_params(self):
+        self.save_dirname = "static_mnist.load_state_dict.save_params"
+        self.params_filename = None
+        orig_param_dict = self.train_and_save_model(True)
+
+        load_param_dict, _ = paddle.load(self.save_dirname)
+        self.check_load_state_dict(orig_param_dict, load_param_dict)
+
 
 if __name__ == '__main__':
     unittest.main()