support multi-batch test in profile tool (#868)

* test batch profile with resnet pspnet yolov3 srcnn

* update doc

* update docs

* fix ut

* fix mmdet

* support batch mmorc and mmrotate

* fix mmcls export to sdk

* resolve comments

* rename to fix #819

* fix conflicts with master

docs/en/02-how-to-run/useful_tools.md

@@ -132,23 +132,25 @@ python tools/onnx2ncnn.py \
 - `output_bin` : The converted `ncnn` bin path.
 - `--log-level` : To set log level which in `'CRITICAL', 'FATAL', 'ERROR', 'WARN', 'WARNING', 'INFO', 'DEBUG', 'NOTSET'`. If not specified, it will be set to `INFO`.
 
-## profile
+## profiler
 
 This tool helps to test latency of models with PyTorch, TensorRT and other backends. Note that the pre- and post-processing is excluded when computing inference latency.
 
 ### Usage
 
 ```bash
-python tools/profile.py \
+python tools/profiler.py \
     ${DEPLOY_CFG} \
     ${MODEL_CFG} \
     ${IMAGE_DIR} \
     --model ${MODEL} \
     --device ${DEVICE} \
     --shape ${SHAPE} \
-    --num-iter {NUM_ITER} \
-    --warmup {WARMUP}
-    --cfg-options ${CFG_OPTIONS}
+    --num-iter ${NUM_ITER} \
+    --warmup ${WARMUP} \
+    --cfg-options ${CFG_OPTIONS} \
+    --batch-size ${BATCH_SIZE} \
+    --img-ext ${IMG_EXT}
 ```
 
 ### Description of all arguments
@@ -162,11 +164,13 @@ python tools/profile.py \
 - `--warmup` : Number of iteration to warm-up the machine. Default is `10`.
 - `--device` : The device type. If not specified, it will be set to `cuda:0`.
 - `--cfg-options` : Optional key-value pairs to be overrode for model config.
+- `--batch-size`: the batch size for test inference. Default is `1`. Note that not all models support `batch_size>1`.
+- `--img-ext`: the file extensions for input images from `image_dir`. Defaults to `['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']`.
 
 ### Example:
 
 ```shell
-python tools/profile.py \
+python tools/profiler.py \
     configs/mmcls/classification_tensorrt_dynamic-224x224-224x224.py \
     ../mmclassification/configs/resnet/resnet18_8xb32_in1k.py \
     ../mmdetection/demo \
@@ -175,6 +179,7 @@ python tools/profile.py \
     --shape 224x224 \
     --num-iter 100 \
     --warmup 10 \
+    --batch-size 1
 ```
 
 And the output look like this:

docs/zh_cn/02-how-to-run/useful_tools.md

@@ -132,23 +132,25 @@ python tools/onnx2ncnn.py \
 - `output_bin` : The converted `ncnn` bin path.
 - `--log-level` : To set log level which in `'CRITICAL', 'FATAL', 'ERROR', 'WARN', 'WARNING', 'INFO', 'DEBUG', 'NOTSET'`. If not specified, it will be set to `INFO`.
 
-## profile
+## profiler
 
 这个工具用来测试 torch 和 trt 等后端的速度，注意测试不包含前后处理。
 
 ### 用法
 
 ```bash
-python tools/profile.py \
+python tools/profiler.py \
     ${DEPLOY_CFG} \
     ${MODEL_CFG} \
     ${IMAGE_DIR} \
     --model ${MODEL} \
     --device ${DEVICE} \
     --shape ${SHAPE} \
-    --num-iter {NUM_ITER} \
-    --warmup {WARMUP}
-    --cfg-options ${CFG_OPTIONS}
+    --num-iter ${NUM_ITER} \
+    --warmup ${WARMUP} \
+    --cfg-options ${CFG_OPTIONS} \
+    --batch-size ${BATCH_SIZE} \
+    --img-ext ${IMG_EXT}
 ```
 
 ### 参数说明
@@ -162,11 +164,13 @@ python tools/profile.py \
 - `--warmup` : Number of iteration to warm-up the machine. Default is `10`.
 - `--device` : The device type. If not specified, it will be set to `cuda:0`.
 - `--cfg-options` : Optional key-value pairs to be overrode for model config.
+- `--batch-size`: the batch size for test inference. Default is `1`. Note that not all models support `batch_size>1`.
+- `--img-ext`: the file extensions for input images from `image_dir`. Defaults to `['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']`.
 
 ### 使用举例
 
 ```shell
-python tools/profile.py \
+python tools/profiler.py \
     configs/mmcls/classification_tensorrt_dynamic-224x224-224x224.py \
     ../mmclassification/configs/resnet/resnet18_8xb32_in1k.py \
     ../mmdetection/demo \
@@ -175,6 +179,7 @@ python tools/profile.py \
     --shape 224x224 \
     --num-iter 100 \
     --warmup 10 \
+    --batch-size 1
 ```
 
 输出：

mmdeploy/codebase/base/task.py

@@ -140,14 +140,14 @@ def single_gpu_test(self,
 
     @abstractmethod
     def create_input(self,
-                     imgs: Union[str, np.ndarray],
+                     imgs: Union[str, np.ndarray, Sequence],
                      input_shape: Sequence[int] = None,
                      **kwargs) -> Tuple[Dict, torch.Tensor]:
         """Create input for model.
 
         Args:
-            imgs (str | np.ndarray): Input image(s), accepted data types are
-                `str`, `np.ndarray`.
+            imgs (str | np.ndarray | Sequence): Input image(s),
+                accepted data types are `str`, `np.ndarray`.
             input_shape (list[int]): Input shape of image in (width, height)
                 format, defaults to `None`.
 

mmdeploy/codebase/mmcls/deploy/classification.py

@@ -14,22 +14,22 @@
 
 
 def process_model_config(model_cfg: mmcv.Config,
-                         imgs: Union[str, np.ndarray],
+                         imgs: Sequence[Union[str, np.ndarray]],
                          input_shape: Optional[Sequence[int]] = None):
     """Process the model config.
 
     Args:
         model_cfg (mmcv.Config): The model config.
-        imgs (str | np.ndarray): Input image(s), accepted data type are `str`,
-            `np.ndarray`.
+        imgs (Sequence[str | np.ndarray]): Input image(s), accepted
+            data type are `str`, `np.ndarray`.
         input_shape (list[int]): A list of two integer in (width, height)
             format specifying input shape. Default: None.
 
     Returns:
         mmcv.Config: the model config after processing.
     """
     cfg = model_cfg.deepcopy()
-    if isinstance(imgs, str):
+    if isinstance(imgs[0], str):
         if cfg.data.test.pipeline[0]['type'] != 'LoadImageFromFile':
             cfg.data.test.pipeline.insert(0, dict(type='LoadImageFromFile'))
     else:
@@ -110,14 +110,14 @@ def init_pytorch_model(self,
         return model.eval()
 
     def create_input(self,
-                     imgs: Union[str, np.ndarray],
+                     imgs: Union[str, np.ndarray, Sequence],
                      input_shape: Optional[Sequence[int]] = None) \
             -> Tuple[Dict, torch.Tensor]:
         """Create input for classifier.
 
         Args:
-            imgs (Any): Input image(s), accepted data type are `str`,
-                `np.ndarray`, `torch.Tensor`.
+            imgs (Union[str, np.ndarray, Sequence]): Input image(s),
+                accepted data type are `str`, `np.ndarray`, Sequence.
             input_shape (list[int]): A list of two integer in (width, height)
                 format specifying input shape. Default: None.
 
@@ -126,15 +126,19 @@ def create_input(self,
         """
         from mmcls.datasets.pipelines import Compose
         from mmcv.parallel import collate, scatter
+        if isinstance(imgs, (str, np.ndarray)):
+            imgs = [imgs]
         cfg = process_model_config(self.model_cfg, imgs, input_shape)
-        if isinstance(imgs, str):
-            data = dict(img_info=dict(filename=imgs), img_prefix=None)
-        else:
-            data = dict(img=imgs)
+        data_list = []
         test_pipeline = Compose(cfg.data.test.pipeline)
-        data = test_pipeline(data)
-        data = collate([data], samples_per_gpu=1)
-        data['img'] = [data['img']]
+        for img in imgs:
+            if isinstance(img, str):
+                data = dict(img_info=dict(filename=img), img_prefix=None)
+            else:
+                data = dict(img=img)
+            data = test_pipeline(data)
+            data_list.append(data)
+        data = collate(data_list, samples_per_gpu=len(data_list))
         if self.device != 'cpu':
             data = scatter(data, [self.device])[0]
         return data, data['img']
@@ -272,7 +276,7 @@ def get_preprocess(self) -> Dict:
             dict: Composed of the preprocess information.
         """
         input_shape = get_input_shape(self.deploy_cfg)
-        cfg = process_model_config(self.model_cfg, '', input_shape)
+        cfg = process_model_config(self.model_cfg, [''], input_shape)
         preprocess = cfg.data.test.pipeline
         return preprocess
 

mmdeploy/codebase/mmdet/deploy/object_detection.py

@@ -99,7 +99,7 @@ def init_pytorch_model(self,
         return model.eval()
 
     def create_input(self,
-                     imgs: Union[str, np.ndarray],
+                     imgs: Union[str, np.ndarray, Sequence],
                      input_shape: Sequence[int] = None) \
             -> Tuple[Dict, torch.Tensor]:
         """Create input for detector.
@@ -115,7 +115,7 @@ def create_input(self,
         """
         from mmcv.parallel import collate, scatter
         from mmdet.datasets.pipelines import Compose
-        if not isinstance(imgs, (list, tuple)):
+        if isinstance(imgs, (str, np.ndarray)):
             imgs = [imgs]
         dynamic_flag = is_dynamic_shape(self.deploy_cfg)
         cfg = process_model_config(self.model_cfg, imgs, input_shape)
@@ -146,10 +146,11 @@ def create_input(self,
 
         data = collate(data_list, samples_per_gpu=len(imgs))
 
-        data['img_metas'] = [
-            img_metas.data[0] for img_metas in data['img_metas']
-        ]
-        data['img'] = [img.data[0] for img in data['img']]
+        for k, v in data.items():
+            # batch_size > 1
+            if isinstance(v[0], DataContainer):
+                data[k] = v[0].data
+
         if self.device != 'cpu':
             data = scatter(data, [self.device])[0]
 

mmdeploy/codebase/mmedit/deploy/super_resolution.py

@@ -154,11 +154,11 @@ def create_input(self,
 
         data = collate(data_arr, samples_per_gpu=len(imgs))
 
-        data['img'] = data['lq']
-
         if self.device != 'cpu':
             data = scatter(data, [self.device])[0]
 
+        data['img'] = data['lq']
+
         return data, data['img']
 
     def visualize(self,

mmdeploy/codebase/mmocr/deploy/text_detection.py

@@ -107,7 +107,7 @@ def init_pytorch_model(self,
         return model.eval()
 
     def create_input(self,
-                     imgs: Union[str, np.ndarray],
+                     imgs: Union[str, np.ndarray, Sequence],
                      input_shape: Sequence[int] = None) \
             -> Tuple[Dict, torch.Tensor]:
         """Create input for segmentor.
@@ -149,32 +149,17 @@ def create_input(self,
             # get tensor from list to stack for batch mode (text detection)
             data_list.append(data)
 
-        if isinstance(data_list[0]['img'], list) and len(data_list) > 1:
-            raise Exception('aug test does not support '
-                            f'inference with batch size '
-                            f'{len(data_list)}')
+        batch_data = collate(data_list, samples_per_gpu=len(imgs))
 
-        data = collate(data_list, samples_per_gpu=len(imgs))
-
-        # process img_metas
-        if isinstance(data['img_metas'], list):
-            data['img_metas'] = [
-                img_metas.data[0] for img_metas in data['img_metas']
-            ]
-        else:
-            data['img_metas'] = data['img_metas'].data
-
-        if isinstance(data['img'], list):
-            data['img'] = [img.data for img in data['img']]
-            if isinstance(data['img'][0], list):
-                data['img'] = [img[0] for img in data['img']]
-        else:
-            data['img'] = data['img'].data
+        for k, v in batch_data.items():
+            # batch_size > 1
+            if isinstance(v[0], DataContainer):
+                batch_data[k] = v[0].data
 
         if self.device != 'cpu':
-            data = scatter(data, [self.device])[0]
+            batch_data = scatter(batch_data, [self.device])[0]
 
-        return data, data['img']
+        return batch_data, batch_data['img']
 
     def visualize(self,
                   model: nn.Module,

mmdeploy/codebase/mmocr/deploy/text_recognition.py

@@ -161,27 +161,17 @@ def create_input(self,
                             f'inference with batch size '
                             f'{len(data_list)}')
 
-        data = collate(data_list, samples_per_gpu=len(imgs))
+        batch_data = collate(data_list, samples_per_gpu=len(imgs))
 
-        # process img_metas
-        if isinstance(data['img_metas'], list):
-            data['img_metas'] = [
-                img_metas.data[0] for img_metas in data['img_metas']
-            ]
-        else:
-            data['img_metas'] = data['img_metas'].data
-
-        if isinstance(data['img'], list):
-            data['img'] = [img.data for img in data['img']]
-            if isinstance(data['img'][0], list):
-                data['img'] = [img[0] for img in data['img']]
-        else:
-            data['img'] = data['img'].data
+        for k, v in batch_data.items():
+            # batch_size > 1
+            if isinstance(v, DataContainer):
+                batch_data[k] = v.data[0]
 
         if self.device != 'cpu':
-            data = scatter(data, [self.device])[0]
+            batch_data = scatter(batch_data, [self.device])[0]
 
-        return data, data['img']
+        return batch_data, batch_data['img']
 
     def visualize(self,
                   model: nn.Module,