From cb4560b79d7cae63652b8c041b993c19ced14355 Mon Sep 17 00:00:00 2001
From: songyouwei <youwei0314@gmail.com>
Date: Sat, 15 Feb 2020 20:21:12 +0800
Subject: [PATCH] python 2 and 3 compatible pickle (#22555)

test=develop
---
 python/paddle/fluid/dygraph/checkpoint.py |  9 ++++++---
 python/paddle/fluid/io.py                 | 18 +++++++++++-------
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/python/paddle/fluid/dygraph/checkpoint.py b/python/paddle/fluid/dygraph/checkpoint.py
index 27658ba3d46..6bda0996694 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='bytes')
 
     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='bytes')
 
     return para_dict, opti_dict
diff --git a/python/paddle/fluid/io.py b/python/paddle/fluid/io.py
index d8d650026be..7e65060829a 100644
--- a/python/paddle/fluid/io.py
+++ b/python/paddle/fluid/io.py
@@ -765,7 +765,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(
@@ -1541,14 +1541,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()
@@ -1695,7 +1695,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='bytes')
     for v in parameter_list:
         assert v.name in load_dict, \
             "Can not find [{}] in model file [{}]".format(
@@ -1715,7 +1716,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='bytes')
         for v in optimizer_var_list:
             assert v.name in load_dict, \
                 "Can not find [{}] in model file [{}]".format(
@@ -1839,12 +1841,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='bytes')
 
     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='bytes')
 
         para_dict.update(opti_dict)
 
-- 
GitLab