Shortcuts

Note

You are reading the documentation for MMClassification 0.x, which will soon be deprecated at the end of 2022. We recommend you upgrade to MMClassification 1.0 to enjoy fruitful new features and better performance brought by OpenMMLab 2.0. Check the installation tutorial, migration tutorial and changelog for more details.

Source code for mmcls.models.backbones.hornet

# 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 mmcls.models.builder import BACKBONES
from ..utils import LayerScale
from .base_backbone import BaseBackbone


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


[docs]@BACKBONES.register_module() class HorNet(BaseBackbone): """HorNet A PyTorch impl of : `HorNet: Efficient High-Order Spatial Interactions with Recursive Gated Convolutions` 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)
[docs] 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
[docs] 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)
Read the Docs v: latest
Versions
master
latest
1.x
dev-1.x
Downloads
html
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.