From d93f39eb32a31e8152e805b874beab5670401251 Mon Sep 17 00:00:00 2001
From: Guanghua Yu <742925032@qq.com>
Date: Mon, 22 Aug 2022 10:22:23 +0800
Subject: [PATCH] [cherry-pick]update paddle2onnx version (#1373)

---
 paddleslim/auto_compression/compressor.py | 5 ++++-
 paddleslim/common/load_model.py           | 8 ++++++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/paddleslim/auto_compression/compressor.py b/paddleslim/auto_compression/compressor.py
index 11b80573..b5459c76 100644
--- a/paddleslim/auto_compression/compressor.py
+++ b/paddleslim/auto_compression/compressor.py
@@ -846,9 +846,12 @@ class AutoCompression:
         assert os.path.exists(
             infer_model_path), 'Not found {}, please check it.'.format(
                 infer_model_path)
+        onnx_save_path = os.path.join(self.final_dir, 'ONNX')
+        if not os.path.exists(onnx_save_path):
+            os.makedirs(onnx_save_path)
         export_onnx(
             self.final_dir,
             model_filename=self.model_filename,
             params_filename=self.params_filename,
-            save_file_path=os.path.join(self.final_dir, model_name),
+            save_file_path=os.path.join(onnx_save_path, model_name),
             deploy_backend=deploy_backend)
diff --git a/paddleslim/common/load_model.py b/paddleslim/common/load_model.py
index cc545b31..1afa312f 100644
--- a/paddleslim/common/load_model.py
+++ b/paddleslim/common/load_model.py
@@ -210,7 +210,7 @@ def export_onnx(model_dir,
         pkg.require('paddle2onnx')
     except:
         from pip._internal import main
-        main(['install', 'paddle2onnx==1.0.0rc3'])
+        main(['install', 'paddle2onnx==1.0.0rc4'])
     import paddle2onnx
     paddle2onnx.command.c_paddle_to_onnx(
         model_file=os.path.join(model_dir, model_filename),
@@ -218,5 +218,9 @@ def export_onnx(model_dir,
         save_file=save_file_path,
         opset_version=opset_version,
         enable_onnx_checker=True,
-        deploy_backend=deploy_backend)
+        deploy_backend=deploy_backend,
+        scale_file=os.path.join(model_dir, 'calibration_table.txt'),
+        calibration_file=os.path.join(
+            save_file_path.rstrip(os.path.split(save_file_path)[-1]),
+            'calibration.cache'))
     _logger.info('Convert model to ONNX: {}'.format(save_file_path))
-- 
GitLab