CVPR 2024 | 最新即插即用注意力机制+ffn层

import torchimport torch.nn as nnfrom timm.models.layers import trunc_normal_from einops import repeat
class LinearProjection(nn.Module):    def __init__(self, dim, heads=8, dim_head=64, bias=True):        super().__init__()        inner_dim = dim_head * heads        self.heads = heads        self.to_q = nn.Linear(dim, inner_dim, bias=bias)        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=bias)        self.dim = dim        self.inner_dim = inner_dim
    def forward(self, x, attn_kv=None):        B_, N, C = x.shape        if attn_kv is not None:            attn_kv = attn_kv.unsqueeze(0).repeat(B_, 1, 1)        else:            attn_kv = x        N_kv = attn_kv.size(1)        q = self.to_q(x).reshape(B_, N, 1, self.heads, C // self.heads).permute(2, 0, 3, 1, 4)        kv = self.to_kv(attn_kv).reshape(B_, N_kv, 2, self.heads, C // self.heads).permute(2, 0, 3, 1, 4)        q = q[0]        k, v = kv[0], kv[1]        return q, k, v
class WindowAttention_sparse(nn.Module):    def __init__(self, dim, win_size, num_heads=8, token_projection='linear', qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):        super().__init__()        self.dim = dim        self.win_size = win_size  # Wh, Ww        self.num_heads = num_heads        head_dim = dim // num_heads        self.scale = qk_scale or head_dim ** -0.5        # define a parameter table of relative position bias        self.relative_position_bias_table = nn.Parameter(            torch.zeros((2 * win_size[0] - 1) * (2 * win_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH        # get pair-wise relative position index for each token inside the window        coords_h = torch.arange(self.win_size[0])  # [0,...,Wh-1]        coords_w = torch.arange(self.win_size[1])  # [0,...,Ww-1]        coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij'))  # 2, Wh, Ww        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2        relative_coords[:, :, 0] += self.win_size[0] - 1  # shift to start from 0        relative_coords[:, :, 1] += self.win_size[1] - 1        relative_coords[:, :, 0] *= 2 * self.win_size[1] - 1        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww        self.register_buffer("relative_position_index", relative_position_index)        trunc_normal_(self.relative_position_bias_table, std=.02)        if token_projection == 'linear':            self.qkv = LinearProjection(dim, num_heads, dim // num_heads, bias=qkv_bias)        else:            raise Exception("Projection error!")        self.token_projection = token_projection        self.attn_drop = nn.Dropout(attn_drop)        self.proj = nn.Linear(dim, dim)        self.proj_drop = nn.Dropout(proj_drop)        self.softmax = nn.Softmax(dim=-1)        self.relu = nn.ReLU()        self.w = nn.Parameter(torch.ones(2))  # 自适应权重参数
    def forward(self, x, attn_kv=None, mask=None):        # 调整输入维度，从 (B, C, H, W) 转为 (B, H, W, C)        x = x.permute(0, 2, 3, 1).reshape(x.shape[0], x.shape[2] * x.shape[3], x.shape[1])        B_, N, C = x.shape        q, k, v = self.qkv(x, attn_kv)        q = q * self.scale        attn = (q @ k.transpose(-2, -1))        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(            self.win_size[0] * self.win_size[1], self.win_size[0] * self.win_size[1], -1)  # Wh*Ww,Wh*Ww,nH        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww        ratio = attn.size(-1) // relative_position_bias.size(-1)        relative_position_bias = repeat(relative_position_bias, 'nH l c -> nH l (c d)', d=ratio)        attn = attn + relative_position_bias.unsqueeze(0)        if mask is not None:            nW = mask.shape[0]            mask = repeat(mask, 'nW m n -> nW m (n d)', d=ratio)            attn = attn.view(B_ // nW, nW, self.num_heads, N, N * ratio) + mask.unsqueeze(1).unsqueeze(0)            attn = attn.view(-1, self.num_heads, N, N * ratio)        attn0 = self.softmax(attn)        attn1 = self.relu(attn) ** 2  # b,h,w,c        w1 = torch.exp(self.w[0]) / torch.sum(torch.exp(self.w))        w2 = torch.exp(self.w[1]) / torch.sum(torch.exp(self.w))        attn = attn0 * w1 + attn1 * w2        attn = self.attn_drop(attn)        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)        x = self.proj(x)        x = self.proj_drop(x)        x = x.reshape(x.shape[0], int(math.sqrt(x.shape[1])), int(math.sqrt(x.shape[1])), x.shape[2]).permute(0, 3, 1, 2)        return x
    def extra_repr(self) -> str:        return f'dim={self.dim}, win_size={self.win_size}, num_heads={self.num_heads}'

import torchimport torch.nn as nnfrom einops import rearrange
class FRFN(nn.Module):    def __init__(self, dim=32, hidden_dim=128, act_layer=nn.GELU, drop=0., use_eca=False):        super().__init__()        self.linear1 = nn.Sequential(nn.Linear(dim, hidden_dim * 2), act_layer())        self.dwconv = nn.Sequential(            nn.Conv2d(hidden_dim, hidden_dim, groups=hidden_dim, kernel_size=3, stride=1, padding=1),            act_layer()        )        self.linear2 = nn.Sequential(nn.Linear(hidden_dim, dim))        self.dim = dim        self.hidden_dim = hidden_dim        self.dim_conv = self.dim // 4        self.dim_untouched = self.dim - self.dim_conv        self.partial_conv3 = nn.Conv2d(self.dim_conv, self.dim_conv, 3, 1, 1, bias=False)
    def forward(self, x):        c, bs, hh, hw = x.size()        x1, x2 = torch.split(x, [self.dim_conv, self.dim_untouched], dim=1)        x1 = self.partial_conv3(x1)        x = torch.cat((x1, x2), 1)        x = rearrange(x, 'b c h w -> b (h w) c', h=hh, w=hw)        x = self.linear1(x)        x_1, x_2 = x.chunk(2, dim=-1)        x_1 = rearrange(x_1, 'b (h w) c -> b c h w', h=hh, w=hw)        x_1 = self.dwconv(x_1)        x_1 = rearrange(x_1, 'b c h w -> b (h w) c', h=hh, w=hw)        x = x_1 * x_2        x = self.linear2(x)        return rearrange(x, 'b (h w) c -> b c h w', h=hh, w=hw)