Source code for mmpretrain.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 mmpretrain.models.backbones.base_backbone import BaseBackbone
from mmpretrain.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
[docs]@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)