Shortcuts

mmcls.models.backbones.levit 源代码

# Copyright (c) OpenMMLab. All rights reserved.
import itertools

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

from mmcls.models.backbones.base_backbone import BaseBackbone
from mmcls.registry import MODELS
from ..utils import build_norm_layer


class HybridBackbone(BaseModule):

    def __init__(
            self,
            embed_dim,
            kernel_size=3,
            stride=2,
            pad=1,
            dilation=1,
            groups=1,
            act_cfg=dict(type='HSwish'),
            conv_cfg=None,
            norm_cfg=dict(type='BN'),
            init_cfg=None,
    ):
        super(HybridBackbone, self).__init__(init_cfg=init_cfg)

        self.input_channels = [
            3, embed_dim // 8, embed_dim // 4, embed_dim // 2
        ]
        self.output_channels = [
            embed_dim // 8, embed_dim // 4, embed_dim // 2, embed_dim
        ]
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg

        self.patch_embed = Sequential()

        for i in range(len(self.input_channels)):
            conv_bn = ConvolutionBatchNorm(
                self.input_channels[i],
                self.output_channels[i],
                kernel_size=kernel_size,
                stride=stride,
                pad=pad,
                dilation=dilation,
                groups=groups,
                norm_cfg=norm_cfg,
            )
            self.patch_embed.add_module('%d' % (2 * i), conv_bn)
            if i < len(self.input_channels) - 1:
                self.patch_embed.add_module('%d' % (i * 2 + 1),
                                            build_activation_layer(act_cfg))

    def forward(self, x):
        x = self.patch_embed(x)
        return x


class ConvolutionBatchNorm(BaseModule):

    def __init__(
            self,
            in_channel,
            out_channel,
            kernel_size=3,
            stride=2,
            pad=1,
            dilation=1,
            groups=1,
            norm_cfg=dict(type='BN'),
    ):
        super(ConvolutionBatchNorm, self).__init__()
        self.conv = nn.Conv2d(
            in_channel,
            out_channel,
            kernel_size=kernel_size,
            stride=stride,
            padding=pad,
            dilation=dilation,
            groups=groups,
            bias=False)
        self.bn = build_norm_layer(norm_cfg, out_channel)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return x

    @torch.no_grad()
    def fuse(self):
        return fuse_conv_bn(self).conv


class LinearBatchNorm(BaseModule):

    def __init__(self, in_feature, out_feature, norm_cfg=dict(type='BN1d')):
        super(LinearBatchNorm, self).__init__()
        self.linear = nn.Linear(in_feature, out_feature, bias=False)
        self.bn = build_norm_layer(norm_cfg, out_feature)

    def forward(self, x):
        x = self.linear(x)
        x = self.bn(x.flatten(0, 1)).reshape_as(x)
        return x

    @torch.no_grad()
    def fuse(self):
        w = self.bn.weight / (self.bn.running_var + self.bn.eps)**0.5
        w = self.linear.weight * w[:, None]
        b = self.bn.bias - self.bn.running_mean * self.bn.weight / \
            (self.bn.running_var + self.bn.eps) ** 0.5

        factory_kwargs = {
            'device': self.linear.weight.device,
            'dtype': self.linear.weight.dtype
        }
        bias = nn.Parameter(
            torch.empty(self.linear.out_features, **factory_kwargs))
        self.linear.register_parameter('bias', bias)
        self.linear.weight.data.copy_(w)
        self.linear.bias.data.copy_(b)
        return self.linear


class Residual(BaseModule):

    def __init__(self, block, drop_path_rate=0.):
        super(Residual, self).__init__()
        self.block = block
        if drop_path_rate > 0:
            self.drop_path = DropPath(drop_path_rate)
        else:
            self.drop_path = nn.Identity()

    def forward(self, x):
        x = x + self.drop_path(self.block(x))
        return x


class Attention(BaseModule):

    def __init__(
            self,
            dim,
            key_dim,
            num_heads=8,
            attn_ratio=4,
            act_cfg=dict(type='HSwish'),
            resolution=14,
    ):
        super(Attention, self).__init__()
        self.num_heads = num_heads
        self.scale = key_dim**-0.5
        self.key_dim = key_dim
        self.nh_kd = nh_kd = key_dim * num_heads
        self.d = int(attn_ratio * key_dim)
        self.dh = int(attn_ratio * key_dim) * num_heads
        self.attn_ratio = attn_ratio
        h = self.dh + nh_kd * 2
        self.qkv = LinearBatchNorm(dim, h)
        self.proj = nn.Sequential(
            build_activation_layer(act_cfg), LinearBatchNorm(self.dh, dim))

        points = list(itertools.product(range(resolution), range(resolution)))
        N = len(points)
        attention_offsets = {}
        idxs = []
        for p1 in points:
            for p2 in points:
                offset = (abs(p1[0] - p2[0]), abs(p1[1] - p2[1]))
                if offset not in attention_offsets:
                    attention_offsets[offset] = len(attention_offsets)
                idxs.append(attention_offsets[offset])
        self.attention_biases = torch.nn.Parameter(
            torch.zeros(num_heads, len(attention_offsets)))
        self.register_buffer('attention_bias_idxs',
                             torch.LongTensor(idxs).view(N, N))

    @torch.no_grad()
    def train(self, mode=True):
        """change the mode of model."""
        super(Attention, self).train(mode)
        if mode and hasattr(self, 'ab'):
            del self.ab
        else:
            self.ab = self.attention_biases[:, self.attention_bias_idxs]

    def forward(self, x):  # x (B,N,C)
        B, N, C = x.shape  # 2 196 128
        qkv = self.qkv(x)  # 2 196 128
        q, k, v = qkv.view(B, N, self.num_heads, -1).split(
            [self.key_dim, self.key_dim, self.d],
            dim=3)  # q 2 196 4 16 ; k 2 196 4 16; v 2 196 4 32
        q = q.permute(0, 2, 1, 3)  # 2 4 196 16
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)

        attn = ((q @ k.transpose(-2, -1)) *
                self.scale  # 2 4 196 16 * 2 4 16 196 -> 2 4 196 196
                + (self.attention_biases[:, self.attention_bias_idxs]
                   if self.training else self.ab))
        attn = attn.softmax(dim=-1)  # 2 4 196 196 -> 2 4 196 196
        x = (attn @ v).transpose(1, 2).reshape(
            B, N,
            self.dh)  # 2 4 196 196 * 2 4 196 32 -> 2 4 196 32 -> 2 196 128
        x = self.proj(x)
        return x


class MLP(nn.Sequential):

    def __init__(self, embed_dim, mlp_ratio, act_cfg=dict(type='HSwish')):
        super(MLP, self).__init__()
        h = embed_dim * mlp_ratio
        self.linear1 = LinearBatchNorm(embed_dim, h)
        self.activation = build_activation_layer(act_cfg)
        self.linear2 = LinearBatchNorm(h, embed_dim)

    def forward(self, x):
        x = self.linear1(x)
        x = self.activation(x)
        x = self.linear2(x)
        return x


class Subsample(BaseModule):

    def __init__(self, stride, resolution):
        super(Subsample, self).__init__()
        self.stride = stride
        self.resolution = resolution

    def forward(self, x):
        B, _, C = x.shape
        # B, N, C -> B, H, W, C
        x = x.view(B, self.resolution, self.resolution, C)
        x = x[:, ::self.stride, ::self.stride]
        x = x.reshape(B, -1, C)  # B, H', W', C -> B, N', C
        return x


class AttentionSubsample(nn.Sequential):

    def __init__(self,
                 in_dim,
                 out_dim,
                 key_dim,
                 num_heads=8,
                 attn_ratio=2,
                 act_cfg=dict(type='HSwish'),
                 stride=2,
                 resolution=14):
        super(AttentionSubsample, self).__init__()
        self.num_heads = num_heads
        self.scale = key_dim**-0.5
        self.key_dim = key_dim
        self.nh_kd = nh_kd = key_dim * num_heads
        self.d = int(attn_ratio * key_dim)
        self.dh = int(attn_ratio * key_dim) * self.num_heads
        self.attn_ratio = attn_ratio
        self.sub_resolution = (resolution - 1) // stride + 1
        h = self.dh + nh_kd
        self.kv = LinearBatchNorm(in_dim, h)

        self.q = nn.Sequential(
            Subsample(stride, resolution), LinearBatchNorm(in_dim, nh_kd))
        self.proj = nn.Sequential(
            build_activation_layer(act_cfg), LinearBatchNorm(self.dh, out_dim))

        self.stride = stride
        self.resolution = resolution
        points = list(itertools.product(range(resolution), range(resolution)))
        sub_points = list(
            itertools.product(
                range(self.sub_resolution), range(self.sub_resolution)))
        N = len(points)
        N_sub = len(sub_points)
        attention_offsets = {}
        idxs = []
        for p1 in sub_points:
            for p2 in points:
                size = 1
                offset = (abs(p1[0] * stride - p2[0] + (size - 1) / 2),
                          abs(p1[1] * stride - p2[1] + (size - 1) / 2))
                if offset not in attention_offsets:
                    attention_offsets[offset] = len(attention_offsets)
                idxs.append(attention_offsets[offset])
        self.attention_biases = torch.nn.Parameter(
            torch.zeros(num_heads, len(attention_offsets)))
        self.register_buffer('attention_bias_idxs',
                             torch.LongTensor(idxs).view(N_sub, N))

    @torch.no_grad()
    def train(self, mode=True):
        super(AttentionSubsample, self).train(mode)
        if mode and hasattr(self, 'ab'):
            del self.ab
        else:
            self.ab = self.attention_biases[:, self.attention_bias_idxs]

    def forward(self, x):
        B, N, C = x.shape
        k, v = self.kv(x).view(B, N, self.num_heads,
                               -1).split([self.key_dim, self.d], dim=3)
        k = k.permute(0, 2, 1, 3)  # BHNC
        v = v.permute(0, 2, 1, 3)  # BHNC
        q = self.q(x).view(B, self.sub_resolution**2, self.num_heads,
                           self.key_dim).permute(0, 2, 1, 3)

        attn = (q @ k.transpose(-2, -1)) * self.scale + \
               (self.attention_biases[:, self.attention_bias_idxs]
                if self.training else self.ab)
        attn = attn.softmax(dim=-1)

        x = (attn @ v).transpose(1, 2).reshape(B, -1, self.dh)
        x = self.proj(x)
        return x


[文档]@MODELS.register_module() class LeViT(BaseBackbone): """LeViT backbone. A PyTorch implementation of `LeViT: A Vision Transformer in ConvNet's Clothing for Faster Inference <https://arxiv.org/abs/2104.01136>`_ Modified from the official implementation: https://github.com/facebookresearch/LeViT Args: arch (str | dict): LeViT architecture. If use string, choose from '128s', '128', '192', '256' and '384'. If use dict, it should have below keys: - **embed_dims** (List[int]): The embed dimensions of each stage. - **key_dims** (List[int]): The embed dimensions of the key in the attention layers of each stage. - **num_heads** (List[int]): The number of heads in each stage. - **depths** (List[int]): The number of blocks in each stage. img_size (int): Input image size patch_size (int | tuple): The patch size. Deault to 16 attn_ratio (int): Ratio of hidden dimensions of the value in attention layers. Defaults to 2. mlp_ratio (int): Ratio of hidden dimensions in MLP layers. Defaults to 2. act_cfg (dict): The config of activation functions. Defaults to ``dict(type='HSwish')``. hybrid_backbone (callable): A callable object to build the patch embed module. Defaults to use :class:`HybridBackbone`. out_indices (Sequence | int): Output from which stages. Defaults to -1, means the last stage. deploy (bool): Whether to switch the model structure to deployment mode. Defaults to False. init_cfg (dict or list[dict], optional): Initialization config dict. Defaults to None. """ arch_zoo = { '128s': { 'embed_dims': [128, 256, 384], 'num_heads': [4, 6, 8], 'depths': [2, 3, 4], 'key_dims': [16, 16, 16], }, '128': { 'embed_dims': [128, 256, 384], 'num_heads': [4, 8, 12], 'depths': [4, 4, 4], 'key_dims': [16, 16, 16], }, '192': { 'embed_dims': [192, 288, 384], 'num_heads': [3, 5, 6], 'depths': [4, 4, 4], 'key_dims': [32, 32, 32], }, '256': { 'embed_dims': [256, 384, 512], 'num_heads': [4, 6, 8], 'depths': [4, 4, 4], 'key_dims': [32, 32, 32], }, '384': { 'embed_dims': [384, 512, 768], 'num_heads': [6, 9, 12], 'depths': [4, 4, 4], 'key_dims': [32, 32, 32], }, } def __init__(self, arch, img_size=224, patch_size=16, attn_ratio=2, mlp_ratio=2, act_cfg=dict(type='HSwish'), hybrid_backbone=HybridBackbone, out_indices=-1, deploy=False, drop_path_rate=0, init_cfg=None): super(LeViT, self).__init__(init_cfg=init_cfg) 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 = self.arch_zoo[arch] elif isinstance(arch, dict): essential_keys = {'embed_dim', 'num_heads', 'depth', 'key_dim'} assert isinstance(arch, dict) and set(arch) == essential_keys, \ f'Custom arch needs a dict with keys {essential_keys}' self.arch = arch else: raise TypeError('Expect "arch" to be either a string ' f'or a dict, got {type(arch)}') self.embed_dims = self.arch['embed_dims'] self.num_heads = self.arch['num_heads'] self.key_dims = self.arch['key_dims'] self.depths = self.arch['depths'] self.num_stages = len(self.embed_dims) self.drop_path_rate = drop_path_rate self.patch_embed = hybrid_backbone(self.embed_dims[0]) self.resolutions = [] resolution = img_size // patch_size self.stages = ModuleList() for i, (embed_dims, key_dims, depth, num_heads) in enumerate( zip(self.embed_dims, self.key_dims, self.depths, self.num_heads)): blocks = [] if i > 0: downsample = AttentionSubsample( in_dim=self.embed_dims[i - 1], out_dim=embed_dims, key_dim=key_dims, num_heads=self.embed_dims[i - 1] // key_dims, attn_ratio=4, act_cfg=act_cfg, stride=2, resolution=resolution) blocks.append(downsample) resolution = downsample.sub_resolution if mlp_ratio > 0: # mlp_ratio blocks.append( Residual( MLP(embed_dims, mlp_ratio, act_cfg=act_cfg), self.drop_path_rate)) self.resolutions.append(resolution) for _ in range(depth): blocks.append( Residual( Attention( embed_dims, key_dims, num_heads, attn_ratio=attn_ratio, act_cfg=act_cfg, resolution=resolution, ), self.drop_path_rate)) if mlp_ratio > 0: blocks.append( Residual( MLP(embed_dims, mlp_ratio, act_cfg=act_cfg), self.drop_path_rate)) self.stages.append(Sequential(*blocks)) if isinstance(out_indices, int): out_indices = [out_indices] elif isinstance(out_indices, tuple): out_indices = list(out_indices) elif not isinstance(out_indices, list): raise TypeError('"out_indices" must by a list, tuple 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 0 <= out_indices[i] < self.num_stages, \ f'Invalid out_indices {index}.' self.out_indices = out_indices self.deploy = False if deploy: self.switch_to_deploy() def switch_to_deploy(self): if self.deploy: return fuse_parameters(self) self.deploy = True def forward(self, x): x = self.patch_embed(x) x = x.flatten(2).transpose(1, 2) # B, C, H, W -> B, L, C outs = [] for i, stage in enumerate(self.stages): x = stage(x) B, _, C = x.shape if i in self.out_indices: out = x.reshape(B, self.resolutions[i], self.resolutions[i], C) out = out.permute(0, 3, 1, 2).contiguous() outs.append(out) return tuple(outs)
def fuse_parameters(module): for child_name, child in module.named_children(): if hasattr(child, 'fuse'): setattr(module, child_name, child.fuse()) else: fuse_parameters(child)
Read the Docs v: mmcls-1.x
Versions
latest
stable
mmcls-1.x
mmcls-0.x
Downloads
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.