| |
| from typing import List |
|
|
| import torch |
| import torch.nn as nn |
| from torch import Tensor |
| import torch.nn.functional as F |
| from mmcv.cnn import ConvModule, DepthwiseSeparableConvModule, Linear |
| from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig |
| from mmengine.model import BaseModule |
| from mmyolo.registry import MODELS |
| from mmyolo.models.layers import CSPLayerWithTwoConv |
|
|
|
|
| @MODELS.register_module() |
| class MaxSigmoidAttnBlock(BaseModule): |
| """Max Sigmoid attention block.""" |
| def __init__(self, |
| in_channels: int, |
| out_channels: int, |
| guide_channels: int, |
| embed_channels: int, |
| kernel_size: int = 3, |
| padding: int = 1, |
| num_heads: int = 1, |
| use_depthwise: bool = False, |
| with_scale: bool = False, |
| conv_cfg: OptConfigType = None, |
| norm_cfg: ConfigType = dict(type='BN', |
| momentum=0.03, |
| eps=0.001), |
| init_cfg: OptMultiConfig = None) -> None: |
| super().__init__(init_cfg=init_cfg) |
| conv = DepthwiseSeparableConvModule if use_depthwise else ConvModule |
|
|
| assert (out_channels % num_heads == 0 and |
| embed_channels % num_heads == 0), \ |
| 'out_channels and embed_channels should be divisible by num_heads.' |
| self.num_heads = num_heads |
| self.head_channels = out_channels // num_heads |
|
|
| self.embed_conv = ConvModule( |
| in_channels, |
| embed_channels, |
| 1, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=None) if embed_channels != in_channels else None |
| self.guide_fc = Linear(guide_channels, embed_channels) |
| self.bias = nn.Parameter(torch.zeros(num_heads)) |
| if with_scale: |
| self.scale = nn.Parameter(torch.ones(1, num_heads, 1, 1)) |
| else: |
| self.scale = 1.0 |
|
|
| self.project_conv = conv(in_channels, |
| out_channels, |
| kernel_size, |
| stride=1, |
| padding=padding, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=None) |
|
|
| def forward(self, x: Tensor, guide: Tensor) -> Tensor: |
| """Forward process.""" |
| B, _, H, W = x.shape |
|
|
| guide = self.guide_fc(guide) |
| guide = guide.reshape(B, -1, self.num_heads, self.head_channels) |
| embed = self.embed_conv(x) if self.embed_conv is not None else x |
| embed = embed.reshape(B, self.num_heads, self.head_channels, H, W) |
|
|
| attn_weight = torch.einsum('bmchw,bnmc->bmhwn', embed, guide) |
| attn_weight = attn_weight.max(dim=-1)[0] |
| attn_weight = attn_weight / (self.head_channels**0.5) |
| attn_weight = attn_weight + self.bias[None, :, None, None] |
| attn_weight = attn_weight.sigmoid() * self.scale |
|
|
| x = self.project_conv(x) |
| x = x.reshape(B, self.num_heads, -1, H, W) |
| x = x * attn_weight.unsqueeze(2) |
| x = x.reshape(B, -1, H, W) |
| return x |
|
|
|
|
| @MODELS.register_module() |
| class MaxSigmoidCSPLayerWithTwoConv(CSPLayerWithTwoConv): |
| """Sigmoid-attention based CSP layer with two convolution layers.""" |
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| guide_channels: int, |
| embed_channels: int, |
| num_heads: int = 1, |
| expand_ratio: float = 0.5, |
| num_blocks: int = 1, |
| with_scale: bool = False, |
| add_identity: bool = True, |
| conv_cfg: OptConfigType = None, |
| norm_cfg: ConfigType = dict(type='BN', momentum=0.03, eps=0.001), |
| act_cfg: ConfigType = dict(type='SiLU', inplace=True), |
| init_cfg: OptMultiConfig = None) -> None: |
| super().__init__(in_channels=in_channels, |
| out_channels=out_channels, |
| expand_ratio=expand_ratio, |
| num_blocks=num_blocks, |
| add_identity=add_identity, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg, |
| init_cfg=init_cfg) |
|
|
| self.final_conv = ConvModule((3 + num_blocks) * self.mid_channels, |
| out_channels, |
| 1, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg) |
|
|
| self.attn_block = MaxSigmoidAttnBlock(self.mid_channels, |
| self.mid_channels, |
| guide_channels=guide_channels, |
| embed_channels=embed_channels, |
| num_heads=num_heads, |
| with_scale=with_scale, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg) |
|
|
| def forward(self, x: Tensor, guide: Tensor) -> Tensor: |
| """Forward process.""" |
| x_main = self.main_conv(x) |
| x_main = list(x_main.split((self.mid_channels, self.mid_channels), 1)) |
| x_main.extend(blocks(x_main[-1]) for blocks in self.blocks) |
| x_main.append(self.attn_block(x_main[-1], guide)) |
| return self.final_conv(torch.cat(x_main, 1)) |
|
|
|
|
| @MODELS.register_module() |
| class ImagePoolingAttentionModule(nn.Module): |
| def __init__(self, |
| image_channels: List[int], |
| text_channels: int, |
| embed_channels: int, |
| with_scale: bool = False, |
| num_feats: int = 3, |
| num_heads: int = 8, |
| pool_size: int = 3): |
| super().__init__() |
|
|
| self.text_channels = text_channels |
| self.embed_channels = embed_channels |
| self.num_heads = num_heads |
| self.num_feats = num_feats |
| self.head_channels = embed_channels // num_heads |
| self.pool_size = pool_size |
|
|
| if with_scale: |
| self.scale = nn.Parameter(torch.tensor([0.]), requires_grad=True) |
| else: |
| self.scale = 1.0 |
| self.projections = nn.ModuleList([ |
| ConvModule(in_channels, embed_channels, 1, act_cfg=None) |
| for in_channels in image_channels |
| ]) |
| self.query = nn.Sequential(nn.LayerNorm(text_channels), |
| Linear(text_channels, embed_channels)) |
| self.key = nn.Sequential(nn.LayerNorm(embed_channels), |
| Linear(embed_channels, embed_channels)) |
| self.value = nn.Sequential(nn.LayerNorm(embed_channels), |
| Linear(embed_channels, embed_channels)) |
| self.proj = Linear(embed_channels, text_channels) |
|
|
| self.image_pools = nn.ModuleList([ |
| nn.AdaptiveMaxPool2d((pool_size, pool_size)) |
| for _ in range(num_feats) |
| ]) |
|
|
| def forward(self, text_features, image_features): |
| B = image_features[0].shape[0] |
| assert len(image_features) == self.num_feats |
| num_patches = self.pool_size**2 |
| mlvl_image_features = [ |
| pool(proj(x)).view(B, -1, num_patches) |
| for (x, proj, pool |
| ) in zip(image_features, self.projections, self.image_pools) |
| ] |
| mlvl_image_features = torch.cat(mlvl_image_features, |
| dim=-1).transpose(1, 2) |
| q = self.query(text_features) |
| k = self.key(mlvl_image_features) |
| v = self.value(mlvl_image_features) |
|
|
| q = q.reshape(B, -1, self.num_heads, self.head_channels) |
| k = k.reshape(B, -1, self.num_heads, self.head_channels) |
| v = v.reshape(B, -1, self.num_heads, self.head_channels) |
|
|
| attn_weight = torch.einsum('bnmc,bkmc->bmnk', q, k) |
| attn_weight = attn_weight / (self.head_channels**0.5) |
| attn_weight = F.softmax(attn_weight, dim=-1) |
|
|
| x = torch.einsum('bmnk,bkmc->bnmc', attn_weight, v) |
| x = self.proj(x.reshape(B, -1, self.embed_channels)) |
| return x * self.scale + text_features |
|
|
|
|
| @MODELS.register_module() |
| class VanillaSigmoidBlock(BaseModule): |
| """Sigmoid attention block.""" |
| def __init__(self, |
| in_channels: int, |
| out_channels: int, |
| guide_channels: int, |
| embed_channels: int, |
| kernel_size: int = 3, |
| padding: int = 1, |
| num_heads: int = 1, |
| use_depthwise: bool = False, |
| with_scale: bool = False, |
| conv_cfg: OptConfigType = None, |
| norm_cfg: ConfigType = dict(type='BN', |
| momentum=0.03, |
| eps=0.001), |
| init_cfg: OptMultiConfig = None) -> None: |
| super().__init__(init_cfg=init_cfg) |
| conv = DepthwiseSeparableConvModule if use_depthwise else ConvModule |
|
|
| assert (out_channels % num_heads == 0 and |
| embed_channels % num_heads == 0), \ |
| 'out_channels and embed_channels should be divisible by num_heads.' |
| self.num_heads = num_heads |
| self.head_channels = out_channels // num_heads |
|
|
| self.project_conv = conv(in_channels, |
| out_channels, |
| kernel_size, |
| stride=1, |
| padding=padding, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=None) |
|
|
| def forward(self, x: Tensor, guide: Tensor) -> Tensor: |
| """Forward process.""" |
| x = self.project_conv(x) |
| x = x * x.sigmoid() |
| return x |
|
|
|
|
| @MODELS.register_module() |
| class EfficientCSPLayerWithTwoConv(CSPLayerWithTwoConv): |
| """Sigmoid-attention based CSP layer with two convolution layers.""" |
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| guide_channels: int, |
| embed_channels: int, |
| num_heads: int = 1, |
| expand_ratio: float = 0.5, |
| num_blocks: int = 1, |
| with_scale: bool = False, |
| add_identity: bool = True, |
| conv_cfg: OptConfigType = None, |
| norm_cfg: ConfigType = dict(type='BN', momentum=0.03, eps=0.001), |
| act_cfg: ConfigType = dict(type='SiLU', inplace=True), |
| init_cfg: OptMultiConfig = None) -> None: |
| super().__init__(in_channels=in_channels, |
| out_channels=out_channels, |
| expand_ratio=expand_ratio, |
| num_blocks=num_blocks, |
| add_identity=add_identity, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg, |
| init_cfg=init_cfg) |
|
|
| self.final_conv = ConvModule((3 + num_blocks) * self.mid_channels, |
| out_channels, |
| 1, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg) |
|
|
| self.attn_block = VanillaSigmoidBlock(self.mid_channels, |
| self.mid_channels, |
| guide_channels=guide_channels, |
| embed_channels=embed_channels, |
| num_heads=num_heads, |
| with_scale=with_scale, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg) |
|
|
| def forward(self, x: Tensor, guide: Tensor) -> Tensor: |
| """Forward process.""" |
| x_main = self.main_conv(x) |
| x_main = list(x_main.split((self.mid_channels, self.mid_channels), 1)) |
| x_main.extend(blocks(x_main[-1]) for blocks in self.blocks) |
| x_main.append(self.attn_block(x_main[-1], guide)) |
| return self.final_conv(torch.cat(x_main, 1)) |
|
|
