提交 a3f17280 编写于 作者: L lujun

fix dy-load bug, test=develop

上级 f0d48e2a
......@@ -75,7 +75,7 @@ def save_persistables(vardict, dirname, filename=None):
_save_var_to_file(vardict, dirname, filename)
def load_persistables(vardict, dirname, filename=None):
def load_persistables(dirname):
"""
This function trys to load persistable variables from the folder
`dirname` or the file `filename`.
......@@ -86,11 +86,7 @@ def load_persistables(vardict, dirname, filename=None):
the file name.
Args:
vardict(dict of Parameters): The parameters will be loaded.
dirname(str): The directory path.
filename(str|None): The file which saved all variables, this file path should be end with '.npz'. If variables were
saved in differnet files, set it to None.
Default: None
Returns:
dict: The parameter-dict resumed from file
......@@ -104,10 +100,7 @@ def load_persistables(vardict, dirname, filename=None):
param_1 = param_dict['PtbModel_0.w_1']
"""
if isinstance(vardict, collections.OrderedDict):
return _load_var_from_file(vardict, dirname, filename)
return {}
return _load_var_from_file(dirname)
def _save_var_to_file(stat_dict, file_dir, file_name):
......@@ -139,42 +132,38 @@ def _save_var_to_file(stat_dict, file_dir, file_name):
})
def _load_var_from_file(stat_dict, file_dir, file_name):
def _load_var_from_file(file_dir):
def walk_filename(file_dir):
var_name_list = []
if os.path.exists(file_dir) and os.path.exists(os.path.join(file_dir)):
base_path = os.path.join(file_dir)
for dirpath, dirnames, filenames in os.walk(os.path.join(file_dir)):
pt = dirpath.replace(base_path, "", 1)[1:]
for fth_name in filenames:
if fth_name[0] != '.':
var_name_list.append(os.path.join(pt, fth_name))
return var_name_list
load_block = default_main_program().global_block()
load_var_map = {}
for var_key, each_var in stat_dict.items():
assert isinstance(each_var, Variable)
if each_var.type == core.VarDesc.VarType.RAW:
continue
new_var = _clone_var_in_block_(load_block, each_var)
if file_name is None:
file_var_list = walk_filename(file_dir)
for var_name in file_var_list:
new_var = Variable(block=load_block, name=var_name)
load_block.append_op(
type='load',
inputs={},
outputs={'Out': [new_var]},
attrs={
'file_path': os.path.join(file_dir,
os.path.normpath(each_var.name))
os.path.normpath(new_var.name))
})
load_var_map[new_var.name] = new_var
if file_name is not None:
load_var_list = []
for name in sorted(load_var_map.keys()):
load_var_list.append(load_var_map[name])
load_block.append_op(
type='load_combine',
inputs={},
outputs={"Out": load_var_list},
attrs={
'file_path': os.path.join(file_dir, os.path.normpath(file_name))
})
for res_var in load_var_list:
load_var_map[res_var.name] = res_var
return load_var_map
......
......@@ -45,6 +45,7 @@ class Layer(core.Layer):
self._dtype = dtype
self._parameters = collections.OrderedDict()
self._sub_layers = collections.OrderedDict()
self._loaddict_holder = collections.OrderedDict()
self._helper = LayerObjectHelper(self._full_name)
......@@ -193,6 +194,9 @@ class Layer(core.Layer):
"""
assert isinstance(parameter, framework.Parameter)
self._parameters[name] = parameter
if parameter.name in self._loaddict_holder:
self._parameters[name] = self._loaddict_holder[parameter.name]
parameter = self._loaddict_holder[parameter.name]
return parameter
def __getattr__(self, name):
......@@ -207,6 +211,9 @@ class Layer(core.Layer):
if params is None:
raise ValueError(
"super(YourLayer, self).__init__() should be called first")
if value.name in self._loaddict_holder:
params[name] = self._loaddict_holder[value.name]
else:
params[name] = value
elif isinstance(value, core.Layer):
layers = self.__dict__.get('_sub_layers', None)
......@@ -244,6 +251,7 @@ class Layer(core.Layer):
return destination
def load_dict(self, stat_dict, include_sublayers=True):
self._loaddict_holder = stat_dict
for name, item in self.__dict__.get('_parameters', None).items():
if item.name in stat_dict:
var = item._ivar.value()
......
......@@ -142,14 +142,11 @@ class TestDygraphCheckpoint(unittest.TestCase):
for param in mnist.parameters():
dy_param_init_value[param.name] = param.numpy()
mnist.load_dict(
fluid.dygraph.load_persistables(mnist.state_dict(),
"save_dir"))
restore = mnist.parameters()
restore = fluid.dygraph.load_persistables("save_dir")
mnist.load_dict(restore)
self.assertEqual(len(dy_param_init_value), len(restore))
for value in restore:
for ky, value in restore.items():
self.assertTrue(
np.allclose(value.numpy(), dy_param_init_value[
value.name]))
......@@ -158,7 +155,7 @@ class TestDygraphCheckpoint(unittest.TestCase):
step += 1
if step > 20:
if step > 10:
break
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册