feat(codebase): support mmdet3d dev-1.x 1.1.0rc1

configs/mmdet3d/voxel-detection/voxel-detection_static.py

@@ -2,5 +2,5 @@
 codebase_config = dict(
     type='mmdet3d', task='VoxelDetection', model_type='end2end')
 onnx_config = dict(
-    input_names=['voxels', 'num_points', 'coors'],
-    output_names=['scores', 'bbox_preds', 'dir_scores'])
+    input_names=['placeholder', 'num_points', 'voxels', 'coors'],
+    output_names=['cls_score', 'bbox_pred', 'dir_cls_pred'])

mmdeploy/apis/onnx/export.py

@@ -119,7 +119,6 @@ def _add_or_update(cfg: dict, key: str, val: Any):
             ), f'Expect input_metas type is dict, get {type(input_metas)}.'
             model_forward = model.forward
             model.forward = partial(model.forward, **input_metas)
-
         torch.onnx.export(
             patched_model,
             args,

mmdeploy/apis/pytorch2onnx.py

@@ -3,7 +3,6 @@
 from typing import Any, Optional, Union
 
 import mmengine
-import torch
 
 from mmdeploy.apis.core.pipeline_manager import no_mp
 from mmdeploy.utils import (Backend, get_backend, get_dynamic_axes,
@@ -64,11 +63,12 @@ def torch2onnx(img: Any,
         img,
         input_shape,
         data_preprocessor=getattr(torch_model, 'data_preprocessor', None))
-    if not isinstance(model_inputs, torch.Tensor) and len(model_inputs) == 1:
+
+    if isinstance(model_inputs, list) and len(model_inputs) == 1:
         model_inputs = model_inputs[0]
     data_samples = data['data_samples']
     patch_metas = {'data_samples': data_samples}
-    input_metas = {'data_samples': data_samples, 'mode': 'predict'}
+    input_metas = {'data_samples': data_samples, 'mode': task_processor.mode()}
 
     # export to onnx
     context_info = dict()

mmdeploy/apis/visualize.py

@@ -68,21 +68,27 @@ def visualize_model(model_cfg: Union[str, mmengine.Config],
             model = task_processor.build_backend_model(model)
 
     model_inputs, _ = task_processor.create_input(img, input_shape)
+
     with torch.no_grad():
         result = model.test_step(model_inputs)[0]
 
     try:
         # check headless
         import tkinter
         tkinter.Tk()
+        data_samples = None
+        if 'data_samples' in model_inputs:
+            data_samples = model_inputs['data_samples']
 
         task_processor.visualize(
             image=img,
             model=model,
             result=result,
+            data_samples=data_samples,
             output_file=output_file,
             window_name=backend.value,
             show_result=show_result)
+
     except Exception as e:
         from mmdeploy.utils import get_root_logger
         logger = get_root_logger()

mmdeploy/backend/tensorrt/utils.py

@@ -43,7 +43,6 @@ def load(path: str) -> trt.ICudaEngine:
 
 def search_cuda_version() -> str:
     """try cmd to get cuda version, then try `torch.cuda`
-
     Returns:
         str: cuda version, for example 10.2
     """
@@ -189,6 +188,15 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto],
         profile.set_shape(input_name, min_shape, opt_shape, max_shape)
     config.add_optimization_profile(profile)
 
+    cuda_version = search_cuda_version()
+    if cuda_version is not None:
+        version_major = int(cuda_version.split('.')[0])
+        if version_major < 11:
+            # cu11 support cublasLt, so cudnn heuristic tactic should disable CUBLAS_LT # noqa E501
+            tactic_source = config.get_tactic_sources() - (
+                1 << int(trt.TacticSource.CUBLAS_LT))
+            config.set_tactic_sources(tactic_source)
+
     if fp16_mode:
         if version.parse(trt.__version__) < version.parse('8'):
             builder.fp16_mode = fp16_mode

mmdeploy/codebase/base/task.py

@@ -75,6 +75,7 @@ def build_data_preprocessor(self):
 
         from mmengine.registry import MODELS
         data_preprocessor = MODELS.build(preprocess_cfg)
+        data_preprocessor.to(self.device)
 
         return data_preprocessor
 
@@ -302,6 +303,14 @@ def visualize(self,
             show=show_result,
             out_file=output_file)
 
+    def mode(self) -> str:
+        """Get pytorch inference mode, it depends on codebase implementation.
+
+        Returns:
+            str: codebase inference mode
+        """
+        return 'predict'
+
     @staticmethod
     @abstractmethod
     def get_partition_cfg(partition_type: str, **kwargs) -> Dict:

mmdeploy/codebase/mmdet3d/deploy/__init__.py

@@ -1,6 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from .mmdetection3d import MMDetection3d
-from .voxel_detection import VoxelDetection
+from .voxel_detection import MMDetection3d, VoxelDetection
 from .voxel_detection_model import VoxelDetectionModel
 
 __all__ = ['MMDetection3d', 'VoxelDetection', 'VoxelDetectionModel']