mmpretrain/models/backbones/edgenext.py

# Copyright (c) OpenMMLab. All rights reserved.
import math
from itertools import chain
from typing import Sequence

import torch
import torch.nn as nn
from mmcv.cnn.bricks import DropPath
from mmengine.model import BaseModule, ModuleList, Sequential

from mmpretrain.registry import MODELS
from ..utils import (ChannelMultiheadAttention, PositionEncodingFourier,
                     build_norm_layer)
from .base_backbone import BaseBackbone
from .convnext import ConvNeXtBlock


class SDTAEncoder(BaseModule):
    """A PyTorch implementation of split depth-wise transpose attention (SDTA)
    encoder.

    Inspiration from
    https://github.com/mmaaz60/EdgeNeXt
    Args:
        in_channel (int): Number of input channels.
        drop_path_rate (float): Stochastic depth dropout rate.
            Defaults to 0.
        layer_scale_init_value (float): Initial value of layer scale.
            Defaults to 1e-6.
        mlp_ratio (int): Number of channels ratio in the MLP.
            Defaults to 4.
        use_pos_emb (bool): Whether to use position encoding.
            Defaults to True.
        num_heads (int): Number of heads in the multihead attention.
            Defaults to 8.
        qkv_bias (bool): Whether to use bias in the multihead attention.
            Defaults to True.
        attn_drop (float): Dropout rate of the attention.
            Defaults to 0.
        proj_drop (float): Dropout rate of the projection.
            Defaults to 0.
        layer_scale_init_value (float): Initial value of layer scale.
            Defaults to 1e-6.
        norm_cfg (dict): Dictionary to construct normalization layer.
            Defaults to ``dict(type='LN')``.
        act_cfg (dict): Dictionary to construct activation layer.
            Defaults to ``dict(type='GELU')``.
        scales (int): Number of scales. Default to 1.
    """

    def __init__(self,
                 in_channel,
                 drop_path_rate=0.,
                 layer_scale_init_value=1e-6,
                 mlp_ratio=4,
                 use_pos_emb=True,
                 num_heads=8,
                 qkv_bias=True,
                 attn_drop=0.,
                 proj_drop=0.,
                 norm_cfg=dict(type='LN'),
                 act_cfg=dict(type='GELU'),
                 scales=1,
                 init_cfg=None):
        super(SDTAEncoder, self).__init__(init_cfg=init_cfg)
        conv_channels = max(
            int(math.ceil(in_channel / scales)),
            int(math.floor(in_channel // scales)))
        self.conv_channels = conv_channels
        self.num_convs = scales if scales == 1 else scales - 1

        self.conv_modules = ModuleList()
        for i in range(self.num_convs):
            self.conv_modules.append(
                nn.Conv2d(
                    conv_channels,
                    conv_channels,
                    kernel_size=3,
                    padding=1,
                    groups=conv_channels))

        self.pos_embed = PositionEncodingFourier(
            embed_dims=in_channel) if use_pos_emb else None

        self.norm_csa = build_norm_layer(norm_cfg, in_channel)
        self.gamma_csa = nn.Parameter(
            layer_scale_init_value * torch.ones(in_channel),
            requires_grad=True) if layer_scale_init_value > 0 else None
        self.csa = ChannelMultiheadAttention(
            embed_dims=in_channel,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            attn_drop=attn_drop,
            proj_drop=proj_drop)

        self.norm = build_norm_layer(norm_cfg, in_channel)
        self.pointwise_conv1 = nn.Linear(in_channel, mlp_ratio * in_channel)
        self.act = MODELS.build(act_cfg)
        self.pointwise_conv2 = nn.Linear(mlp_ratio * in_channel, in_channel)
        self.gamma = nn.Parameter(
            layer_scale_init_value * torch.ones(in_channel),
            requires_grad=True) if layer_scale_init_value > 0 else None
        self.drop_path = DropPath(
            drop_path_rate) if drop_path_rate > 0. else nn.Identity()

    def forward(self, x):
        shortcut = x
        spx = torch.split(x, self.conv_channels, dim=1)
        for i in range(self.num_convs):
            if i == 0:
                sp = spx[i]
            else:
                sp = sp + spx[i]
            sp = self.conv_modules[i](sp)
            if i == 0:
                out = sp
            else:
                out = torch.cat((out, sp), 1)

        x = torch.cat((out, spx[self.num_convs]), 1)

        # Channel Self-attention
        B, C, H, W = x.shape
        x = x.reshape(B, C, H * W).permute(0, 2, 1)
        if self.pos_embed:
            pos_encoding = self.pos_embed((B, H, W))
            pos_encoding = pos_encoding.reshape(B, -1,
                                                x.shape[1]).permute(0, 2, 1)
            x += pos_encoding

        x = x + self.drop_path(self.gamma_csa * self.csa(self.norm_csa(x)))
        x = x.reshape(B, H, W, C)

        # Inverted Bottleneck
        x = self.norm(x)
        x = self.pointwise_conv1(x)
        x = self.act(x)
        x = self.pointwise_conv2(x)

        if self.gamma is not None:
            x = self.gamma * x
        x = x.permute(0, 3, 1, 2)  # (B, H, W, C) -> (B, C, H, W)

        x = shortcut + self.drop_path(x)

        return x


@MODELS.register_module()
class EdgeNeXt(BaseBackbone):
    """EdgeNeXt.

    A PyTorch implementation of: `EdgeNeXt: Efficiently Amalgamated
    CNN-Transformer Architecture for Mobile Vision Applications
    <https://arxiv.org/abs/2206.10589>`_

    Inspiration from
    https://github.com/mmaaz60/EdgeNeXt

    Args:
        arch (str | dict): The model's architecture. If string, it should be
            one of architectures in ``EdgeNeXt.arch_settings``.
            And if dict, it should include the following keys:

            - channels (list[int]): The number of channels at each stage.
            - depths (list[int]): The number of blocks at each stage.
            - num_heads (list[int]): The number of heads at each stage.

            Defaults to 'xxsmall'.
        in_channels (int): The number of input channels.
            Defaults to 3.
        global_blocks (list[int]): The number of global blocks.
            Defaults to [0, 1, 1, 1].
        global_block_type (list[str]): The type of global blocks.
            Defaults to ['None', 'SDTA', 'SDTA', 'SDTA'].
        drop_path_rate (float): Stochastic depth dropout rate.
            Defaults to 0.
        layer_scale_init_value (float): Initial value of layer scale.
            Defaults to 1e-6.
        linear_pw_conv (bool): Whether to use linear layer to do pointwise
            convolution. Defaults to False.
        mlp_ratio (int): The number of channel ratio in MLP layers.
            Defaults to 4.
        conv_kernel_size (list[int]): The kernel size of convolutional layers
            at each stage. Defaults to [3, 5, 7, 9].
        use_pos_embd_csa (list[bool]): Whether to use positional embedding in
            Channel Self-Attention. Defaults to [False, True, False, False].
        use_pos_emebd_global (bool): Whether to use positional embedding for
            whole network. Defaults to False.
        d2_scales (list[int]): The number of channel groups used for SDTA at
            each stage. Defaults to [2, 2, 3, 4].
        norm_cfg (dict): The config of normalization layer.
            Defaults to ``dict(type='LN2d', eps=1e-6)``.
        out_indices (Sequence | int): Output from which stages.
            Defaults to -1, means the last stage.
        frozen_stages (int): Stages to be frozen (all param fixed).
            Defaults to 0, which means not freezing any parameters.
        gap_before_final_norm (bool): Whether to globally average the feature
            map before the final norm layer. Defaults to True.
        act_cfg (dict): The config of activation layer.
            Defaults to ``dict(type='GELU')``.
        init_cfg (dict, optional): Config for initialization.
            Defaults to None.
    """
    arch_settings = {
        'xxsmall': {  # parameters: 1.3M
            'channels': [24, 48, 88, 168],
            'depths': [2, 2, 6, 2],
            'num_heads': [4, 4, 4, 4]
        },
        'xsmall': {  # parameters: 2.3M
            'channels': [32, 64, 100, 192],
            'depths': [3, 3, 9, 3],
            'num_heads': [4, 4, 4, 4]
        },
        'small': {  # parameters: 5.6M
            'channels': [48, 96, 160, 304],
            'depths': [3, 3, 9, 3],
            'num_heads': [8, 8, 8, 8]
        },
        'base': {  # parameters: 18.51M
            'channels': [80, 160, 288, 584],
            'depths': [3, 3, 9, 3],
            'num_heads': [8, 8, 8, 8]
        },
    }

    def __init__(self,
                 arch='xxsmall',
                 in_channels=3,
                 global_blocks=[0, 1, 1, 1],
                 global_block_type=['None', 'SDTA', 'SDTA', 'SDTA'],
                 drop_path_rate=0.,
                 layer_scale_init_value=1e-6,
                 linear_pw_conv=True,
                 mlp_ratio=4,
                 conv_kernel_sizes=[3, 5, 7, 9],
                 use_pos_embd_csa=[False, True, False, False],
                 use_pos_embd_global=False,
                 d2_scales=[2, 2, 3, 4],
                 norm_cfg=dict(type='LN2d', eps=1e-6),
                 out_indices=-1,
                 frozen_stages=0,
                 gap_before_final_norm=True,
                 act_cfg=dict(type='GELU'),
                 init_cfg=None):
        super(EdgeNeXt, self).__init__(init_cfg=init_cfg)

        if isinstance(arch, str):
            arch = arch.lower()
            assert arch in self.arch_settings, \
                f'Arch {arch} is not in default archs ' \
                f'{set(self.arch_settings)}'
            self.arch_settings = self.arch_settings[arch]
        elif isinstance(arch, dict):
            essential_keys = {'channels', 'depths', 'num_heads'}
            assert isinstance(arch, dict) and set(arch) == essential_keys, \
                f'Custom arch needs a dict with keys {essential_keys}'
            self.arch_settings = arch

        self.channels = self.arch_settings['channels']
        self.depths = self.arch_settings['depths']
        self.num_heads = self.arch_settings['num_heads']
        self.num_layers = len(self.depths)
        self.use_pos_embd_global = use_pos_embd_global

        for g in global_block_type:
            assert g in ['None',
                         'SDTA'], f'Global block type {g} is not supported'

        self.num_stages = len(self.depths)

        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] = 4 + index
                assert out_indices[i] >= 0, f'Invalid out_indices {index}'
        self.out_indices = out_indices

        self.frozen_stages = frozen_stages
        self.gap_before_final_norm = gap_before_final_norm

        if self.use_pos_embd_global:
            self.pos_embed = PositionEncodingFourier(
                embed_dims=self.channels[0])
        else:
            self.pos_embed = None

        # stochastic depth decay rule
        dpr = [
            x.item()
            for x in torch.linspace(0, drop_path_rate, sum(self.depths))
        ]

        self.downsample_layers = ModuleList()
        stem = nn.Sequential(
            nn.Conv2d(in_channels, self.channels[0], kernel_size=4, stride=4),
            build_norm_layer(norm_cfg, self.channels[0]),
        )
        self.downsample_layers.append(stem)

        self.stages = ModuleList()
        block_idx = 0
        for i in range(self.num_stages):
            depth = self.depths[i]
            channels = self.channels[i]

            if i >= 1:
                downsample_layer = nn.Sequential(
                    build_norm_layer(norm_cfg, self.channels[i - 1]),
                    nn.Conv2d(
                        self.channels[i - 1],
                        channels,
                        kernel_size=2,
                        stride=2,
                    ))
                self.downsample_layers.append(downsample_layer)

            stage_blocks = []
            for j in range(depth):
                if j > depth - global_blocks[i] - 1:
                    stage_blocks.append(
                        SDTAEncoder(
                            in_channel=channels,
                            drop_path_rate=dpr[block_idx + j],
                            mlp_ratio=mlp_ratio,
                            scales=d2_scales[i],
                            use_pos_emb=use_pos_embd_csa[i],
                            num_heads=self.num_heads[i],
                        ))
                else:
                    dw_conv_cfg = dict(
                        kernel_size=conv_kernel_sizes[i],
                        padding=conv_kernel_sizes[i] // 2,
                    )
                    stage_blocks.append(
                        ConvNeXtBlock(
                            in_channels=channels,
                            dw_conv_cfg=dw_conv_cfg,
                            norm_cfg=norm_cfg,
                            act_cfg=act_cfg,
                            linear_pw_conv=linear_pw_conv,
                            drop_path_rate=dpr[block_idx + j],
                            layer_scale_init_value=layer_scale_init_value,
                        ))
            block_idx += depth

            stage_blocks = Sequential(*stage_blocks)
            self.stages.append(stage_blocks)

            if i in self.out_indices:
                out_norm_cfg = dict(type='LN') if self.gap_before_final_norm \
                    else norm_cfg
                norm_layer = build_norm_layer(out_norm_cfg, channels)
                self.add_module(f'norm{i}', norm_layer)

    def init_weights(self) -> None:
        # TODO: need to be implemented in the future
        return super().init_weights()

    def forward(self, x):
        outs = []
        for i, stage in enumerate(self.stages):
            x = self.downsample_layers[i](x)
            x = stage(x)
            if self.pos_embed and i == 0:
                B, _, H, W = x.shape
                x += self.pos_embed((B, H, W))

            if i in self.out_indices:
                norm_layer = getattr(self, f'norm{i}')
                if self.gap_before_final_norm:
                    gap = x.mean([-2, -1], keepdim=True)
                    outs.append(norm_layer(gap.flatten(1)))
                else:
                    # The output of LayerNorm2d may be discontiguous, which
                    # may cause some problem in the downstream tasks
                    outs.append(norm_layer(x).contiguous())

        return tuple(outs)

    def _freeze_stages(self):
        for i in range(self.frozen_stages):
            downsample_layer = self.downsample_layers[i]
            stage = self.stages[i]
            downsample_layer.eval()
            stage.eval()
            for param in chain(downsample_layer.parameters(),
                               stage.parameters()):
                param.requires_grad = False

    def train(self, mode=True):
        super(EdgeNeXt, self).train(mode)
        self._freeze_stages()