Fix demo of pruning. (#87)

demo/prune/train.py

@@ -8,7 +8,7 @@ import math
 import time
 import numpy as np
 import paddle.fluid as fluid
-from paddleslim.prune import Pruner
+from paddleslim.prune import Pruner, save_model
 from paddleslim.common import get_logger
 from paddleslim.analysis import flops
 sys.path.append(sys.path[0] + "/../")
@@ -223,7 +223,7 @@ def compress(args):
         train(i, pruned_program)
         if i % args.test_period == 0:
             test(i, pruned_val_program)
-            save_model(pruned_val_program,
+            save_model(exe, pruned_val_program,
                        os.path.join(args.model_path, str(i)))
 
 

paddleslim/prune/__init__.py

@@ -23,8 +23,8 @@ from .sensitive import *
 from ..prune import sensitive
 from .prune_walker import *
 from ..prune import prune_walker
-from io import *
-from ..prune import io
+from .prune_io import *
+from ..prune import prune_io
 
 __all__ = []
 
@@ -33,4 +33,4 @@ __all__ += auto_pruner.__all__
 __all__ += sensitive_pruner.__all__
 __all__ += sensitive.__all__
 __all__ += prune_walker.__all__
-__all__ += io.__all__
+__all__ += prune_io.__all__