Shortcuts

mmpretrain.models.backbones.mobilevit 源代码

# Copyright (c) OpenMMLab. All rights reserved.
import math
from typing import Callable, Optional, Sequence

import torch
import torch.nn.functional as F
from mmcv.cnn import ConvModule, build_norm_layer
from torch import nn

from mmpretrain.registry import MODELS
from .base_backbone import BaseBackbone
from .mobilenet_v2 import InvertedResidual
from .vision_transformer import TransformerEncoderLayer


class MobileVitBlock(nn.Module):
    """MobileViT block.

    According to the paper, the MobileViT block has a local representation.
    a transformer-as-convolution layer which consists of a global
    representation with unfolding and folding, and a final fusion layer.

    Args:
        in_channels (int): Number of input image channels.
        transformer_dim (int): Number of transformer channels.
        ffn_dim (int): Number of ffn channels in transformer block.
        out_channels (int): Number of channels in output.
        conv_ksize (int): Conv kernel size in local representation
            and fusion. Defaults to 3.
        conv_cfg (dict, optional): Config dict for convolution layer.
            Defaults to None, which means using conv2d.
        norm_cfg (dict, optional): Config dict for normalization layer.
            Defaults to dict(type='BN').
        act_cfg (dict, optional): Config dict for activation layer.
            Defaults to dict(type='Swish').
        num_transformer_blocks (int): Number of transformer blocks in
            a MobileViT block. Defaults to 2.
        patch_size (int): Patch size for unfolding and folding.
             Defaults to 2.
        num_heads (int): Number of heads in global representation.
             Defaults to 4.
        drop_rate (float): Probability of an element to be zeroed
            after the feed forward layer. Defaults to 0.
        attn_drop_rate (float): The drop out rate for attention output weights.
            Defaults to 0.
        drop_path_rate (float): Stochastic depth rate. Defaults to 0.
        no_fusion (bool): Whether to remove the fusion layer.
            Defaults to False.
        transformer_norm_cfg (dict, optional): Config dict for normalization
            layer in transformer. Defaults to dict(type='LN').
    """

    def __init__(
            self,
            in_channels: int,
            transformer_dim: int,
            ffn_dim: int,
            out_channels: int,
            conv_ksize: int = 3,
            conv_cfg: Optional[dict] = None,
            norm_cfg: Optional[dict] = dict(type='BN'),
            act_cfg: Optional[dict] = dict(type='Swish'),
            num_transformer_blocks: int = 2,
            patch_size: int = 2,
            num_heads: int = 4,
            drop_rate: float = 0.,
            attn_drop_rate: float = 0.,
            drop_path_rate: float = 0.,
            no_fusion: bool = False,
            transformer_norm_cfg: Callable = dict(type='LN'),
    ):
        super(MobileVitBlock, self).__init__()

        self.local_rep = nn.Sequential(
            ConvModule(
                in_channels=in_channels,
                out_channels=in_channels,
                kernel_size=conv_ksize,
                padding=int((conv_ksize - 1) / 2),
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=act_cfg),
            ConvModule(
                in_channels=in_channels,
                out_channels=transformer_dim,
                kernel_size=1,
                bias=False,
                conv_cfg=conv_cfg,
                norm_cfg=None,
                act_cfg=None),
        )

        global_rep = [
            TransformerEncoderLayer(
                embed_dims=transformer_dim,
                num_heads=num_heads,
                feedforward_channels=ffn_dim,
                drop_rate=drop_rate,
                attn_drop_rate=attn_drop_rate,
                drop_path_rate=drop_path_rate,
                qkv_bias=True,
                act_cfg=dict(type='Swish'),
                norm_cfg=transformer_norm_cfg)
            for _ in range(num_transformer_blocks)
        ]
        global_rep.append(
            build_norm_layer(transformer_norm_cfg, transformer_dim)[1])
        self.global_rep = nn.Sequential(*global_rep)

        self.conv_proj = ConvModule(
            in_channels=transformer_dim,
            out_channels=out_channels,
            kernel_size=1,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)

        if no_fusion:
            self.conv_fusion = None
        else:
            self.conv_fusion = ConvModule(
                in_channels=in_channels + out_channels,
                out_channels=out_channels,
                kernel_size=conv_ksize,
                padding=int((conv_ksize - 1) / 2),
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=act_cfg)

        self.patch_size = (patch_size, patch_size)
        self.patch_area = self.patch_size[0] * self.patch_size[1]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shortcut = x

        # Local representation
        x = self.local_rep(x)

        # Unfold (feature map -> patches)
        patch_h, patch_w = self.patch_size
        B, C, H, W = x.shape
        new_h, new_w = math.ceil(H / patch_h) * patch_h, math.ceil(
            W / patch_w) * patch_w
        num_patch_h, num_patch_w = new_h // patch_h, new_w // patch_w  # n_h, n_w # noqa
        num_patches = num_patch_h * num_patch_w  # N
        interpolate = False
        if new_h != H or new_w != W:
            # Note: Padding can be done, but then it needs to be handled in attention function. # noqa
            x = F.interpolate(
                x, size=(new_h, new_w), mode='bilinear', align_corners=False)
            interpolate = True

        # [B, C, H, W] --> [B * C * n_h, n_w, p_h, p_w]
        x = x.reshape(B * C * num_patch_h, patch_h, num_patch_w,
                      patch_w).transpose(1, 2)
        # [B * C * n_h, n_w, p_h, p_w] --> [BP, N, C] where P = p_h * p_w and N = n_h * n_w # noqa
        x = x.reshape(B, C, num_patches,
                      self.patch_area).transpose(1, 3).reshape(
                          B * self.patch_area, num_patches, -1)

        # Global representations
        x = self.global_rep(x)

        # Fold (patch -> feature map)
        # [B, P, N, C] --> [B*C*n_h, n_w, p_h, p_w]
        x = x.contiguous().view(B, self.patch_area, num_patches, -1)
        x = x.transpose(1, 3).reshape(B * C * num_patch_h, num_patch_w,
                                      patch_h, patch_w)
        # [B*C*n_h, n_w, p_h, p_w] --> [B*C*n_h, p_h, n_w, p_w] --> [B, C, H, W] # noqa
        x = x.transpose(1, 2).reshape(B, C, num_patch_h * patch_h,
                                      num_patch_w * patch_w)
        if interpolate:
            x = F.interpolate(
                x, size=(H, W), mode='bilinear', align_corners=False)

        x = self.conv_proj(x)
        if self.conv_fusion is not None:
            x = self.conv_fusion(torch.cat((shortcut, x), dim=1))
        return x


[文档]@MODELS.register_module() class MobileViT(BaseBackbone): """MobileViT backbone. A PyTorch implementation of : `MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer <https://arxiv.org/pdf/2110.02178.pdf>`_ Modified from the `official repo <https://github.com/apple/ml-cvnets/blob/main/cvnets/models/classification/mobilevit.py>`_ and `timm <https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/mobilevit.py>`_. Args: arch (str | List[list]): Architecture of MobileViT. - If a string, choose from "small", "x_small" and "xx_small". - If a list, every item should be also a list, and the first item of the sub-list can be chosen from "moblienetv2" and "mobilevit", which indicates the type of this layer sequence. If "mobilenetv2", the other items are the arguments of :attr:`~MobileViT.make_mobilenetv2_layer` (except ``in_channels``) and if "mobilevit", the other items are the arguments of :attr:`~MobileViT.make_mobilevit_layer` (except ``in_channels``). Defaults to "small". in_channels (int): Number of input image channels. Defaults to 3. stem_channels (int): Channels of stem layer. Defaults to 16. last_exp_factor (int): Channels expand factor of last layer. Defaults to 4. out_indices (Sequence[int]): Output from which stages. Defaults to (4, ). frozen_stages (int): Stages to be frozen (all param fixed). Defaults to -1, which means not freezing any parameters. conv_cfg (dict, optional): Config dict for convolution layer. Defaults to None, which means using conv2d. norm_cfg (dict, optional): Config dict for normalization layer. Defaults to dict(type='BN'). act_cfg (dict, optional): Config dict for activation layer. Defaults to dict(type='Swish'). init_cfg (dict, optional): Initialization config dict. """ # noqa # Parameters to build layers. The first param is the type of layer. # For `mobilenetv2` layer, the rest params from left to right are: # out channels, stride, num of blocks, expand_ratio. # For `mobilevit` layer, the rest params from left to right are: # out channels, stride, transformer_channels, ffn channels, # num of transformer blocks, expand_ratio. arch_settings = { 'small': [ ['mobilenetv2', 32, 1, 1, 4], ['mobilenetv2', 64, 2, 3, 4], ['mobilevit', 96, 2, 144, 288, 2, 4], ['mobilevit', 128, 2, 192, 384, 4, 4], ['mobilevit', 160, 2, 240, 480, 3, 4], ], 'x_small': [ ['mobilenetv2', 32, 1, 1, 4], ['mobilenetv2', 48, 2, 3, 4], ['mobilevit', 64, 2, 96, 192, 2, 4], ['mobilevit', 80, 2, 120, 240, 4, 4], ['mobilevit', 96, 2, 144, 288, 3, 4], ], 'xx_small': [ ['mobilenetv2', 16, 1, 1, 2], ['mobilenetv2', 24, 2, 3, 2], ['mobilevit', 48, 2, 64, 128, 2, 2], ['mobilevit', 64, 2, 80, 160, 4, 2], ['mobilevit', 80, 2, 96, 192, 3, 2], ] } def __init__(self, arch='small', in_channels=3, stem_channels=16, last_exp_factor=4, out_indices=(4, ), frozen_stages=-1, conv_cfg=None, norm_cfg=dict(type='BN'), act_cfg=dict(type='Swish'), init_cfg=[ dict(type='Kaiming', layer=['Conv2d']), dict( type='Constant', val=1, layer=['_BatchNorm', 'GroupNorm']) ]): super(MobileViT, self).__init__(init_cfg) if isinstance(arch, str): arch = arch.lower() assert arch in self.arch_settings, \ f'Unavailable arch, please choose from ' \ f'({set(self.arch_settings)}) or pass a list.' arch = self.arch_settings[arch] self.arch = arch self.num_stages = len(arch) # check out indices and frozen stages if isinstance(out_indices, int): out_indices = [out_indices] assert isinstance(out_indices, Sequence), \ f'"out_indices" must by a sequence or int, ' \ f'get {type(out_indices)} instead.' for i, index in enumerate(out_indices): if index < 0: out_indices[i] = self.num_stages + index assert out_indices[i] >= 0, f'Invalid out_indices {index}' self.out_indices = out_indices if frozen_stages not in range(-1, self.num_stages): raise ValueError('frozen_stages must be in range(-1, ' f'{self.num_stages}). ' f'But received {frozen_stages}') self.frozen_stages = frozen_stages _make_layer_func = { 'mobilenetv2': self.make_mobilenetv2_layer, 'mobilevit': self.make_mobilevit_layer, } self.stem = ConvModule( in_channels=in_channels, out_channels=stem_channels, kernel_size=3, stride=2, padding=1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg) in_channels = stem_channels layers = [] for i, layer_settings in enumerate(arch): layer_type, settings = layer_settings[0], layer_settings[1:] layer, out_channels = _make_layer_func[layer_type](in_channels, *settings) layers.append(layer) in_channels = out_channels self.layers = nn.Sequential(*layers) self.conv_1x1_exp = ConvModule( in_channels=in_channels, out_channels=last_exp_factor * in_channels, kernel_size=1, stride=1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg)
[文档] @staticmethod def make_mobilevit_layer(in_channels, out_channels, stride, transformer_dim, ffn_dim, num_transformer_blocks, expand_ratio=4): """Build mobilevit layer, which consists of one InvertedResidual and one MobileVitBlock. Args: in_channels (int): The input channels. out_channels (int): The output channels. stride (int): The stride of the first 3x3 convolution in the ``InvertedResidual`` layers. transformer_dim (int): The channels of the transformer layers. ffn_dim (int): The mid-channels of the feedforward network in transformer layers. num_transformer_blocks (int): The number of transformer blocks. expand_ratio (int): adjusts number of channels of the hidden layer in ``InvertedResidual`` by this amount. Defaults to 4. """ layer = [] layer.append( InvertedResidual( in_channels=in_channels, out_channels=out_channels, stride=stride, expand_ratio=expand_ratio, act_cfg=dict(type='Swish'), )) layer.append( MobileVitBlock( in_channels=out_channels, transformer_dim=transformer_dim, ffn_dim=ffn_dim, out_channels=out_channels, num_transformer_blocks=num_transformer_blocks, )) return nn.Sequential(*layer), out_channels
[文档] @staticmethod def make_mobilenetv2_layer(in_channels, out_channels, stride, num_blocks, expand_ratio=4): """Build mobilenetv2 layer, which consists of several InvertedResidual layers. Args: in_channels (int): The input channels. out_channels (int): The output channels. stride (int): The stride of the first 3x3 convolution in the ``InvertedResidual`` layers. num_blocks (int): The number of ``InvertedResidual`` blocks. expand_ratio (int): adjusts number of channels of the hidden layer in ``InvertedResidual`` by this amount. Defaults to 4. """ layer = [] for i in range(num_blocks): stride = stride if i == 0 else 1 layer.append( InvertedResidual( in_channels=in_channels, out_channels=out_channels, stride=stride, expand_ratio=expand_ratio, act_cfg=dict(type='Swish'), )) in_channels = out_channels return nn.Sequential(*layer), out_channels
def _freeze_stages(self): for i in range(0, self.frozen_stages): layer = self.layers[i] layer.eval() for param in layer.parameters(): param.requires_grad = False def train(self, mode=True): super(MobileViT, self).train(mode) self._freeze_stages() def forward(self, x): x = self.stem(x) outs = [] for i, layer in enumerate(self.layers): x = layer(x) if i == len(self.layers) - 1: x = self.conv_1x1_exp(x) if i in self.out_indices: outs.append(x) return tuple(outs)
Read the Docs v: dev
Versions
latest
stable
mmcls-1.x
mmcls-0.x
dev
Downloads
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.