| |
| import torch.nn as nn |
| from mmcv.cnn import build_norm_layer |
|
|
| from mmseg.registry import MODELS |
|
|
|
|
| @MODELS.register_module() |
| class Feature2Pyramid(nn.Module): |
| """Feature2Pyramid. |
| |
| A neck structure connect ViT backbone and decoder_heads. |
| |
| Args: |
| embed_dims (int): Embedding dimension. |
| rescales (list[float]): Different sampling multiples were |
| used to obtain pyramid features. Default: [4, 2, 1, 0.5]. |
| norm_cfg (dict): Config dict for normalization layer. |
| Default: dict(type='SyncBN', requires_grad=True). |
| """ |
|
|
| def __init__(self, |
| embed_dim, |
| rescales=[4, 2, 1, 0.5], |
| norm_cfg=dict(type='SyncBN', requires_grad=True)): |
| super().__init__() |
| self.rescales = rescales |
| self.upsample_4x = None |
| for k in self.rescales: |
| if k == 4: |
| self.upsample_4x = nn.Sequential( |
| nn.ConvTranspose2d( |
| embed_dim, embed_dim, kernel_size=2, stride=2), |
| build_norm_layer(norm_cfg, embed_dim)[1], |
| nn.GELU(), |
| nn.ConvTranspose2d( |
| embed_dim, embed_dim, kernel_size=2, stride=2), |
| ) |
| elif k == 2: |
| self.upsample_2x = nn.Sequential( |
| nn.ConvTranspose2d( |
| embed_dim, embed_dim, kernel_size=2, stride=2)) |
| elif k == 1: |
| self.identity = nn.Identity() |
| elif k == 0.5: |
| self.downsample_2x = nn.MaxPool2d(kernel_size=2, stride=2) |
| elif k == 0.25: |
| self.downsample_4x = nn.MaxPool2d(kernel_size=4, stride=4) |
| else: |
| raise KeyError(f'invalid {k} for feature2pyramid') |
|
|
| def forward(self, inputs): |
| assert len(inputs) == len(self.rescales) |
| outputs = [] |
| if self.upsample_4x is not None: |
| ops = [ |
| self.upsample_4x, self.upsample_2x, self.identity, |
| self.downsample_2x |
| ] |
| else: |
| ops = [ |
| self.upsample_2x, self.identity, self.downsample_2x, |
| self.downsample_4x |
| ] |
| for i in range(len(inputs)): |
| outputs.append(ops[i](inputs[i])) |
| return tuple(outputs) |
|
|
