mmpretrain/models/backbones/hornet.py

# Copyright (c) OpenMMLab. All rights reserved.
# Adapted from official impl at https://github.com/raoyongming/HorNet.
try:
    import torch.fft
    fft = True
except ImportError:
    fft = None

import copy
from functools import partial
from typing import Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
from mmcv.cnn.bricks import DropPath

from mmpretrain.models.backbones.base_backbone import BaseBackbone
from mmpretrain.registry import MODELS
from ..utils import LayerScale


def get_dwconv(dim, kernel_size, bias=True):
    """build a pepth-wise convolution."""
    return nn.Conv2d(
        dim,
        dim,
        kernel_size=kernel_size,
        padding=(kernel_size - 1) // 2,
        bias=bias,
        groups=dim)


class HorNetLayerNorm(nn.Module):
    """An implementation of LayerNorm of HorNet.

    The differences between HorNetLayerNorm & torch LayerNorm:
        1. Supports two data formats channels_last or channels_first.
    Args:
        normalized_shape (int or list or torch.Size): input shape from an
            expected input of size.
        eps (float): a value added to the denominator for numerical stability.
            Defaults to 1e-5.
        data_format (str): The ordering of the dimensions in the inputs.
            channels_last corresponds to inputs with shape (batch_size, height,
            width, channels) while channels_first corresponds to inputs with
            shape (batch_size, channels, height, width).
            Defaults to 'channels_last'.
    """

    def __init__(self,
                 normalized_shape,
                 eps=1e-6,
                 data_format='channels_last'):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.eps = eps
        self.data_format = data_format
        if self.data_format not in ['channels_last', 'channels_first']:
            raise ValueError(
                'data_format must be channels_last or channels_first')
        self.normalized_shape = (normalized_shape, )

    def forward(self, x):
        if self.data_format == 'channels_last':
            return F.layer_norm(x, self.normalized_shape, self.weight,
                                self.bias, self.eps)
        elif self.data_format == 'channels_first':
            u = x.mean(1, keepdim=True)
            s = (x - u).pow(2).mean(1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.eps)
            x = self.weight[:, None, None] * x + self.bias[:, None, None]
            return x


class GlobalLocalFilter(nn.Module):
    """A GlobalLocalFilter of HorNet.

    Args:
        dim (int): Number of input channels.
        h (int): Height of complex_weight.
            Defaults to 14.
        w (int): Width of complex_weight.
            Defaults to 8.
    """

    def __init__(self, dim, h=14, w=8):
        super().__init__()
        self.dw = nn.Conv2d(
            dim // 2,
            dim // 2,
            kernel_size=3,
            padding=1,
            bias=False,
            groups=dim // 2)
        self.complex_weight = nn.Parameter(
            torch.randn(dim // 2, h, w, 2, dtype=torch.float32) * 0.02)
        self.pre_norm = HorNetLayerNorm(
            dim, eps=1e-6, data_format='channels_first')
        self.post_norm = HorNetLayerNorm(
            dim, eps=1e-6, data_format='channels_first')

    def forward(self, x):
        x = self.pre_norm(x)
        x1, x2 = torch.chunk(x, 2, dim=1)
        x1 = self.dw(x1)

        x2 = x2.to(torch.float32)
        B, C, a, b = x2.shape
        x2 = torch.fft.rfft2(x2, dim=(2, 3), norm='ortho')

        weight = self.complex_weight
        if not weight.shape[1:3] == x2.shape[2:4]:
            weight = F.interpolate(
                weight.permute(3, 0, 1, 2),
                size=x2.shape[2:4],
                mode='bilinear',
                align_corners=True).permute(1, 2, 3, 0)

        weight = torch.view_as_complex(weight.contiguous())

        x2 = x2 * weight
        x2 = torch.fft.irfft2(x2, s=(a, b), dim=(2, 3), norm='ortho')

        x = torch.cat([x1.unsqueeze(2), x2.unsqueeze(2)],
                      dim=2).reshape(B, 2 * C, a, b)
        x = self.post_norm(x)
        return x


class gnConv(nn.Module):
    """A gnConv of HorNet.

    Args:
        dim (int): Number of input channels.
        order (int): Order of gnConv.
            Defaults to 5.
        dw_cfg (dict): The Config for dw conv.
            Defaults to ``dict(type='DW', kernel_size=7)``.
        scale (float): Scaling parameter of gflayer outputs.
            Defaults to 1.0.
    """

    def __init__(self,
                 dim,
                 order=5,
                 dw_cfg=dict(type='DW', kernel_size=7),
                 scale=1.0):
        super().__init__()
        self.order = order
        self.dims = [dim // 2**i for i in range(order)]
        self.dims.reverse()
        self.proj_in = nn.Conv2d(dim, 2 * dim, 1)

        cfg = copy.deepcopy(dw_cfg)
        dw_type = cfg.pop('type')
        assert dw_type in ['DW', 'GF'],\
            'dw_type should be `DW` or `GF`'
        if dw_type == 'DW':
            self.dwconv = get_dwconv(sum(self.dims), **cfg)
        elif dw_type == 'GF':
            self.dwconv = GlobalLocalFilter(sum(self.dims), **cfg)

        self.proj_out = nn.Conv2d(dim, dim, 1)

        self.projs = nn.ModuleList([
            nn.Conv2d(self.dims[i], self.dims[i + 1], 1)
            for i in range(order - 1)
        ])

        self.scale = scale

    def forward(self, x):
        x = self.proj_in(x)
        y, x = torch.split(x, (self.dims[0], sum(self.dims)), dim=1)

        x = self.dwconv(x) * self.scale

        dw_list = torch.split(x, self.dims, dim=1)
        x = y * dw_list[0]

        for i in range(self.order - 1):
            x = self.projs[i](x) * dw_list[i + 1]

        x = self.proj_out(x)

        return x


class HorNetBlock(nn.Module):
    """A block of HorNet.

    Args:
        dim (int): Number of input channels.
        order (int): Order of gnConv.
            Defaults to 5.
        dw_cfg (dict): The Config for dw conv.
            Defaults to ``dict(type='DW', kernel_size=7)``.
        scale (float): Scaling parameter of gflayer outputs.
            Defaults to 1.0.
        drop_path_rate (float): Stochastic depth rate. Defaults to 0.
        use_layer_scale (bool): Whether to use use_layer_scale in HorNet
             block. Defaults to True.
    """

    def __init__(self,
                 dim,
                 order=5,
                 dw_cfg=dict(type='DW', kernel_size=7),
                 scale=1.0,
                 drop_path_rate=0.,
                 use_layer_scale=True):
        super().__init__()
        self.out_channels = dim

        self.norm1 = HorNetLayerNorm(
            dim, eps=1e-6, data_format='channels_first')
        self.gnconv = gnConv(dim, order, dw_cfg, scale)
        self.norm2 = HorNetLayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(dim, 4 * dim)
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(4 * dim, dim)

        if use_layer_scale:
            self.gamma1 = LayerScale(dim, data_format='channels_first')
            self.gamma2 = LayerScale(dim)
        else:
            self.gamma1, self.gamma2 = nn.Identity(), nn.Identity()

        self.drop_path = DropPath(
            drop_path_rate) if drop_path_rate > 0. else nn.Identity()

    def forward(self, x):
        x = x + self.drop_path(self.gamma1(self.gnconv(self.norm1(x))))

        input = x
        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
        x = self.norm2(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        x = self.gamma2(x)
        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)

        x = input + self.drop_path(x)
        return x


@MODELS.register_module()
class HorNet(BaseBackbone):
    """HorNet backbone.

    A PyTorch implementation of paper `HorNet: Efficient High-Order Spatial
    Interactions with Recursive Gated Convolutions
    <https://arxiv.org/abs/2207.14284>`_ .
    Inspiration from https://github.com/raoyongming/HorNet

    Args:
        arch (str | dict): HorNet architecture.

            If use string, choose from 'tiny', 'small', 'base' and 'large'.
            If use dict, it should have below keys:

            - **base_dim** (int): The base dimensions of embedding.
            - **depths** (List[int]): The number of blocks in each stage.
            - **orders** (List[int]): The number of order of gnConv in each
                stage.
            - **dw_cfg** (List[dict]): The Config for dw conv.

            Defaults to 'tiny'.
        in_channels (int): Number of input image channels. Defaults to 3.
        drop_path_rate (float): Stochastic depth rate. Defaults to 0.
        scale (float): Scaling parameter of gflayer outputs. Defaults to 1/3.
        use_layer_scale (bool): Whether to use use_layer_scale in HorNet
            block. Defaults to True.
        out_indices (Sequence[int]): Output from which stages.
            Default: ``(3, )``.
        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
            -1 means not freezing any parameters. Defaults to -1.
        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
            memory while slowing down the training speed. Defaults to False.
        gap_before_final_norm (bool): Whether to globally average the feature
            map before the final norm layer. In the official repo, it's only
            used in classification task. Defaults to True.
        init_cfg (dict, optional): The Config for initialization.
            Defaults to None.
    """
    arch_zoo = {
        **dict.fromkeys(['t', 'tiny'],
                        {'base_dim': 64,
                         'depths': [2, 3, 18, 2],
                         'orders': [2, 3, 4, 5],
                         'dw_cfg': [dict(type='DW', kernel_size=7)] * 4}),
        **dict.fromkeys(['t-gf', 'tiny-gf'],
                        {'base_dim': 64,
                         'depths': [2, 3, 18, 2],
                         'orders': [2, 3, 4, 5],
                         'dw_cfg': [
                             dict(type='DW', kernel_size=7),
                             dict(type='DW', kernel_size=7),
                             dict(type='GF', h=14, w=8),
                             dict(type='GF', h=7, w=4)]}),
        **dict.fromkeys(['s', 'small'],
                        {'base_dim': 96,
                         'depths': [2, 3, 18, 2],
                         'orders': [2, 3, 4, 5],
                         'dw_cfg': [dict(type='DW', kernel_size=7)] * 4}),
        **dict.fromkeys(['s-gf', 'small-gf'],
                        {'base_dim': 96,
                         'depths': [2, 3, 18, 2],
                         'orders': [2, 3, 4, 5],
                         'dw_cfg': [
                             dict(type='DW', kernel_size=7),
                             dict(type='DW', kernel_size=7),
                             dict(type='GF', h=14, w=8),
                             dict(type='GF', h=7, w=4)]}),
        **dict.fromkeys(['b', 'base'],
                        {'base_dim': 128,
                         'depths': [2, 3, 18, 2],
                         'orders': [2, 3, 4, 5],
                         'dw_cfg': [dict(type='DW', kernel_size=7)] * 4}),
        **dict.fromkeys(['b-gf', 'base-gf'],
                        {'base_dim': 128,
                         'depths': [2, 3, 18, 2],
                         'orders': [2, 3, 4, 5],
                         'dw_cfg': [
                             dict(type='DW', kernel_size=7),
                             dict(type='DW', kernel_size=7),
                             dict(type='GF', h=14, w=8),
                             dict(type='GF', h=7, w=4)]}),
        **dict.fromkeys(['b-gf384', 'base-gf384'],
                        {'base_dim': 128,
                         'depths': [2, 3, 18, 2],
                         'orders': [2, 3, 4, 5],
                         'dw_cfg': [
                             dict(type='DW', kernel_size=7),
                             dict(type='DW', kernel_size=7),
                             dict(type='GF', h=24, w=12),
                             dict(type='GF', h=13, w=7)]}),
        **dict.fromkeys(['l', 'large'],
                        {'base_dim': 192,
                         'depths': [2, 3, 18, 2],
                         'orders': [2, 3, 4, 5],
                         'dw_cfg': [dict(type='DW', kernel_size=7)] * 4}),
        **dict.fromkeys(['l-gf', 'large-gf'],
                        {'base_dim': 192,
                         'depths': [2, 3, 18, 2],
                         'orders': [2, 3, 4, 5],
                         'dw_cfg': [
                             dict(type='DW', kernel_size=7),
                             dict(type='DW', kernel_size=7),
                             dict(type='GF', h=14, w=8),
                             dict(type='GF', h=7, w=4)]}),
        **dict.fromkeys(['l-gf384', 'large-gf384'],
                        {'base_dim': 192,
                         'depths': [2, 3, 18, 2],
                         'orders': [2, 3, 4, 5],
                         'dw_cfg': [
                             dict(type='DW', kernel_size=7),
                             dict(type='DW', kernel_size=7),
                             dict(type='GF', h=24, w=12),
                             dict(type='GF', h=13, w=7)]}),
    }  # yapf: disable

    def __init__(self,
                 arch='tiny',
                 in_channels=3,
                 drop_path_rate=0.,
                 scale=1 / 3,
                 use_layer_scale=True,
                 out_indices=(3, ),
                 frozen_stages=-1,
                 with_cp=False,
                 gap_before_final_norm=True,
                 init_cfg=None):
        super().__init__(init_cfg=init_cfg)
        if fft is None:
            raise RuntimeError(
                'Failed to import torch.fft. Please install "torch>=1.7".')

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

        self.scale = scale
        self.out_indices = out_indices
        self.frozen_stages = frozen_stages
        self.with_cp = with_cp
        self.gap_before_final_norm = gap_before_final_norm

        base_dim = self.arch_settings['base_dim']
        dims = list(map(lambda x: 2**x * base_dim, range(4)))

        self.downsample_layers = nn.ModuleList()
        stem = nn.Sequential(
            nn.Conv2d(in_channels, dims[0], kernel_size=4, stride=4),
            HorNetLayerNorm(dims[0], eps=1e-6, data_format='channels_first'))
        self.downsample_layers.append(stem)
        for i in range(3):
            downsample_layer = nn.Sequential(
                HorNetLayerNorm(
                    dims[i], eps=1e-6, data_format='channels_first'),
                nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
            )
            self.downsample_layers.append(downsample_layer)

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

        cur_block_idx = 0
        self.stages = nn.ModuleList()
        for i in range(4):
            stage = nn.Sequential(*[
                HorNetBlock(
                    dim=dims[i],
                    order=self.arch_settings['orders'][i],
                    dw_cfg=self.arch_settings['dw_cfg'][i],
                    scale=self.scale,
                    drop_path_rate=dpr[cur_block_idx + j],
                    use_layer_scale=use_layer_scale)
                for j in range(self.arch_settings['depths'][i])
            ])
            self.stages.append(stage)
            cur_block_idx += self.arch_settings['depths'][i]

        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.'
        out_indices = list(out_indices)
        for i, index in enumerate(out_indices):
            if index < 0:
                out_indices[i] = len(self.stages) + index
            assert 0 <= out_indices[i] <= len(self.stages), \
                f'Invalid out_indices {index}.'
        self.out_indices = out_indices

        norm_layer = partial(
            HorNetLayerNorm, eps=1e-6, data_format='channels_first')
        for i_layer in out_indices:
            layer = norm_layer(dims[i_layer])
            layer_name = f'norm{i_layer}'
            self.add_module(layer_name, layer)

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

    def _freeze_stages(self):
        for i in range(0, self.frozen_stages + 1):
            # freeze patch embed
            m = self.downsample_layers[i]
            m.eval()
            for param in m.parameters():
                param.requires_grad = False

            # freeze blocks
            m = self.stages[i]
            m.eval()
            for param in m.parameters():
                param.requires_grad = False

            if i in self.out_indices:
                # freeze norm
                m = getattr(self, f'norm{i + 1}')
                m.eval()
                for param in m.parameters():
                    param.requires_grad = False

    def forward(self, x):
        outs = []
        for i in range(4):
            x = self.downsample_layers[i](x)
            if self.with_cp:
                x = checkpoint.checkpoint_sequential(self.stages[i],
                                                     len(self.stages[i]), x)
            else:
                x = self.stages[i](x)
            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)