Add pantheon import in all & Fix write error in py3 (PaddlePaddle#112)

lizexu123 · Feb 14, 2020 · a4f4298 · a4f4298
1 parent a7b93e0
commit a4f4298
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 4 deletions.
diff --git a/paddleslim/__init__.py b/paddleslim/__init__.py
@@ -19,4 +19,5 @@
 from paddleslim import analysis
 from paddleslim import dist
 from paddleslim import quant
-__all__ = ['models', 'prune', 'nas', 'analysis', 'dist', 'quant']
+from paddleslim import pantheon
+__all__ = ['models', 'prune', 'nas', 'analysis', 'dist', 'quant', 'pantheon']
diff --git a/paddleslim/pantheon/README.md b/paddleslim/pantheon/README.md
@@ -13,7 +13,7 @@ The illustration below shows an application of Pantheon, where the sudent model
 ## Prerequisites
 
 - Python 2.7.x or 3.x
-- PaddlePaddle >= 1.6.0
+- PaddlePaddle >= 1.7.0
 
 ## APIs
 

diff --git a/paddleslim/pantheon/student.py b/paddleslim/pantheon/student.py
@@ -158,7 +158,7 @@ def get_cmd():
 
                     if end_recved:
                         break
-                    with open(in_path, 'r') as fin:
+                    with open(in_path, 'rb') as fin:
                         # get knowledge desc
                         desc = pickle.load(fin)
                         out_queue.put(desc)

diff --git a/paddleslim/pantheon/teacher.py b/paddleslim/pantheon/teacher.py
@@ -151,7 +151,7 @@ def start(self):
             self._t2s_queue = None
             self._cmd_queue = None
 
-        self._out_file = open(self._out_path, "w") if self._out_path else None
+        self._out_file = open(self._out_path, "wb") if self._out_path else None
         if self._out_file:
             return