67 lines
2.6 KiB
Python
67 lines
2.6 KiB
Python
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
from torch import nn
|
|||
|
|
from torch.nn import init
|
|||
|
|
import math
|
|||
|
|
from torch.nn import functional as F
|
|||
|
|
|
|||
|
|
|
|||
|
|
class OutlookAttention(nn.Module):
|
|||
|
|
|
|||
|
|
def __init__(self, dim, num_heads=1, kernel_size=3, padding=1, stride=1, qkv_bias=False,
|
|||
|
|
attn_drop=0.1):
|
|||
|
|
super().__init__()
|
|||
|
|
self.dim = dim
|
|||
|
|
self.num_heads = num_heads
|
|||
|
|
self.head_dim = dim // num_heads
|
|||
|
|
self.kernel_size = kernel_size
|
|||
|
|
self.padding = padding
|
|||
|
|
self.stride = stride
|
|||
|
|
self.scale = self.head_dim ** (-0.5)
|
|||
|
|
|
|||
|
|
self.v_pj = nn.Linear(dim, dim, bias=qkv_bias)
|
|||
|
|
self.attn = nn.Linear(dim, kernel_size ** 4 * num_heads)
|
|||
|
|
|
|||
|
|
self.attn_drop = nn.Dropout(attn_drop)
|
|||
|
|
self.proj = nn.Linear(dim, dim)
|
|||
|
|
self.proj_drop = nn.Dropout(attn_drop)
|
|||
|
|
|
|||
|
|
self.unflod = nn.Unfold(kernel_size, padding, stride) # 手动卷积
|
|||
|
|
self.pool = nn.AvgPool2d(kernel_size=stride, stride=stride, ceil_mode=True)
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
B, H, W, C = x.shape
|
|||
|
|
|
|||
|
|
# 映射到新的特征v
|
|||
|
|
v = self.v_pj(x).permute(0, 3, 1, 2) # B,C,H,W
|
|||
|
|
h, w = math.ceil(H / self.stride), math.ceil(W / self.stride)
|
|||
|
|
v = self.unflod(v).reshape(B, self.num_heads, self.head_dim, self.kernel_size * self.kernel_size,
|
|||
|
|
h * w).permute(0, 1, 4, 3, 2) # B,num_head,H*W,kxk,head_dim
|
|||
|
|
|
|||
|
|
# 生成Attention Map
|
|||
|
|
attn = self.pool(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1) # B,H,W,C
|
|||
|
|
attn = self.attn(attn).reshape(B, h * w, self.num_heads, self.kernel_size * self.kernel_size \
|
|||
|
|
, self.kernel_size * self.kernel_size).permute(0, 2, 1, 3,
|
|||
|
|
4) # B，num_head，H*W,kxk,kxk
|
|||
|
|
attn = self.scale * attn
|
|||
|
|
attn = attn.softmax(-1)
|
|||
|
|
attn = self.attn_drop(attn)
|
|||
|
|
|
|||
|
|
# 获取weighted特征
|
|||
|
|
out = (attn @ v).permute(0, 1, 4, 3, 2).reshape(B, C * self.kernel_size * self.kernel_size,
|
|||
|
|
h * w) # B,dimxkxk,H*W
|
|||
|
|
out = F.fold(out, output_size=(H, W), kernel_size=self.kernel_size,
|
|||
|
|
padding=self.padding, stride=self.stride) # B,C,H,W
|
|||
|
|
out = self.proj(out.permute(0, 2, 3, 1)) # B,H,W,C
|
|||
|
|
out = self.proj_drop(out)
|
|||
|
|
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 输入 B, H, W, C, 输出 B, H, W, C
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
block = OutlookAttention(dim=512).cuda()
|
|||
|
|
input = torch.rand(1, 64, 64, 512).cuda()
|
|||
|
|
output = block(input)
|
|||
|
|
print(input.size(), output.size())
|