diff --git a/python/paddle/fluid/dygraph/checkpoint.py b/python/paddle/fluid/dygraph/checkpoint.py index 27658ba3d46baffa4b64cdf07931c638b9dee086..a51c431c515f23c5e0ec10de71ee8e65b83c492a 100644 --- a/python/paddle/fluid/dygraph/checkpoint.py +++ b/python/paddle/fluid/dygraph/checkpoint.py @@ -18,6 +18,7 @@ import os import collections from ..framework import Variable, default_main_program, in_dygraph_mode, dygraph_only, Parameter, ParamBase import pickle +import six from . import learning_rate_scheduler import warnings from .. import core @@ -88,7 +89,7 @@ def save_dygraph(state_dict, model_path): os.makedirs(dir_name) with open(file_name, 'wb') as f: - pickle.dump(model_dict, f) + pickle.dump(model_dict, f, protocol=2) @dygraph_only @@ -130,7 +131,8 @@ def load_dygraph(model_path, keep_name_table=False): params_file_path)) with open(params_file_path, 'rb') as f: - para_dict = pickle.load(f) + para_dict = pickle.load(f) if six.PY2 else pickle.load( + f, encoding='latin1') if not keep_name_table and "StructuredToParameterName@@" in para_dict: del para_dict["StructuredToParameterName@@"] @@ -138,6 +140,7 @@ def load_dygraph(model_path, keep_name_table=False): opti_file_path = model_path + ".pdopt" if os.path.exists(opti_file_path): with open(opti_file_path, 'rb') as f: - opti_dict = pickle.load(f) + opti_dict = pickle.load(f) if six.PY2 else pickle.load( + f, encoding='latin1') return para_dict, opti_dict diff --git a/python/paddle/fluid/io.py b/python/paddle/fluid/io.py index 09adb1fb13dcbf33ede404873c7a16528b3cecbd..448eb49563aa8caed389b43c0650c4f3e68685be 100644 --- a/python/paddle/fluid/io.py +++ b/python/paddle/fluid/io.py @@ -800,7 +800,7 @@ def load_vars(executor, var_temp = paddle.fluid.global_scope().find_var(each_var.name) assert var_temp != None, "can't not find var: " + each_var.name new_shape = (np.array(var_temp.get_tensor())).shape - assert each_var.name in orig_para_shape, earch_var.name + "MUST in var list" + assert each_var.name in orig_para_shape, each_var.name + "MUST in var list" orig_shape = orig_para_shape.get(each_var.name) if new_shape != orig_shape: raise RuntimeError( @@ -1579,14 +1579,14 @@ def save(program, model_path): parameter_list = list(filter(is_parameter, program.list_vars())) param_dict = {p.name: get_tensor(p) for p in parameter_list} with open(model_path + ".pdparams", 'wb') as f: - pickle.dump(param_dict, f) + pickle.dump(param_dict, f, protocol=2) optimizer_var_list = list( filter(is_belong_to_optimizer, program.list_vars())) opt_dict = {p.name: get_tensor(p) for p in optimizer_var_list} with open(model_path + ".pdopt", 'wb') as f: - pickle.dump(opt_dict, f) + pickle.dump(opt_dict, f, protocol=2) main_program = program.clone() program.desc.flush() @@ -1733,7 +1733,8 @@ def load(program, model_path, executor=None, var_list=None): global_scope(), executor._default_executor) with open(parameter_file_name, 'rb') as f: - load_dict = pickle.load(f) + load_dict = pickle.load(f) if six.PY2 else pickle.load( + f, encoding='latin1') for v in parameter_list: assert v.name in load_dict, \ "Can not find [{}] in model file [{}]".format( @@ -1753,7 +1754,8 @@ def load(program, model_path, executor=None, var_list=None): optimizer_var_list, global_scope(), executor._default_executor) with open(opt_file_name, 'rb') as f: - load_dict = pickle.load(f) + load_dict = pickle.load(f) if six.PY2 else pickle.load( + f, encoding='latin1') for v in optimizer_var_list: assert v.name in load_dict, \ "Can not find [{}] in model file [{}]".format( @@ -1877,12 +1879,14 @@ def load_program_state(model_path, var_list=None): "Parameter file [{}] not exits".format(parameter_file_name) with open(parameter_file_name, 'rb') as f: - para_dict = pickle.load(f) + para_dict = pickle.load(f) if six.PY2 else pickle.load( + f, encoding='latin1') opt_file_name = model_prefix + ".pdopt" if os.path.exists(opt_file_name): with open(opt_file_name, 'rb') as f: - opti_dict = pickle.load(f) + opti_dict = pickle.load(f) if six.PY2 else pickle.load( + f, encoding='latin1') para_dict.update(opti_dict)