[WIP]support compress

paddlex/engine.py

@@ -53,5 +53,8 @@ def run(self):
                 res.print()
                 if self._output:
                     res.save_all(save_path=self._output)
+        elif self._mode == "compress":
+            disable_pir_bydefault()
+            return self._model.compress()
         else:
             raise_unsupported_api_error(f"{self._mode}", self.__class__)

paddlex/inference/utils/pp_option.py

@@ -87,7 +87,7 @@ def _get_default_config(self):
         device_type, device_ids = parse_device(get_default_device())
 
         default_config = {
-            "run_mode": "paddle",
+            "run_mode": "trt_fp32",
             "device_type": device_type,
             "device_id": None if device_ids is None else device_ids[0],
             "cpu_threads": 8,

paddlex/model.py

@@ -21,6 +21,7 @@
     build_trainer,
     build_evaluater,
     build_exportor,
+    build_compressor,
 )
 
 
@@ -115,3 +116,7 @@ def export(self):
     def predict(self):
         predict_kwargs, predictor = self._build_predictor()
         yield from predictor(**predict_kwargs)
+
+    def compress(self):
+        compressor = build_compressor(self._config)
+        return compressor.compress()

paddlex/modules/__init__.py

@@ -18,13 +18,15 @@
     build_trainer,
     build_evaluater,
     build_exportor,
+    build_compressor,
 )
 
 from .image_classification import (
     ClsDatasetChecker,
     ClsTrainer,
     ClsEvaluator,
     ClsExportor,
+    ClsCompressor,
 )
 
 from .multilabel_classification import (

paddlex/modules/base/__init__.py

@@ -16,3 +16,4 @@
 from .trainer import build_trainer, BaseTrainer
 from .evaluator import build_evaluater, BaseEvaluator
 from .exportor import build_exportor, BaseExportor
+from .compressor import build_compressor, BaseCompressor

paddlex/modules/image_classification/__init__.py

@@ -16,3 +16,4 @@
 from .dataset_checker import ClsDatasetChecker
 from .evaluator import ClsEvaluator
 from .exportor import ClsExportor
+from .compressor import ClsCompressor

paddlex/modules/object_detection/__init__.py

@@ -16,3 +16,4 @@
 from .dataset_checker import COCODatasetChecker
 from .evaluator import DetEvaluator
 from .exportor import DetExportor
+from .compressor import DetCompressor

paddlex/repo_apis/PaddleClas_api/cls/model.py

@@ -325,6 +325,8 @@ def compression(
             config._update_use_vdl(use_vdl)
             config._update_slim_config(self.model_info["auto_compression_config_path"])
             config.update_pretrained_weights(weight_path)
+            uniform_output_enabled = kwargs.pop("uniform_output_enabled", False)
+            config.update([f"Global.uniform_output_enabled={uniform_output_enabled}"])
 
             if batch_size is not None:
                 config.update_batch_size(batch_size)

paddlex/repo_apis/PaddleClas_api/cls/register.py

@@ -147,7 +147,10 @@
         "model_name": "PP-LCNet_x1_0",
         "suite": "Cls",
         "config_path": osp.join(PDX_CONFIG_DIR, "PP-LCNet_x1_0.yaml"),
-        "supported_apis": ["train", "evaluate", "predict", "export"],
+        "auto_compression_config_path": osp.join(
+            PDX_CONFIG_DIR, "slim/PP_LCNet_x1_0_quantization.yaml"
+        ),
+        "supported_apis": ["train", "evaluate", "predict", "export", "compression"],
     }
 )
 

paddlex/repo_apis/PaddleClas_api/cls/runner.py

@@ -174,7 +174,7 @@ def compression(
         ]
         cp_export = self.export(config_path, export_cli_args, device)
 
-        return cp_train, cp_export
+        return cp_export
 
 
 def _extract_eval_metrics(stdout: str) -> dict:

paddlex/repo_apis/PaddleDetection_api/object_det/model.py

@@ -407,6 +407,8 @@ def compression(
         else:
             save_dir = abspath(save_dir)
         cps_config.update_save_dir(save_dir)
+        uniform_output_enabled = kwargs.pop("uniform_output_enabled", False)
+        config.update({"uniform_output_enabled": uniform_output_enabled})
         if use_vdl:
             train_cli_args.append(CLIArgument("--use_vdl", use_vdl))
             train_cli_args.append(CLIArgument("--vdl_log_dir", save_dir))
@@ -425,6 +427,6 @@ def compression(
 
             self._assert_empty_kwargs(kwargs)
 
-            self.runner.compression(
+            return self.runner.compression(
                 config_path, train_cli_args, export_cli_args, device, save_dir
             )

paddlex/repo_apis/PaddleDetection_api/object_det/register.py

@@ -132,7 +132,10 @@
         "model_name": "RT-DETR-L",
         "suite": "Det",
         "config_path": osp.join(PDX_CONFIG_DIR, "RT-DETR-L.yaml"),
-        "supported_apis": ["train", "evaluate", "predict", "export"],
+        "auto_compression_config_path": osp.join(
+            PDX_CONFIG_DIR, "slim", "rtdetr_qat.yml"
+        ),
+        "supported_apis": ["train", "evaluate", "predict", "export", "compression"],
         "supported_dataset_types": ["COCODetDataset"],
         "supported_train_opts": {
             "device": ["cpu", "gpu_nxcx", "xpu", "npu", "mlu"],
@@ -1080,7 +1083,17 @@
         "model_name": "PP-DocLayout-M",
         "suite": "Det",
         "config_path": osp.join(PDX_CONFIG_DIR, "PP-DocLayout-M.yaml"),
-        "supported_apis": ["train", "evaluate", "predict", "export", "infer"],
+        "supported_apis": [
+            "train",
+            "evaluate",
+            "predict",
+            "export",
+            "infer",
+            "compression",
+        ],
+        "auto_compression_config_path": osp.join(
+            PDX_CONFIG_DIR, "slim", "PP-DocLayout-M_qat.yaml"
+        ),
         "supported_dataset_types": ["COCODetDataset"],
         "supported_train_opts": {
             "device": ["cpu", "gpu_nxcx", "xpu", "npu", "mlu"],

paddlex/repo_apis/PaddleDetection_api/object_det/runner.py

@@ -192,7 +192,7 @@ def compression(
         ]
         cp_export = self.run_cmd(cmd, switch_wdir=True, echo=True, silent=False)
 
-        return cp_train, cp_export
+        return cp_export
 
     def _gather_opts_args(self, args):
         """_gather_opts_args"""