[Feature] Segformer backbone re-implementation #594

clownrat6 · 2021-06-09T04:54:36Z

Add backbone MixVisionTransformer;

~~- [ ] Add head SegFormerHead;~~

~~- [ ] Add dataset transform pipeline AlignedResize;~~

~~- [ ] Add some config for segformer;~~

codecov · 2021-06-09T18:11:39Z

Codecov Report

Merging #594 (e7039ff) into master (e610ed1) will increase coverage by 0.09%.
The diff coverage is 71.42%.

@@            Coverage Diff             @@
##           master     #594      +/-   ##
==========================================
+ Coverage   85.18%   85.28%   +0.09%     
==========================================
  Files         105      107       +2     
  Lines        5671     5817     +146     
  Branches      923      951      +28     
==========================================
+ Hits         4831     4961     +130     
- Misses        662      673      +11     
- Partials      178      183       +5

Flag	Coverage Δ
unittests	`85.26% <71.42%> (+0.09%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
mmseg/models/backbones/swin.py	`86.89% <ø> (ø)`
mmseg/models/backbones/vit.py	`84.84% <ø> (ø)`
mmseg/models/utils/ckpt_convert.py	`4.42% <4.65%> (+0.13%)`	⬆️
mmseg/models/backbones/mit.py	`88.80% <88.80%> (ø)`
mmseg/models/backbones/__init__.py	`100.00% <100.00%> (ø)`
mmseg/models/utils/__init__.py	`100.00% <100.00%> (ø)`
mmseg/models/utils/embed.py	`81.57% <100.00%> (+1.02%)`	⬆️
mmseg/models/utils/shape_convert.py	`100.00% <100.00%> (ø)`
... and 1 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update e610ed1...e7039ff. Read the comment docs.

2. Remove rebundant functions and unit tests;

Junjun2016 · 2021-06-20T12:50:09Z

mmseg/models/backbones/mit.py

+        self.dwconv = DWConv(feedforward_channels)
+        self.act = build_activation_layer(act_cfg)
+        self.fc2 = Linear(feedforward_channels, in_channels)
+        self.drop = nn.Dropout(drop_rate)


We can replace FC with 1x1 conv, so we can avoid dimension transform and integrate depthwise conv in MLP.

Junjun2016 · 2021-06-20T12:56:00Z

mmseg/models/backbones/mit.py

+        return x
+
+
+class Mlp(BaseModule):


Be careful that the variable or class name needs to be consistent with the paper.
Rename Mlp to MixFFN

Junjun2016 · 2021-06-20T13:01:55Z

mmseg/models/backbones/mit.py

+        return x
+
+
+class Attention(BaseModule):


Rename Attention to EfficientMultiheadAttention.
We can also inherit from MultiheadAttention in MMCV and pass different query, key, and value in forward phrase according to different spatial reduction.

Junjun2016 · 2021-06-20T13:04:31Z

mmseg/models/backbones/mit.py

+        self.sr_ratio = sr_ratio
+        if sr_ratio > 1:
+            self.sr = ConvModule(
+                in_channels=dim,
+                out_channels=dim,
+                kernel_size=sr_ratio,
+                stride=sr_ratio)
+            _, self.norm = build_norm_layer(norm_cfg, dim)


If we inherit from MultiheadAttention in MMCV, we only need to add these lines.

Junjun2016 · 2021-06-20T13:11:10Z

mmseg/models/backbones/mit.py

+        return x
+
+
+class Block(BaseModule):


Junjun2016 · 2021-06-20T13:42:00Z

mmseg/models/backbones/mit.py

+            proj_drop=drop_rate,
+            sr_ratio=sr_ratio)
+        # NOTE: drop path for stochastic depth, we shall see if this is better
+        # than dropout here


Is this comment necessary?

Junjun2016 · 2021-06-20T13:54:59Z

mmseg/models/backbones/mit.py

+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+
+        self.img_size = img_size
+        self.patch_size = patch_size
+        num_rows, num_cols = img_size[0] // patch_size[0], img_size[
+            1] // patch_size[1]
+        self.num_patches = num_rows * num_cols


Are these lines necessary?

Junjun2016 · 2021-06-20T13:55:30Z

mmseg/models/backbones/mit.py

+        num_rows, num_cols = img_size[0] // patch_size[0], img_size[
+            1] // patch_size[1]
+        self.num_patches = num_rows * num_cols
+        self.proj = nn.Conv2d(


Use ConvModule

Junjun2016 · 2021-06-20T14:30:41Z

mmseg/models/backbones/mit.py

+        self.pretrained = pretrained
+        self.depths = depths
+        # patch_embed
+        self.patch_embed1 = OverlapPatchEmbed(


Use ModuleList

mmseg/models/backbones/mit.py

Junjun2016 · 2021-06-20T14:46:21Z

mmseg/models/backbones/mit.py

+                strict=False,
+                logger=logger)
+
+    def reset_drop_path(self, drop_path_rate):


Where is this used?

2. Modify module name;

2. Add some unit tests for MixVisionTransformer;

xvjiarui · 2021-07-14T01:41:48Z

mmseg/models/utils/shape_convert.py

@@ -0,0 +1,10 @@
+def nlc_to_nchw(tensor, H, W):


xvjiarui · 2021-07-14T01:43:31Z

mmseg/models/backbones/mit.py

+    A PyTorch implement of : `SegFormer: Simple and Efficient Design for
+    Semantic Segmentation with Transformers` -
+        https://arxiv.org/pdf/2105.15203.pdf
+
+    in_channels (int): Number of input channels. Default: 3.
+    embed_dims (int): Embedding dimension. Default: 768.
+    num_stags (int): The num of stages. Default: 4.
+    num_layers (list[int]): The layer number of each transformer encode
+        layer. Default: [3, 4, 6, 3].
+    num_heads (list[int]): The attention heads of each transformer
+        encode layer. Default: [1, 2, 4, 8].
+    patch_sizes (list[int]): The patch_size of each overlapped patch embedding.
+        Default: [7, 3, 3, 3].
+    strides (list[int]): The stride of each overlapped patch embedding.
+        Default: [4, 2, 2, 2].
+    sr_ratios (list[int]): The spatial reduction rate of each transformer
+        encode layer. Default: [8, 4, 2, 1].
+    out_indices (list[int] | tuple[int] | int): Output from which stages.
+            Default: (0, 1, 2, 3).
+    mlp_ratio (int): ratio of mlp hidden dim to embedding dim.
+        Default: 4.
+    out_indices (list | tuple | int): Output from which stages.
+        Default: -1.
+    qkv_bias (bool): Enable bias for qkv if True. Default: True.
+    drop_rate (float): Probability of an element to be zeroed.
+        Default 0.0
+    attn_drop_rate (float): The drop out rate for attention layer.
+        Default 0.0
+    drop_path_rate (float): stochastic depth rate. Default 0.0
+    norm_cfg (dict): Config dict for normalization layer.
+        Default: dict(type='LN')
+    act_cfg (dict): The activation config for FFNs.
+        Defalut: dict(type='GELU').
+    pretrain_style (str): Choose to use official or mmcls pretrain weights.
+        Default: official.
+    pretrained (str, optional): model pretrained path. Default: None.
+    init_cfg (dict or list[dict], optional): Initialization config dict.
+        Default: None.
+    """


The format of the docstring is not correct.

xvjiarui · 2021-07-14T01:44:36Z

mmseg/models/backbones/mit.py

+            dropout_layer=dict(type='DropPath', drop_prob=drop_path_rate),
+            act_cfg=act_cfg)
+
+    def forward(self, x, H, W):


Suggested change

def forward(self, x, H, W):

def forward(self, x, hw_shape):

xvjiarui · 2021-07-14T01:47:11Z

mmseg/models/backbones/mit.py

+        conv1x1 = partial(
+            ConvModule,
+            kernel_size=1,
+            stride=1,
+            bias=True,
+            norm_cfg=None,
+            act_cfg=None)


No need to use partial.

xvjiarui · 2021-07-14T01:47:41Z

mmseg/models/backbones/mit.py

+from ..utils import PatchEmbed, mit_convert, nchw_to_nlc, nlc_to_nchw
+
+
+class PEConv(BaseModule):


Why do we need to use PEConv to wrap conv?

xvjiarui · 2021-07-14T01:54:03Z

mmseg/models/backbones/mit.py

+                 ffn_drop=0.,
+                 pe_index=1,
+                 dropout_layer=None,
+                 add_identity=True,


Is add_identity argument necessary?

xvjiarui · 2021-07-14T01:58:18Z

mmseg/models/backbones/mit.py

+        # first position of MixFFN
+        if pe_index == 0:
+            layers.append(PEConv(in_channels))
+        for idx in range(num_fcs - 1):
+            container = []
+            container.append(
+                conv1x1(
+                    in_channels=in_channels,
+                    out_channels=feedforward_channels))
+            # middle position of MixFFN
+            if pe_index == idx + 1:
+                container.append(PEConv(feedforward_channels))
+            container.append(self.activate)
+            container.append(nn.Dropout(ffn_drop))
+            layers.append(Sequential(*container))
+        layers.append(
+            conv1x1(
+                in_channels=feedforward_channels, out_channels=in_channels))
+        # Last position of MixFFN
+        if pe_index == num_fcs:
+            layers.append(PEConv(feedforward_channels))
+        layers.append(nn.Dropout(ffn_drop))
+        self.layers = Sequential(*layers)
+        self.dropout_layer = build_dropout(
+            dropout_layer) if dropout_layer else torch.nn.Identity()
+        self.add_identity = add_identity


This part is too long.
Consider simplifying it.

xvjiarui · 2021-07-14T01:58:50Z

mmseg/models/backbones/mit.py

+                 num_fcs=2,
+                 act_cfg=dict(type='GELU'),
+                 ffn_drop=0.,
+                 pe_index=1,


If it is too complicated, we may remove pe_index.

Just fix thenum_fcs to 2, and insert PE in the middle.

2. Simplify MixFFN; 3. Modify H, W to hw_shape;

xvjiarui · 2021-07-15T03:50:21Z

mmseg/models/utils/shape_convert.py

+        hw_shape (Sequence[int]): The height and width of output feature map.
+    """


xvjiarui · 2021-07-15T03:50:27Z

mmseg/models/utils/shape_convert.py

+    """Flatten [N, C, H, W] shape tensor to [N, L, C] shape tensor.
+
+    Args:
+        x (Tensor): The input tensor for convertion.


xvjiarui · 2021-07-15T03:52:16Z

mmseg/models/utils/embed.py

-                      (0, self.patch_size[1] - W % self.patch_size[1], 0, 0))
+
+        # TODO: Process overlapping op
+        if not self.overlapping:


Overlapping is not precise for this.
We may consider something like auto_pad or pad_to_patch_size.

We also need to make it an argument and add a docstring for that.

…gformer

mmseg/models/utils/embed.py

xvjiarui · 2021-07-19T04:09:43Z

mmseg/models/utils/embed.py

+        pad_to_patch_size (bool, optional): Whether to pad feature map shape
+            to multiple patch size. Default: False.


We may make it True by default.

* [Feature]Segformer re-implementation * Using act_cfg and norm_cfg to control activation and normalization * Split this PR into several little PRs * Fix lint error * Remove SegFormerHead * parameters init refactor * 1. Refactor segformer backbone parameters init; 2. Remove rebundant functions and unit tests; * Remove rebundant codes * 1. Remove rebundant codes; 2. Modify module name; * Refactor the backbone of segformer using mmcv.cnn.bricks.transformer.py * Fix some code logic bugs. * Add mit_convert.py to match pretrain keys of segformer. * Resolve some comments. * 1. Add some assert to ensure right params; 2. Support flexible peconv position; * Add pe_index assert and fix unit test. * 1. Add doc string for MixVisionTransformer; 2. Add some unit tests for MixVisionTransformer; * Use hw_shape to pass shape of feature map. * 1. Fix doc string of MixVisionTransformer; 2. Simplify MixFFN; 3. Modify H, W to hw_shape; * Add more unit tests. * Add doc string for shape convertion functions. * Add some unit tests to improve code coverage. * Fix Segformer backbone pretrain weights match bug. * resolve the shape convertion functions doc string. * Add pad_to_patch_size arg. * Modify default value of pad_to_patch_size arg.

* [Flax] Fix unet and ddim scheduler * correct * finish

* Update README_cn.md * Update README_cn.md

[Feature]Segformer re-implementation

b56ceb2

clownrat6 requested review from xvjiarui and Junjun2016 June 9, 2021 04:55

Using act_cfg and norm_cfg to control activation and normalization

5a54664

clownrat6 changed the title ~~[Feature]Segformer re-implementation~~ [Feature]Segformer backbone re-implementation Jun 10, 2021

sennnnn added 3 commits June 10, 2021 18:15

Split this PR into several little PRs

0d60aeb

Fix lint error

c2b0ffd

Remove SegFormerHead

c7e676e

AncientRemember approved these changes Jun 10, 2021

View reviewed changes

Junjun2016 changed the title ~~[Feature]Segformer backbone re-implementation~~ [Feature] Segformer backbone re-implementation Jun 18, 2021

sennnnn added 5 commits June 19, 2021 03:46

Merge Master

237bc9e

parameters init refactor

9edbddd

1. Refactor segformer backbone parameters init;

99d7308

2. Remove rebundant functions and unit tests;

Merge branch 'master' into segformer

02b09bf

Remove rebundant codes

3373827

Junjun2016 reviewed Jun 20, 2021

View reviewed changes

mmseg/models/backbones/mit.py Outdated

return x

class Block(BaseModule):

Copy link

Collaborator

Junjun2016 Jun 20, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rename it.

Junjun2016 reviewed Jun 20, 2021

View reviewed changes

mmseg/models/backbones/mit.py Show resolved Hide resolved

Junjun2016 reviewed Jun 20, 2021

View reviewed changes

sennnnn added 3 commits June 21, 2021 01:06

1. Remove rebundant codes;

d967a9f

2. Modify module name;

Refactor the backbone of segformer using mmcv.cnn.bricks.transformer.py

47326de

Fix some code logic bugs.

589af69

sennnnn added 2 commits July 10, 2021 01:36

Add pe_index assert and fix unit test.

e137df3

1. Add doc string for MixVisionTransformer;

b4ae0b5

2. Add some unit tests for MixVisionTransformer;

xvjiarui reviewed Jul 14, 2021

View reviewed changes

mmseg/models/utils/shape_convert.py Outdated

@@ -0,0 +1,10 @@

def nlc_to_nchw(tensor, H, W):

Copy link

Collaborator

xvjiarui Jul 14, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Docstring.

xvjiarui reviewed Jul 14, 2021

View reviewed changes

sennnnn added 5 commits July 14, 2021 11:06

Use hw_shape to pass shape of feature map.

38f7bbe

1. Fix doc string of MixVisionTransformer;

443438f

2. Simplify MixFFN; 3. Modify H, W to hw_shape;

Add more unit tests.

1affeb6

Add doc string for shape convertion functions.

e969496

Add some unit tests to improve code coverage.

582939b

xvjiarui reviewed Jul 15, 2021

View reviewed changes

sennnnn added 3 commits July 15, 2021 12:55

Fix Segformer backbone pretrain weights match bug.

e5acfe8

resolve the shape convertion functions doc string.

71df617

Merge branch 'master' of github.com:open-mmlab/mmsegmentation into se…

5ac60d4

…gformer

xvjiarui reviewed Jul 16, 2021

View reviewed changes

mmseg/models/utils/embed.py Show resolved Hide resolved

Add pad_to_patch_size arg.

38b2db8

xvjiarui reviewed Jul 19, 2021

View reviewed changes

Modify default value of pad_to_patch_size arg.

e7039ff

xvjiarui approved these changes Jul 19, 2021

View reviewed changes

xvjiarui merged commit 4d34581 into open-mmlab:master Jul 19, 2021

aravind-h-v pushed a commit to aravind-h-v/mmsegmentation that referenced this pull request Mar 27, 2023

[Flax] Fix unet and ddim scheduler (open-mmlab#594)

2345481

* [Flax] Fix unet and ddim scheduler * correct * finish

sibozhang pushed a commit to sibozhang/mmsegmentation that referenced this pull request Mar 22, 2024

Update README_cn.md (open-mmlab#594)

e409809

* Update README_cn.md * Update README_cn.md

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Feature] Segformer backbone re-implementation #594

[Feature] Segformer backbone re-implementation #594

clownrat6 commented Jun 9, 2021 •

edited

Loading

codecov bot commented Jun 9, 2021 •

edited

Loading

Junjun2016 Jun 20, 2021 •

edited

Loading

Junjun2016 Jun 20, 2021

Junjun2016 Jun 20, 2021

Junjun2016 Jun 20, 2021

Junjun2016 Jun 20, 2021

Junjun2016 Jun 20, 2021

Junjun2016 Jun 20, 2021

Junjun2016 Jun 20, 2021

Junjun2016 Jun 20, 2021

Junjun2016 Jun 20, 2021

xvjiarui Jul 14, 2021

xvjiarui Jul 14, 2021

xvjiarui Jul 14, 2021

xvjiarui Jul 14, 2021

xvjiarui Jul 14, 2021 •

edited

Loading

xvjiarui Jul 14, 2021 •

edited

Loading

xvjiarui Jul 14, 2021

xvjiarui Jul 14, 2021

xvjiarui Jul 14, 2021

xvjiarui Jul 15, 2021

xvjiarui Jul 15, 2021

xvjiarui Jul 15, 2021

xvjiarui Jul 19, 2021

		from ..utils import PatchEmbed, mit_convert, nchw_to_nlc, nlc_to_nchw


		class PEConv(BaseModule):

		hw_shape (Sequence[int]): The height and width of output feature map.
		"""

		pad_to_patch_size (bool, optional): Whether to pad feature map shape
		to multiple patch size. Default: False.

[Feature] Segformer backbone re-implementation #594

[Feature] Segformer backbone re-implementation #594

Conversation

clownrat6 commented Jun 9, 2021 • edited Loading

codecov bot commented Jun 9, 2021 • edited Loading

Codecov Report

Junjun2016 Jun 20, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

xvjiarui Jul 14, 2021 • edited Loading

Choose a reason for hiding this comment

xvjiarui Jul 14, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

clownrat6 commented Jun 9, 2021 •

edited

Loading

codecov bot commented Jun 9, 2021 •

edited

Loading

Junjun2016 Jun 20, 2021 •

edited

Loading

xvjiarui Jul 14, 2021 •

edited

Loading

xvjiarui Jul 14, 2021 •

edited

Loading