diff --git a/paddleslim/auto_compression/compressor.py b/paddleslim/auto_compression/compressor.py index 5885387f0e325a6fc11aaf68f83ba58a63665503..144290e75ed42417a8d1ca3285a52df8fcb86346 100644 --- a/paddleslim/auto_compression/compressor.py +++ b/paddleslim/auto_compression/compressor.py @@ -586,28 +586,11 @@ class AutoCompression: tmp_model_path = os.path.join( self.tmp_dir, 'strategy_{}'.format(str(strategy_idx + 1))) final_model_path = os.path.join(self.final_dir) - if not os.path.exists(final_model_path): - os.makedirs(final_model_path) - - tmp_model_file = ".".join([tmp_model_path, "pdmodel"]) - if not os.path.exists(tmp_model_file): - tmp_model_file = os.path.join(tmp_model_path, self.model_filename) - - tmp_params_file = ".".join([tmp_model_path, "pdiparams"]) - if not os.path.exists(tmp_params_file): - tmp_params_file = os.path.join(tmp_model_path, self.params_filename) - - if self.model_filename is None: - self.model_filename = "infer.pdmodel" - if self.params_filename is None: - self.params_filename = "infer.pdiparams" - - final_model_file = os.path.join(final_model_path, self.model_filename) - final_params_file = os.path.join(final_model_path, self.params_filename) - if paddle.distributed.get_rank() == 0: - shutil.move(tmp_model_file, final_model_file) - shutil.move(tmp_params_file, final_params_file) + for _file in os.listdir(tmp_model_path): + _file_path = os.path.join(tmp_model_path, _file) + if os.path.isfile(_file_path): + shutil.copy(_file_path, final_model_path) shutil.rmtree(self.tmp_dir) _logger.info( "==> The ACT compression has been completed and the final model is saved in `{}`".