From 6b82fd9786ff8efb601707b941c58cf38a12ba3e Mon Sep 17 00:00:00 2001 From: maningsheng Date: Thu, 17 Mar 2022 12:31:32 +0800 Subject: [PATCH 1/5] support mmseg fcn+unet dynamic shape --- docs/en/codebases/mmseg.md | 2 +- docs/en/supported_models.md | 2 +- mmdeploy/codebase/mmseg/models/__init__.py | 1 + .../codebase/mmseg/models/utils/__init__.py | 4 ++ .../mmseg/models/utils/up_conv_block.py | 37 +++++++++++++++++++ 5 files changed, 44 insertions(+), 2 deletions(-) create mode 100644 mmdeploy/codebase/mmseg/models/utils/__init__.py create mode 100644 mmdeploy/codebase/mmseg/models/utils/up_conv_block.py diff --git a/docs/en/codebases/mmseg.md b/docs/en/codebases/mmseg.md index a2b6859cac..8cb30994b0 100644 --- a/docs/en/codebases/mmseg.md +++ b/docs/en/codebases/mmseg.md @@ -15,7 +15,7 @@ Please refer to [get_started.md](https://github.com/open-mmlab/mmsegmentation/bl | DeepLabV3 | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/deeplabv3) | | DeepLabV3+ | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/deeplabv3plus) | | Fast-SCNN[*](#static_shape) | Y | Y | N | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/fastscnn) | -| UNet[*](#static_shape) | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/unet) | +| UNet | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/unet) | | ANN[*](#static_shape) | Y | Y | N | N | N | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/ann) | | APCNet | Y | Y | Y | N | N | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/apcnet) | | BiSeNetV1 | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/bisenetv1) | diff --git a/docs/en/supported_models.md b/docs/en/supported_models.md index edf51c6591..fb6c42b254 100644 --- a/docs/en/supported_models.md +++ b/docs/en/supported_models.md @@ -29,7 +29,7 @@ The table below lists the models that are guaranteed to be exportable to other b | DeepLabV3 | MMSegmentation | Y | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/deeplabv3) | | DeepLabV3+ | MMSegmentation | Y | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/deeplabv3plus) | | Fast-SCNN[*static](#note) | MMSegmentation | Y | Y | Y | N | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/fastscnn) | -| UNet[*static](#note) | MMSegmentation | Y | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/unet) | +| UNet | MMSegmentation | Y | Y | Y | Y | Y | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/unet) | | ANN[*](#note) | MMSegmentation | ? | Y | Y | N | N | N | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/ann) | | APCNet | MMSegmentation | ? | Y | Y | Y | N | N | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/apcnet) | | BiSeNetV1 | MMSegmentation | ? | Y | Y | Y | N | Y | [config](https://github.com/open-mmlab/mmsegmentation/tree/master/configs/bisenetv1) | diff --git a/mmdeploy/codebase/mmseg/models/__init__.py b/mmdeploy/codebase/mmseg/models/__init__.py index 77b260b1c3..f8c63589a9 100644 --- a/mmdeploy/codebase/mmseg/models/__init__.py +++ b/mmdeploy/codebase/mmseg/models/__init__.py @@ -1,3 +1,4 @@ # Copyright (c) OpenMMLab. All rights reserved. from .decode_heads import * # noqa: F401,F403 from .segmentors import * # noqa: F401,F403 +from .utils import * # noqa: F401,F403 diff --git a/mmdeploy/codebase/mmseg/models/utils/__init__.py b/mmdeploy/codebase/mmseg/models/utils/__init__.py new file mode 100644 index 0000000000..954eaa3487 --- /dev/null +++ b/mmdeploy/codebase/mmseg/models/utils/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from .up_conv_block import up_conv_block__forward + +__all__ = ['up_conv_block__forward'] diff --git a/mmdeploy/codebase/mmseg/models/utils/up_conv_block.py b/mmdeploy/codebase/mmseg/models/utils/up_conv_block.py new file mode 100644 index 0000000000..fab00e31eb --- /dev/null +++ b/mmdeploy/codebase/mmseg/models/utils/up_conv_block.py @@ -0,0 +1,37 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +import torch + +from mmdeploy.core import FUNCTION_REWRITER +from mmdeploy.utils import is_dynamic_shape + + +@FUNCTION_REWRITER.register_rewriter( + func_name='mmseg.models.utils.UpConvBlock.forward') +def up_conv_block__forward(ctx, self, skip, x): + """Rewrite `forward` for default backend. + + To support dynamic shape for UNet backbone, + upsample feature maps with `size` instead of `scale_factor` + + Args: + ctx (ContextCaller): The context with additional information. + self: The instance of the original class. + skip (Tensor): Skip branch feature. + x (Tensor): Input feature to be upsampled. + + Returns: + Tensor: Upsampled output feature map. + """ + if is_dynamic_shape(ctx.cfg): + # upsample with `size` instead of `scale_factor` + from mmseg.ops import Upsample + for c in self.upsample.interp_upsample: + if isinstance(c, Upsample): + c.size = skip.shape[-2:] + c.scale_factor = None + + x = self.upsample(x) + out = torch.cat([skip, x], dim=1) + out = self.conv_block(out) + return out From 54bd8804ae58b5772d28aff3c467cb6558c49e8d Mon Sep 17 00:00:00 2001 From: maningsheng Date: Fri, 18 Mar 2022 14:28:39 +0800 Subject: [PATCH 2/5] add test --- mmdeploy/utils/test.py | 8 ++- .../test_mmseg/test_mmseg_models.py | 55 ++++++++++++++++++- 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/mmdeploy/utils/test.py b/mmdeploy/utils/test.py index 5234d3e495..15fa8b4261 100644 --- a/mmdeploy/utils/test.py +++ b/mmdeploy/utils/test.py @@ -358,7 +358,11 @@ def get_onnx_model(wrapped_model: nn.Module, patched_model = patch_model( wrapped_model, cfg=deploy_cfg, backend=backend.value) flatten_model_inputs = get_flatten_inputs(model_inputs) - input_names = [k for k, v in flatten_model_inputs.items() if k != 'ctx'] + input_names = onnx_cfg.get('input_names', None) + if input_names is None: + input_names = [ + k for k, v in flatten_model_inputs.items() if k != 'ctx' + ] output_names = onnx_cfg.get('output_names', None) dynamic_axes = get_dynamic_axes(deploy_cfg, input_names) @@ -366,7 +370,7 @@ def get_onnx_model(wrapped_model: nn.Module, cfg=deploy_cfg, backend=backend.value, opset=11), torch.no_grad(): torch.onnx.export( patched_model, - tuple([v for k, v in model_inputs.items()]), + tuple([flatten_model_inputs[name] for name in input_names]), onnx_file_path, export_params=True, input_names=input_names, diff --git a/tests/test_codebase/test_mmseg/test_mmseg_models.py b/tests/test_codebase/test_mmseg/test_mmseg_models.py index a29abb898f..dfcd5b4cdb 100644 --- a/tests/test_codebase/test_mmseg/test_mmseg_models.py +++ b/tests/test_codebase/test_mmseg/test_mmseg_models.py @@ -9,7 +9,7 @@ from mmseg.models.decode_heads.decode_head import BaseDecodeHead from mmdeploy.codebase import import_codebase -from mmdeploy.utils import Backend, Codebase +from mmdeploy.utils import Backend, Codebase, Task from mmdeploy.utils.test import (WrapModel, check_backend, get_model_outputs, get_rewrite_outputs) @@ -261,3 +261,56 @@ def test_emamodule_forward(backend): model_outputs.shape) assert torch.allclose( rewrite_outputs, model_outputs, rtol=1e-03, atol=1e-05) + + +@pytest.mark.parametrize('is_dynamic_shape', [True, False]) +@pytest.mark.parametrize('backend', [Backend.ONNXRUNTIME]) +def test_upconvblock_forward(backend, is_dynamic_shape): + check_backend(backend) + from mmseg.models.backbones.unet import BasicConvBlock + from mmseg.models.utils import UpConvBlock + + head = UpConvBlock(BasicConvBlock, 16, 8, 8).eval() + dynamic_axes = { + 'x': { + 0: 'b', + 2: 'h', + 3: 'w' + }, + 'skip': { + 0: 'b', + 2: 'h', + 3: 'w' + }, + 'output': { + 0: 'b', + 2: 'h', + 3: 'w' + }, + } if is_dynamic_shape else None + deploy_cfg = mmcv.Config( + dict( + backend_config=dict(type=backend.value), + onnx_config=dict( + input_names=['skip', 'x'], + output_names=['output'], + dynamic_axes=dynamic_axes), + codebase_config=dict( + type=Codebase.MMSEG.value, task=Task.SEGMENTATION.value))) + x = torch.randn(1, 16, 16, 16) + skip = torch.randn(1, 8, 32, 32) + model_inputs = {'x': x, 'skip': skip} + with torch.no_grad(): + model_outputs = get_model_outputs(head, 'forward', model_inputs) + + wrapped_model = WrapModel(head, 'forward') + rewrite_outputs, is_backend_output = get_rewrite_outputs( + wrapped_model=wrapped_model, + model_inputs=model_inputs, + deploy_cfg=deploy_cfg) + if is_backend_output: + rewrite_outputs = rewrite_outputs[0] + rewrite_outputs = rewrite_outputs.to(model_outputs).reshape( + model_outputs.shape) + assert torch.allclose( + rewrite_outputs, model_outputs, rtol=1e-03, atol=1e-05) From 13901911d26612d6c1f1ea66efce4ac13aca3359 Mon Sep 17 00:00:00 2001 From: maningsheng Date: Fri, 18 Mar 2022 16:53:12 +0800 Subject: [PATCH 3/5] fix ci --- mmdeploy/utils/test.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/mmdeploy/utils/test.py b/mmdeploy/utils/test.py index 15fa8b4261..3b8cb864d3 100644 --- a/mmdeploy/utils/test.py +++ b/mmdeploy/utils/test.py @@ -333,9 +333,12 @@ def get_flatten_inputs( if isinstance(value, torch.Tensor): flatten_inputs[name] = value elif isinstance(value, (list, tuple)): - for i, tensor in enumerate(value): - name_i = f'{name}_{i}' - flatten_inputs[name_i] = tensor + if len(value) == 1: + flatten_inputs[name] = value[0] + else: + for i, tensor in enumerate(value): + name_i = f'{name}_{i}' + flatten_inputs[name_i] = tensor return flatten_inputs @@ -363,6 +366,7 @@ def get_onnx_model(wrapped_model: nn.Module, input_names = [ k for k, v in flatten_model_inputs.items() if k != 'ctx' ] + assert len(input_names) == 1, 'Only support one input' output_names = onnx_cfg.get('output_names', None) dynamic_axes = get_dynamic_axes(deploy_cfg, input_names) @@ -370,7 +374,7 @@ def get_onnx_model(wrapped_model: nn.Module, cfg=deploy_cfg, backend=backend.value, opset=11), torch.no_grad(): torch.onnx.export( patched_model, - tuple([flatten_model_inputs[name] for name in input_names]), + tuple([model_inputs[name] for name in input_names]), onnx_file_path, export_params=True, input_names=input_names, From a50f6a9f8298bd5989ae1ca3512a26db23e5ea2e Mon Sep 17 00:00:00 2001 From: maningsheng Date: Mon, 21 Mar 2022 16:42:50 +0800 Subject: [PATCH 4/5] fix units --- mmdeploy/utils/test.py | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/mmdeploy/utils/test.py b/mmdeploy/utils/test.py index 3b8cb864d3..6f1decaa6b 100644 --- a/mmdeploy/utils/test.py +++ b/mmdeploy/utils/test.py @@ -366,15 +366,24 @@ def get_onnx_model(wrapped_model: nn.Module, input_names = [ k for k, v in flatten_model_inputs.items() if k != 'ctx' ] - assert len(input_names) == 1, 'Only support one input' output_names = onnx_cfg.get('output_names', None) dynamic_axes = get_dynamic_axes(deploy_cfg, input_names) + class DummyModel(torch.nn.Module): + + def __init__(self): + super(DummyModel, self).__init__() + self.model = patched_model + + def forward(self, inputs: dict): + return self.model(**inputs) + + model = DummyModel().eval() + with RewriterContext( cfg=deploy_cfg, backend=backend.value, opset=11), torch.no_grad(): torch.onnx.export( - patched_model, - tuple([model_inputs[name] for name in input_names]), + model, (model_inputs, {}), onnx_file_path, export_params=True, input_names=input_names, @@ -429,8 +438,13 @@ def get_backend_outputs(ir_file_path: str, """ backend = get_backend(deploy_cfg) flatten_model_inputs = get_flatten_inputs(model_inputs) - input_names = [k for k, v in flatten_model_inputs.items() if k != 'ctx'] - output_names = get_ir_config(deploy_cfg).get('output_names', None) + ir_config = get_ir_config(deploy_cfg) + input_names = ir_config.get('input_names', None) + output_names = ir_config.get('output_names', None) + if input_names is None: + input_names = [ + k for k, v in flatten_model_inputs.items() if k != 'ctx' + ] # prepare backend model and input features if backend == Backend.TENSORRT: From fcf8fbcdaaacb34cd1aa9bcdd93a8e1d153cf5fa Mon Sep 17 00:00:00 2001 From: maningsheng Date: Thu, 24 Mar 2022 11:39:02 +0800 Subject: [PATCH 5/5] resolve comments --- mmdeploy/codebase/mmseg/models/utils/up_conv_block.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/mmdeploy/codebase/mmseg/models/utils/up_conv_block.py b/mmdeploy/codebase/mmseg/models/utils/up_conv_block.py index fab00e31eb..2ca7592851 100644 --- a/mmdeploy/codebase/mmseg/models/utils/up_conv_block.py +++ b/mmdeploy/codebase/mmseg/models/utils/up_conv_block.py @@ -23,7 +23,10 @@ def up_conv_block__forward(ctx, self, skip, x): Returns: Tensor: Upsampled output feature map. """ - if is_dynamic_shape(ctx.cfg): + from mmcv.cnn import ConvModule + + # only valid when self.upsample is from build_upsample_layer + if is_dynamic_shape(ctx.cfg) and not isinstance(self.upsample, ConvModule): # upsample with `size` instead of `scale_factor` from mmseg.ops import Upsample for c in self.upsample.interp_upsample: