optim performance and cls/rec batchsize for inference (#480)

Author: liangxhao

deploy/py_infer/src/data_process/utils/constants.py

@@ -1,6 +1,9 @@
 import argparse
+import gc
 from abc import ABCMeta, abstractmethod
 
+from ..core import Model
+
 
 class InferBase(metaclass=ABCMeta):
     """
@@ -10,10 +13,47 @@ class InferBase(metaclass=ABCMeta):
     def __init__(self, args: argparse.Namespace, **kwargs):
         super().__init__()
         self.args = args
+
         self.model = None
+        self._bs_list = []
+        self._hw_list = []
+
+    def init(self, *, preprocess=True, model=True, postprocess=True):
+        if preprocess or model:
+            self._init_model()
+
+        if preprocess:
+            self._init_preprocess()
+
+        if postprocess:
+            self._init_postprocess()
+
+        if not model:
+            self.free_model()
+
+        if model:
+            if isinstance(self.model, dict):
+                for _model in self.model.values():
+                    _model.warmup()
+            elif isinstance(self.model, Model):
+                self.model.warmup()
+            else:
+                pass
+
+    @abstractmethod
+    def _init_preprocess(self):
+        pass
 
     @abstractmethod
-    def init(self, **kwargs):
+    def _init_model(self):
+        pass
+
+    @abstractmethod
+    def _init_postprocess(self):
+        pass
+
+    @abstractmethod
+    def get_params(self):
         pass
 
     @abstractmethod
@@ -34,12 +74,13 @@ def postprocess(self, *args, **kwargs):
 
     def free_model(self):
         if hasattr(self, "model") and self.model:
-            if isinstance(self.model, (tuple, list)):
-                for model in self.model:
-                    del model
-            else:
-                del self.model
-        self.model = None
+            if isinstance(self.model, dict):
+                self.model.clear()
+
+            del self.model
+            gc.collect()
+
+            self.model = None
 
     def __del__(self):
         self.free_model()

deploy/py_infer/src/infer/infer_base.py

@@ -11,9 +11,12 @@
 class TextClassifier(InferBase):
     def __init__(self, args):
         super(TextClassifier, self).__init__(args)
-        self._bs_list = []
 
-    def init(self, warmup=False):
+    def _init_preprocess(self):
+        preprocess_ops = build_preprocess(self.args.cls_config_path)
+        self.preprocess_ops = functools.partial(preprocess_ops, target_size=self._hw_list[0])
+
+    def _init_model(self):
         self.model = Model(
             backend=self.args.backend, model_path=self.args.cls_model_path, device_id=self.args.device_id
         )
@@ -28,12 +31,13 @@ def init(self, warmup=False):
             batchsize, _, model_height, model_width = shape_info
             self._bs_list = [batchsize]
 
-        preprocess_ops = build_preprocess(self.args.cls_config_path)
-        self.preprocess_ops = functools.partial(preprocess_ops, target_size=(model_height, model_width))
+        self._hw_list = [(model_height, model_width)]
+
+    def _init_postprocess(self):
         self.postprocess_ops = build_postprocess(self.args.cls_config_path)
 
-        if warmup:
-            self.model.warmup()
+    def get_params(self):
+        return {"cls_batch_num": self._bs_list}
 
     def __call__(self, image: Union[np.ndarray, List[np.ndarray]]):
         images = [image] if isinstance(image, np.ndarray) else image

deploy/py_infer/src/infer/infer_cls.py

@@ -10,9 +10,11 @@
 class TextDetector(InferBase):
     def __init__(self, args):
         super(TextDetector, self).__init__(args)
-        self._hw_list = []
 
-    def init(self, warmup=False):
+    def _init_preprocess(self):
+        self.preprocess_ops = build_preprocess(self.args.det_config_path)
+
+    def _init_model(self):
         self.model = Model(
             backend=self.args.backend,
             model_path=self.args.det_model_path,
@@ -34,11 +36,13 @@ def init(self, warmup=False):
             raise ValueError("Input batch size must be 1 for detection model.")
 
         self._hw_list = hw_list
-        self.preprocess_ops = build_preprocess(self.args.det_config_path)
+        self._bs_list = [batchsize]
+
+    def _init_postprocess(self):
         self.postprocess_ops = build_postprocess(self.args.det_config_path, rescale_fields=["polys"])
 
-        if warmup:
-            self.model.warmup()
+    def get_params(self):
+        return {"det_batch_num": self._bs_list}
 
     def __call__(self, image: np.ndarray):
         data = self.preprocess(image)
@@ -54,9 +58,7 @@ def preprocess(self, image: np.ndarray) -> Dict:
     def model_infer(self, data: Dict) -> List[np.ndarray]:
         return self.model.infer([data["image"]])  # model infer for single input
 
-    def postprocess(self, pred, shape_list: np.ndarray) -> np.ndarray:
+    def postprocess(self, pred, shape_list: np.ndarray) -> List[np.ndarray]:
         polys = self.postprocess_ops(tuple(pred), shape_list)["polys"][0]  # {'polys': [img0_polys, ...], ...}
-        polys = np.array(polys)
-        # polys.shape may be (0,), (polys_num, points_num, 2), (1, polys_num, points_num, 2)
-        polys_shape = (-1, *polys.shape[-2:]) if polys.size != 0 else (0, 0, 2)
-        return polys.reshape(*polys_shape)  # (polys_num, points_num, 2), because bs=1
+        polys = [np.array(x) for x in polys]
+        return polys  # [poly(points_num, 2), ...], bs=1

deploy/py_infer/src/infer/infer_det.py

@@ -1,6 +1,5 @@
 import math
 import os
-from collections import defaultdict
 from typing import Dict, List, Tuple, Union
 
 import numpy as np
@@ -14,11 +13,7 @@
 class TextRecognizer(InferBase):
     def __init__(self, args):
         super(TextRecognizer, self).__init__(args)
-
-        self._hw_list = []
-        self._bs_list = []
-
-        self.model = defaultdict()
+        self.model: Dict[int, Model] = {}
         self.shape_type = None
 
     def __get_shape_for_single_model(self, filename):
@@ -65,7 +60,10 @@ def __get_resized_hw(self, image_list):
 
         return max_h, max_w
 
-    def init(self, warmup=False):
+    def _init_preprocess(self):
+        self.preprocess_ops = build_preprocess(self.args.rec_config_path)
+
+    def _init_model(self):
         model_path = self.args.rec_model_path
 
         if os.path.isfile(model_path):
@@ -91,13 +89,12 @@ def init(self, warmup=False):
 
         self._bs_list.sort()
 
-        self.preprocess_ops = build_preprocess(self.args.rec_config_path)
+    def _init_postprocess(self):
         params = {"character_dict_path": self.args.character_dict_path}
         self.postprocess_ops = build_postprocess(self.args.rec_config_path, **params)
 
-        if warmup:
-            for model in self.model.values():
-                model.warmup()
+    def get_params(self):
+        return {"rec_batch_num": self._bs_list}
 
     def __call__(self, image: Union[np.ndarray, List[np.ndarray]]):
         images = [image] if isinstance(image, np.ndarray) else image
@@ -129,10 +126,10 @@ def preprocess(self, image: List[np.ndarray]) -> Tuple[List[int], List[Dict]]:
         return split_bs, split_data
 
     def model_infer(self, data: Dict) -> List[np.ndarray]:
-        input = data["image"]
-        bs, *_ = input.shape
+        input_data = data["image"]
+        bs, *_ = input_data.shape
         n = bs if bs in self._bs_list else -1
-        return self.model[n].infer([input])
+        return self.model[n].infer([input_data])
 
     def postprocess(self, pred, batch=None):
         pred = gear_utils.get_batch_from_padding(pred, batch)

deploy/py_infer/src/infer/infer_rec.py

@@ -122,9 +122,6 @@ def update_task_info(args):
     task_order = (det, cls, rec)
     if task_order in task_map:
         setattr(args, "task_type", task_map[task_order])
-        setattr(args, "save_vis_det_save_dir", bool(args.vis_det_save_dir))
-        setattr(args, "save_vis_pipeline_save_dir", bool(args.vis_pipeline_save_dir))
-        setattr(args, "save_crop_res_dir", bool(args.crop_save_dir))
     else:
         unsupported_task_map = {
             (False, False, False): "empty",
@@ -163,9 +160,6 @@ def check_and_update_args(args):
     if args.rec_model_path and not args.rec_model_name_or_config:
         raise ValueError("rec_model_name_or_config can't be emtpy when set rec_model_path for recognition.")
 
-    if args.parallel_num < 1 or args.parallel_num > 4:
-        raise ValueError(f"parallel_num must between [1,4], current: {args.parallel_num}.")
-
     if args.crop_save_dir and args.task_type not in (TaskType.DET_REC, TaskType.DET_CLS_REC):
         raise ValueError("det_model_path and rec_model_path can't be empty when set crop_save_dir.")
 
@@ -219,11 +213,11 @@ def check_and_update_args(args):
 def init_save_dir(args):
     if args.res_save_dir:
         save_path_init(args.res_save_dir)
-    if args.save_crop_res_dir:
+    if args.crop_save_dir:
         save_path_init(args.crop_save_dir)
-    if args.save_vis_pipeline_save_dir:
+    if args.vis_pipeline_save_dir:
         save_path_init(args.vis_pipeline_save_dir)
-    if args.save_vis_det_save_dir:
+    if args.vis_det_save_dir:
         save_path_init(args.vis_det_save_dir)
     if args.save_log_dir:
         save_path_init(args.save_log_dir, exist_ok=True)

deploy/py_infer/src/infer_args.py

@@ -1,3 +1,3 @@
-from .message_data import ProfilingData, StartSign, StopSign
+from .message_data import ProfilingData, StopSign
 from .module_data import ModuleConnectDesc, ModuleDesc, ModuleInitArgs
 from .process_data import ProcessData, StopData

deploy/py_infer/src/parallel/datatype/__init__.py

@@ -1,11 +1,6 @@
 from dataclasses import dataclass
 
 
-@dataclass
-class StartSign:
-    start: bool = True
-
-
 @dataclass
 class StopSign:
     stop: bool = True

deploy/py_infer/src/parallel/datatype/message_data.py

@@ -6,20 +6,21 @@
 
 @dataclass
 class ProcessData:
-    # infer_test info
-    sub_image_total: int = 0
-    image_total: int = 0
-    infer_result: list = field(default_factory=lambda: [])
+    # skip each compute node
     skip: bool = False
+    # prediction results of each image
+    infer_result: list = field(default_factory=lambda: [])
 
     # image basic info
-    image_path: str = ""
-    image_name: str = ""
-    image_id: int = ""
-    frame: np.ndarray = None
+    image_path: List[str] = field(default_factory=lambda: [])
+    frame: List[np.ndarray] = field(default_factory=lambda: [])
 
+    # sub image of detection box, for det (+ cls) + rec
+    sub_image_total: int = 0  # len(sub_image_list_0) + len(sub_image_list_1) + ...
     sub_image_list: list = field(default_factory=lambda: [])
-    sub_image_size: int = 0
+    sub_image_size: int = 0  # len of sub_image_list
+
+    # data for preprocess -> infer -> postprocess
     data: Union[np.ndarray, List[np.ndarray], Dict] = None
 
 

deploy/py_infer/src/parallel/datatype/process_data.py

@@ -1,2 +1,3 @@
 from .module_base import ModuleBase
 from .module_manager import ModuleManager
+from .pipeline_manager import ParallelPipelineManager