Fix pruning export (PaddlePaddle#2350)

heavengate · Mar 15, 2021 · 8635791 · 8635791
1 parent 73c861e
commit 8635791
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/slim/prune/export_model.py b/slim/prune/export_model.py
@@ -63,6 +63,9 @@ def main():
             test_fetches = model.test(feed_vars)
     infer_prog = infer_prog.clone(True)
 
+    exe.run(startup_prog)
+    checkpoint.load_checkpoint(exe, infer_prog, cfg.weights)
+
     pruned_params = FLAGS.pruned_params
     assert (
         FLAGS.pruned_params is not None
@@ -90,13 +93,9 @@ def main():
     logger.info("pruned FLOPS: {}".format(
         float(base_flops - pruned_flops) / base_flops))
 
-    exe.run(startup_prog)
-    checkpoint.load_checkpoint(exe, infer_prog, cfg.weights)
-
     dump_infer_config(FLAGS, cfg)
     save_infer_model(FLAGS, exe, feed_vars, test_fetches, infer_prog)
 
-
 if __name__ == '__main__':
     enable_static_mode()
     parser = ArgsParser()