diff --git a/configs/_base_/backends/tensorrt_fp16.py b/configs/_base_/backends/tensorrt_fp16.py new file mode 100644 index 000000000..1517c4096 --- /dev/null +++ b/configs/_base_/backends/tensorrt_fp16.py @@ -0,0 +1,6 @@ +import tensorrt as trt + +backend_config = dict( + type='tensorrt', + common_config=dict( + fp16_mode=True, log_level=trt.Logger.INFO, max_workspace_size=0)) diff --git a/configs/mmcls/classification_tensorrt_fp16_dynamic-224x224-224x224.py b/configs/mmcls/classification_tensorrt_fp16_dynamic-224x224-224x224.py new file mode 100644 index 000000000..66af49fae --- /dev/null +++ b/configs/mmcls/classification_tensorrt_fp16_dynamic-224x224-224x224.py @@ -0,0 +1,13 @@ +_base_ = ['./classification_dynamic.py', '../_base_/backends/tensorrt_fp16.py'] + +onnx_config = dict(input_shape=[224, 224]) +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 224, 224], + opt_shape=[4, 3, 224, 224], + max_shape=[64, 3, 224, 224]))) + ]) diff --git a/configs/mmcls/classification_tensorrt_fp16_static-224x224.py b/configs/mmcls/classification_tensorrt_fp16_static-224x224.py new file mode 100644 index 000000000..0eab7ae6a --- /dev/null +++ b/configs/mmcls/classification_tensorrt_fp16_static-224x224.py @@ -0,0 +1,13 @@ +_base_ = ['./classification_static.py', '../_base_/backends/tensorrt_fp16.py'] + +onnx_config = dict(input_shape=[224, 224]) +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 224, 224], + opt_shape=[1, 3, 224, 224], + max_shape=[1, 3, 224, 224]))) + ]) diff --git a/configs/mmdet/_base_/base_tensorrt_fp16_dynamic-320x320-1344x1344.py b/configs/mmdet/_base_/base_tensorrt_fp16_dynamic-320x320-1344x1344.py new file mode 100644 index 000000000..fb37f6e20 --- /dev/null +++ b/configs/mmdet/_base_/base_tensorrt_fp16_dynamic-320x320-1344x1344.py @@ -0,0 +1,12 @@ +_base_ = ['./base_dynamic.py', '../../_base_/backends/tensorrt_fp16.py'] + +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 320, 320], + opt_shape=[1, 3, 800, 1344], + max_shape=[1, 3, 1344, 1344]))) + ]) diff --git a/configs/mmdet/_base_/base_tensorrt_fp16_static-800x1344.py b/configs/mmdet/_base_/base_tensorrt_fp16_static-800x1344.py new file mode 100644 index 000000000..0215627cd --- /dev/null +++ b/configs/mmdet/_base_/base_tensorrt_fp16_static-800x1344.py @@ -0,0 +1,14 @@ +_base_ = ['./base_static.py', '../../_base_/backends/tensorrt_fp16.py'] + +onnx_config = dict(input_shape=(1344, 800)) + +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 800, 1344], + opt_shape=[1, 3, 800, 1344], + max_shape=[1, 3, 800, 1344]))) + ]) diff --git a/configs/mmdet/single-stage/single-stage_tensorrt_fp16_dynamic-320x320-1344x1344.py b/configs/mmdet/single-stage/single-stage_tensorrt_fp16_dynamic-320x320-1344x1344.py new file mode 100644 index 000000000..d946e8ade --- /dev/null +++ b/configs/mmdet/single-stage/single-stage_tensorrt_fp16_dynamic-320x320-1344x1344.py @@ -0,0 +1 @@ +_base_ = ['../_base_/base_tensorrt_fp16_dynamic-320x320-1344x1344.py'] diff --git a/configs/mmdet/single-stage/single-stage_tensorrt_fp16_static-800x1344.py b/configs/mmdet/single-stage/single-stage_tensorrt_fp16_static-800x1344.py new file mode 100644 index 000000000..cf7fd7707 --- /dev/null +++ b/configs/mmdet/single-stage/single-stage_tensorrt_fp16_static-800x1344.py @@ -0,0 +1 @@ +_base_ = ['../_base_/base_tensorrt_fp16_static-800x1344.py'] diff --git a/configs/mmdet/two-stage/two-stage_partition_tensorrt-fp16_dynamic-320x320-1344x1344.py b/configs/mmdet/two-stage/two-stage_partition_tensorrt-fp16_dynamic-320x320-1344x1344.py new file mode 100644 index 000000000..13802234b --- /dev/null +++ b/configs/mmdet/two-stage/two-stage_partition_tensorrt-fp16_dynamic-320x320-1344x1344.py @@ -0,0 +1,19 @@ +_base_ = ['./two-stage_tensorrt_fp16_dynamic-320x320-1344x1344.py'] + +partition_config = dict(type='two_stage', apply_marks=True) +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 320, 320], + opt_shape=[1, 3, 800, 1344], + max_shape=[1, 3, 1344, 1344]))), + dict( + input_shapes=dict( + bbox_feats=dict( + min_shape=[500, 256, 7, 7], + opt_shape=[1000, 256, 7, 7], + max_shape=[2000, 256, 7, 7]))) + ]) diff --git a/configs/mmdet/two-stage/two-stage_tensorrt_fp16_dynamic-320x320-1344x1344.py b/configs/mmdet/two-stage/two-stage_tensorrt_fp16_dynamic-320x320-1344x1344.py new file mode 100644 index 000000000..d946e8ade --- /dev/null +++ b/configs/mmdet/two-stage/two-stage_tensorrt_fp16_dynamic-320x320-1344x1344.py @@ -0,0 +1 @@ +_base_ = ['../_base_/base_tensorrt_fp16_dynamic-320x320-1344x1344.py'] diff --git a/configs/mmdet/two-stage/two-stage_tensorrt_fp16_static-800x1344.py b/configs/mmdet/two-stage/two-stage_tensorrt_fp16_static-800x1344.py new file mode 100644 index 000000000..cf7fd7707 --- /dev/null +++ b/configs/mmdet/two-stage/two-stage_tensorrt_fp16_static-800x1344.py @@ -0,0 +1 @@ +_base_ = ['../_base_/base_tensorrt_fp16_static-800x1344.py'] diff --git a/configs/mmdet/with-mask/mask_tensorrt_fp16_dynamic-320x320-1344x1344.py b/configs/mmdet/with-mask/mask_tensorrt_fp16_dynamic-320x320-1344x1344.py new file mode 100644 index 000000000..d1f06b802 --- /dev/null +++ b/configs/mmdet/with-mask/mask_tensorrt_fp16_dynamic-320x320-1344x1344.py @@ -0,0 +1,14 @@ +_base_ = [ + '../_base_/mask_base_dynamic.py', '../../_base_/backends/tensorrt_fp16.py' +] + +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 320, 320], + opt_shape=[1, 3, 800, 1344], + max_shape=[1, 3, 1344, 1344]))) + ]) diff --git a/configs/mmdet/with-mask/mask_tensorrt_fp16_static-800x1344.py b/configs/mmdet/with-mask/mask_tensorrt_fp16_static-800x1344.py new file mode 100644 index 000000000..b2d4c2da1 --- /dev/null +++ b/configs/mmdet/with-mask/mask_tensorrt_fp16_static-800x1344.py @@ -0,0 +1,15 @@ +_base_ = [ + '../_base_/mask_base_static.py', '../../_base_/backends/tensorrt_fp16.py' +] + +onnx_config = dict(input_shape=(1344, 800)) +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 800, 1344], + opt_shape=[1, 3, 800, 1344], + max_shape=[1, 3, 800, 1344]))) + ]) diff --git a/configs/mmdet/with-mask/mask_tensorrt_int8_dynamic-320x320-1344x1344.py b/configs/mmdet/with-mask/mask_tensorrt_int8_dynamic-320x320-1344x1344.py new file mode 100644 index 000000000..a121157ff --- /dev/null +++ b/configs/mmdet/with-mask/mask_tensorrt_int8_dynamic-320x320-1344x1344.py @@ -0,0 +1,14 @@ +_base_ = [ + '../_base_/mask_base_dynamic.py', '../../_base_/backends/tensorrt_int8.py' +] + +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 320, 320], + opt_shape=[1, 3, 800, 1344], + max_shape=[1, 3, 1344, 1344]))) + ]) diff --git a/configs/mmdet/with-mask/mask_tensorrt_int8_static-800x1344.py b/configs/mmdet/with-mask/mask_tensorrt_int8_static-800x1344.py new file mode 100644 index 000000000..7e6af3ca8 --- /dev/null +++ b/configs/mmdet/with-mask/mask_tensorrt_int8_static-800x1344.py @@ -0,0 +1,15 @@ +_base_ = [ + '../_base_/mask_base_static.py', '../../_base_/backends/tensorrt_int8.py' +] + +onnx_config = dict(input_shape=(1344, 800)) +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 800, 1344], + opt_shape=[1, 3, 800, 1344], + max_shape=[1, 3, 800, 1344]))) + ]) diff --git a/configs/mmedit/super-resolution/super-resolution_tensorrt_fp16_dynamic-32x32-512x512.py b/configs/mmedit/super-resolution/super-resolution_tensorrt_fp16_dynamic-32x32-512x512.py new file mode 100644 index 000000000..faa9013b5 --- /dev/null +++ b/configs/mmedit/super-resolution/super-resolution_tensorrt_fp16_dynamic-32x32-512x512.py @@ -0,0 +1,11 @@ +_base_ = [ + './super-resolution_dynamic.py', '../../_base_/backends/tensorrt_fp16.py' +] +backend_config = dict(model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 32, 32], + opt_shape=[1, 3, 256, 256], + max_shape=[1, 3, 512, 512]))) +]) diff --git a/configs/mmedit/super-resolution/super-resolution_tensorrt_fp16_static-256x256.py b/configs/mmedit/super-resolution/super-resolution_tensorrt_fp16_static-256x256.py new file mode 100644 index 000000000..d828a38e7 --- /dev/null +++ b/configs/mmedit/super-resolution/super-resolution_tensorrt_fp16_static-256x256.py @@ -0,0 +1,12 @@ +_base_ = [ + './super-resolution_static.py', '../../_base_/backends/tensorrt_fp16.py' +] +onnx_config = dict(input_shape=[256, 256]) +backend_config = dict(model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 256, 256], + opt_shape=[1, 3, 256, 256], + max_shape=[1, 3, 256, 256]))) +]) diff --git a/configs/mmedit/super-resolution/super-resolution_tensorrt_int8_dynamic-32x32-512x512.py b/configs/mmedit/super-resolution/super-resolution_tensorrt_int8_dynamic-32x32-512x512.py new file mode 100644 index 000000000..ddd786c5b --- /dev/null +++ b/configs/mmedit/super-resolution/super-resolution_tensorrt_int8_dynamic-32x32-512x512.py @@ -0,0 +1,11 @@ +_base_ = [ + './super-resolution_dynamic.py', '../../_base_/backends/tensorrt_int8.py' +] +backend_config = dict(model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 32, 32], + opt_shape=[1, 3, 256, 256], + max_shape=[1, 3, 512, 512]))) +]) diff --git a/configs/mmedit/super-resolution/super-resolution_tensorrt_int8_static-256x256.py b/configs/mmedit/super-resolution/super-resolution_tensorrt_int8_static-256x256.py new file mode 100644 index 000000000..d8220c10b --- /dev/null +++ b/configs/mmedit/super-resolution/super-resolution_tensorrt_int8_static-256x256.py @@ -0,0 +1,12 @@ +_base_ = [ + './super-resolution_static.py', '../../_base_/backends/tensorrt_int8.py' +] +onnx_config = dict(input_shape=[256, 256]) +backend_config = dict(model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 256, 256], + opt_shape=[1, 3, 256, 256], + max_shape=[1, 3, 256, 256]))) +]) diff --git a/configs/mmocr/text-detection/text-detection_tensorrt_fp16_dynamic-320x320-1024x1824.py b/configs/mmocr/text-detection/text-detection_tensorrt_fp16_dynamic-320x320-1024x1824.py new file mode 100644 index 000000000..336b9a0e3 --- /dev/null +++ b/configs/mmocr/text-detection/text-detection_tensorrt_fp16_dynamic-320x320-1024x1824.py @@ -0,0 +1,13 @@ +_base_ = [ + './text-detection_dynamic.py', '../../_base_/backends/tensorrt_fp16.py' +] +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 320, 320], + opt_shape=[1, 3, 600, 800], + max_shape=[1, 3, 1024, 1824]))) + ]) diff --git a/configs/mmocr/text-detection/text-detection_tensorrt_fp16_static-512x512.py b/configs/mmocr/text-detection/text-detection_tensorrt_fp16_static-512x512.py new file mode 100644 index 000000000..bd5f55cd7 --- /dev/null +++ b/configs/mmocr/text-detection/text-detection_tensorrt_fp16_static-512x512.py @@ -0,0 +1,15 @@ +_base_ = [ + './text-detection_static.py', '../../_base_/backends/tensorrt_fp16.py' +] + +onnx_config = dict(input_shape=[512, 512]) +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 512, 512], + opt_shape=[1, 3, 512, 512], + max_shape=[1, 3, 512, 512]))) + ]) diff --git a/configs/mmocr/text-detection/text-detection_tensorrt_int8_dynamic-320x320-1024x1824.py b/configs/mmocr/text-detection/text-detection_tensorrt_int8_dynamic-320x320-1024x1824.py new file mode 100644 index 000000000..01d16b859 --- /dev/null +++ b/configs/mmocr/text-detection/text-detection_tensorrt_int8_dynamic-320x320-1024x1824.py @@ -0,0 +1,13 @@ +_base_ = [ + './text-detection_dynamic.py', '../../_base_/backends/tensorrt_int8.py' +] +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 320, 320], + opt_shape=[1, 3, 600, 800], + max_shape=[1, 3, 1024, 1824]))) + ]) diff --git a/configs/mmocr/text-detection/text-detection_tensorrt_int8_static-512x512.py b/configs/mmocr/text-detection/text-detection_tensorrt_int8_static-512x512.py new file mode 100644 index 000000000..df5fdf3e5 --- /dev/null +++ b/configs/mmocr/text-detection/text-detection_tensorrt_int8_static-512x512.py @@ -0,0 +1,15 @@ +_base_ = [ + './text-detection_static.py', '../../_base_/backends/tensorrt_int8.py' +] + +onnx_config = dict(input_shape=[512, 512]) +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 512, 512], + opt_shape=[1, 3, 512, 512], + max_shape=[1, 3, 512, 512]))) + ]) diff --git a/configs/mmocr/text-recognition/text-recognition_tensorrt_fp16_dynamic-1x32x32-1x32x640.py b/configs/mmocr/text-recognition/text-recognition_tensorrt_fp16_dynamic-1x32x32-1x32x640.py new file mode 100644 index 000000000..6f0ad7a38 --- /dev/null +++ b/configs/mmocr/text-recognition/text-recognition_tensorrt_fp16_dynamic-1x32x32-1x32x640.py @@ -0,0 +1,13 @@ +_base_ = [ + './text-recognition_dynamic.py', '../../_base_/backends/tensorrt_fp16.py' +] +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 1, 32, 32], + opt_shape=[1, 1, 32, 64], + max_shape=[1, 1, 32, 640]))) + ]) diff --git a/configs/mmocr/text-recognition/text-recognition_tensorrt_fp16_static-1x32x32.py b/configs/mmocr/text-recognition/text-recognition_tensorrt_fp16_static-1x32x32.py new file mode 100644 index 000000000..cb483077a --- /dev/null +++ b/configs/mmocr/text-recognition/text-recognition_tensorrt_fp16_static-1x32x32.py @@ -0,0 +1,15 @@ +_base_ = [ + './text-recognition_static.py', '../../_base_/backends/tensorrt_fp16.py' +] + +onnx_config = dict(input_shape=[32, 32]) +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 1, 32, 32], + opt_shape=[1, 1, 32, 32], + max_shape=[1, 1, 32, 32]))) + ]) diff --git a/configs/mmocr/text-recognition/text-recognition_tensorrt_int8_dynamic-1x32x32-1x32x640.py b/configs/mmocr/text-recognition/text-recognition_tensorrt_int8_dynamic-1x32x32-1x32x640.py new file mode 100644 index 000000000..c1ff2ce48 --- /dev/null +++ b/configs/mmocr/text-recognition/text-recognition_tensorrt_int8_dynamic-1x32x32-1x32x640.py @@ -0,0 +1,13 @@ +_base_ = [ + './text-recognition_dynamic.py', '../../_base_/backends/tensorrt_int8.py' +] +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 1, 32, 32], + opt_shape=[1, 1, 32, 64], + max_shape=[1, 1, 32, 640]))) + ]) diff --git a/configs/mmocr/text-recognition/text-recognition_tensorrt_int8_static-1x32x32.py b/configs/mmocr/text-recognition/text-recognition_tensorrt_int8_static-1x32x32.py new file mode 100644 index 000000000..c5502d5c4 --- /dev/null +++ b/configs/mmocr/text-recognition/text-recognition_tensorrt_int8_static-1x32x32.py @@ -0,0 +1,15 @@ +_base_ = [ + './text-recognition_static.py', '../../_base_/backends/tensorrt_int8.py' +] + +onnx_config = dict(input_shape=[32, 32]) +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 1, 32, 32], + opt_shape=[1, 1, 32, 32], + max_shape=[1, 1, 32, 32]))) + ]) diff --git a/configs/mmseg/segmentation_tensorrt_fp16_dynamic-512x1024-2048x2048.py b/configs/mmseg/segmentation_tensorrt_fp16_dynamic-512x1024-2048x2048.py new file mode 100644 index 000000000..7d67e2311 --- /dev/null +++ b/configs/mmseg/segmentation_tensorrt_fp16_dynamic-512x1024-2048x2048.py @@ -0,0 +1,11 @@ +_base_ = ['./segmentation_dynamic.py', '../_base_/backends/tensorrt_fp16.py'] +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 512, 1024], + opt_shape=[1, 3, 1024, 2048], + max_shape=[1, 3, 2048, 2048]))) + ]) diff --git a/configs/mmseg/segmentation_tensorrt_fp16_static-512x1024.py b/configs/mmseg/segmentation_tensorrt_fp16_static-512x1024.py new file mode 100644 index 000000000..9fa62ab31 --- /dev/null +++ b/configs/mmseg/segmentation_tensorrt_fp16_static-512x1024.py @@ -0,0 +1,13 @@ +_base_ = ['./segmentation_static.py', '../_base_/backends/tensorrt_fp16.py'] + +onnx_config = dict(input_shape=[1024, 512]) +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 512, 1024], + opt_shape=[1, 3, 512, 1024], + max_shape=[1, 3, 512, 1024]))) + ]) diff --git a/configs/mmseg/segmentation_tensorrt_int8_dynamic-512x1024-2048x2048.py b/configs/mmseg/segmentation_tensorrt_int8_dynamic-512x1024-2048x2048.py new file mode 100644 index 000000000..db0d90eae --- /dev/null +++ b/configs/mmseg/segmentation_tensorrt_int8_dynamic-512x1024-2048x2048.py @@ -0,0 +1,11 @@ +_base_ = ['./segmentation_dynamic.py', '../_base_/backends/tensorrt_int8.py'] +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 512, 1024], + opt_shape=[1, 3, 1024, 2048], + max_shape=[1, 3, 2048, 2048]))) + ]) diff --git a/configs/mmseg/segmentation_tensorrt_int8_static-512x1024.py b/configs/mmseg/segmentation_tensorrt_int8_static-512x1024.py new file mode 100644 index 000000000..b84cb7fa6 --- /dev/null +++ b/configs/mmseg/segmentation_tensorrt_int8_static-512x1024.py @@ -0,0 +1,13 @@ +_base_ = ['./segmentation_static.py', '../_base_/backends/tensorrt_int8.py'] + +onnx_config = dict(input_shape=[1024, 512]) +backend_config = dict( + common_config=dict(max_workspace_size=1 << 30), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + min_shape=[1, 3, 512, 1024], + opt_shape=[1, 3, 512, 1024], + max_shape=[1, 3, 512, 1024]))) + ]) diff --git a/mmdeploy/apis/utils.py b/mmdeploy/apis/utils.py index 09d4618d4..cc603bd8c 100644 --- a/mmdeploy/apis/utils.py +++ b/mmdeploy/apis/utils.py @@ -359,5 +359,9 @@ def get_tensor_from_input(codebase: Codebase, input_data: Dict[str, Any]): from mmdeploy.mmocr.export import get_tensor_from_input \ as get_tensor_from_input_mmocr return get_tensor_from_input_mmocr(input_data) + elif codebase == Codebase.MMEDIT: + from mmdeploy.mmedit.export import get_tensor_from_input \ + as get_tensor_from_input_mmedit + return get_tensor_from_input_mmedit(input_data) else: raise NotImplementedError(f'Unknown codebase type: {codebase.value}') diff --git a/mmdeploy/mmdet/models/roi_heads/test_mixins.py b/mmdeploy/mmdet/models/roi_heads/test_mixins.py index 7d5c7e280..c73e0ffbc 100644 --- a/mmdeploy/mmdet/models/roi_heads/test_mixins.py +++ b/mmdeploy/mmdet/models/roi_heads/test_mixins.py @@ -45,9 +45,11 @@ def simple_test_mask_of_mask_test_mixin(ctx, self, x, img_metas, det_bboxes, bboxes_shape, labels_shape = list(det_bboxes.shape), list( det_labels.shape) bboxes_shape[1], labels_shape[1] = 1, 1 - det_bboxes = torch.tensor([[[0., 0., 1., 1., - 0.]]]).expand(*bboxes_shape) - det_labels = torch.tensor([[0]]).expand(*labels_shape) + det_bboxes = torch.tensor( + [[[0., 0., 1., 1., 0.]]], + device=det_bboxes.device).expand(*bboxes_shape) + det_labels = torch.tensor( + [[0]], device=det_bboxes.device).expand(*labels_shape) batch_size = det_bboxes.size(0) det_bboxes = det_bboxes[..., :4] diff --git a/mmdeploy/mmedit/export/__init__.py b/mmdeploy/mmedit/export/__init__.py index 7cc21c387..0ed7292c7 100644 --- a/mmdeploy/mmedit/export/__init__.py +++ b/mmdeploy/mmedit/export/__init__.py @@ -1,3 +1,7 @@ -from .prepare_input import build_dataloader, build_dataset, create_input +from .prepare_input import (build_dataloader, build_dataset, create_input, + get_tensor_from_input) -__all__ = ['create_input', 'build_dataset', 'build_dataloader'] +__all__ = [ + 'create_input', 'build_dataset', 'build_dataloader', + 'get_tensor_from_input' +] diff --git a/mmdeploy/mmedit/export/prepare_input.py b/mmdeploy/mmedit/export/prepare_input.py index c7293eca6..63efc9491 100644 --- a/mmdeploy/mmedit/export/prepare_input.py +++ b/mmdeploy/mmedit/export/prepare_input.py @@ -1,4 +1,4 @@ -from typing import Optional, Sequence, Union +from typing import Any, Dict, Optional, Sequence, Union import mmcv import numpy as np @@ -188,3 +188,14 @@ def build_dataloader(dataset: Dataset, return build_dataloader_mmedit(dataset, samples_per_gpu, workers_per_gpu, num_gpus, dist, shuffle, seed, drop_last, pin_memory, persistent_workers, **kwargs) + + +def get_tensor_from_input(input_data: Dict[str, Any]): + """Get input tensor from input data. + + Args: + input_data (dict): Input data containing meta info and image tensor. + Returns: + torch.Tensor: An image in `Tensor`. + """ + return input_data['lq'] diff --git a/mmdeploy/mmocr/export/__init__.py b/mmdeploy/mmocr/export/__init__.py index 7cc21c387..0ed7292c7 100644 --- a/mmdeploy/mmocr/export/__init__.py +++ b/mmdeploy/mmocr/export/__init__.py @@ -1,3 +1,7 @@ -from .prepare_input import build_dataloader, build_dataset, create_input +from .prepare_input import (build_dataloader, build_dataset, create_input, + get_tensor_from_input) -__all__ = ['create_input', 'build_dataset', 'build_dataloader'] +__all__ = [ + 'create_input', 'build_dataset', 'build_dataloader', + 'get_tensor_from_input' +] diff --git a/mmdeploy/mmocr/export/prepare_input.py b/mmdeploy/mmocr/export/prepare_input.py index cd02ebce4..f56236695 100644 --- a/mmdeploy/mmocr/export/prepare_input.py +++ b/mmdeploy/mmocr/export/prepare_input.py @@ -2,7 +2,7 @@ from typing import Any, Optional, Sequence, Union import mmcv import numpy as np -from mmcv.parallel import collate, scatter +from mmcv.parallel import DataContainer, collate, scatter from mmdet.datasets import replace_ImageToTensor from mmocr.datasets import build_dataloader as build_dataloader_mmocr from mmocr.datasets import build_dataset as build_dataset_mmocr @@ -189,4 +189,6 @@ def get_tensor_from_input(input_data: tuple): Returns: torch.Tensor: An image in `Tensor`. """ - return input_data['img'] + if isinstance(input_data['img'], DataContainer): + return input_data['img'].data[0] + return input_data['img'][0] diff --git a/mmdeploy/mmseg/export/prepare_input.py b/mmdeploy/mmseg/export/prepare_input.py index 8f9426721..aa32d2dc5 100644 --- a/mmdeploy/mmseg/export/prepare_input.py +++ b/mmdeploy/mmseg/export/prepare_input.py @@ -138,4 +138,4 @@ def get_tensor_from_input(input_data: tuple): Returns: torch.Tensor: An image in `Tensor`. """ - return input_data['img'] + return input_data['img'][0] diff --git a/mmdeploy/mmseg/models/segmentors/base.py b/mmdeploy/mmseg/models/segmentors/base.py index 9e10ab8e8..e420e30f2 100644 --- a/mmdeploy/mmseg/models/segmentors/base.py +++ b/mmdeploy/mmseg/models/segmentors/base.py @@ -10,7 +10,11 @@ def forward_of_base_segmentor(ctx, self, img, img_metas=None, **kwargs): """Rewrite `forward` for default backend.""" if img_metas is None: img_metas = {} - assert isinstance(img_metas, dict) + while isinstance(img_metas, list): + img_metas = img_metas[0] + + if isinstance(img, list): + img = torch.cat(img, 0) assert isinstance(img, torch.Tensor) deploy_cfg = ctx.cfg