156 lines
4.9 KiB
Python
Executable File
156 lines
4.9 KiB
Python
Executable File
import math
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class SelfAttention(nn.Module):
|
|
def __init__(
|
|
self, dim, heads=8, qkv_bias=False, qk_scale=None, dropout_rate=0.0
|
|
):
|
|
super().__init__()
|
|
self.num_heads = heads
|
|
head_dim = dim // heads
|
|
self.scale = qk_scale or head_dim ** -0.5
|
|
|
|
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
|
self.attn_drop = nn.Dropout(dropout_rate)
|
|
self.proj = nn.Linear(dim, dim)
|
|
self.proj_drop = nn.Dropout(dropout_rate)
|
|
|
|
def forward(self, x):
|
|
B, N, C = x.shape
|
|
qkv = (
|
|
self.qkv(x)
|
|
.reshape(B, N, 3, self.num_heads, C // self.num_heads)
|
|
.permute(2, 0, 3, 1, 4)
|
|
)
|
|
q, k, v = (
|
|
qkv[0],
|
|
qkv[1],
|
|
qkv[2],
|
|
) # make torchscript happy (cannot use tensor as tuple)
|
|
|
|
attn = (q @ k.transpose(-2, -1)) * self.scale
|
|
attn = attn.softmax(dim=-1)
|
|
attn = self.attn_drop(attn)
|
|
|
|
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
|
|
x = self.proj(x)
|
|
x = self.proj_drop(x)
|
|
return x
|
|
|
|
|
|
class AxialAttention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_planes,
|
|
out_planes,
|
|
groups=8,
|
|
kernel_size=56,
|
|
stride=1,
|
|
bias=False,
|
|
width=False,
|
|
):
|
|
assert (in_planes % groups == 0) and (out_planes % groups == 0)
|
|
super(AxialAttention, self).__init__()
|
|
self.in_planes = in_planes
|
|
self.out_planes = out_planes
|
|
self.groups = groups
|
|
self.group_planes = out_planes // groups
|
|
self.kernel_size = kernel_size
|
|
self.stride = stride
|
|
self.bias = bias
|
|
self.width = width
|
|
|
|
# Multi-head self attention
|
|
self.qkv_transform = nn.Conv1d(
|
|
in_planes,
|
|
out_planes * 2,
|
|
kernel_size=1,
|
|
stride=1,
|
|
padding=0,
|
|
bias=False,
|
|
)
|
|
self.bn_qkv = nn.BatchNorm1d(out_planes * 2)
|
|
self.bn_similarity = nn.BatchNorm2d(groups * 3)
|
|
self.bn_output = nn.BatchNorm1d(out_planes * 2)
|
|
|
|
# Position embedding
|
|
self.relative = nn.Parameter(
|
|
torch.randn(self.group_planes * 2, kernel_size * 2 - 1),
|
|
requires_grad=True,
|
|
)
|
|
query_index = torch.arange(kernel_size).unsqueeze(0)
|
|
key_index = torch.arange(kernel_size).unsqueeze(1)
|
|
relative_index = key_index - query_index + kernel_size - 1
|
|
self.register_buffer('flatten_index', relative_index.view(-1))
|
|
if stride > 1:
|
|
self.pooling = nn.AvgPool2d(stride, stride=stride)
|
|
|
|
self.reset_parameters()
|
|
|
|
def forward(self, x):
|
|
if self.width:
|
|
x = x.permute(0, 2, 1, 3)
|
|
else:
|
|
x = x.permute(0, 3, 1, 2) # N, W, C, H
|
|
N, W, C, H = x.shape
|
|
x = x.contiguous().view(N * W, C, H)
|
|
|
|
# Transformations
|
|
qkv = self.bn_qkv(self.qkv_transform(x))
|
|
q, k, v = torch.split(
|
|
qkv.reshape(N * W, self.groups, self.group_planes * 2, H),
|
|
[self.group_planes // 2, self.group_planes // 2, self.group_planes],
|
|
dim=2,
|
|
)
|
|
|
|
# Calculate position embedding
|
|
all_embeddings = torch.index_select(
|
|
self.relative, 1, self.flatten_index
|
|
).view(self.group_planes * 2, self.kernel_size, self.kernel_size)
|
|
q_embedding, k_embedding, v_embedding = torch.split(
|
|
all_embeddings,
|
|
[self.group_planes // 2, self.group_planes // 2, self.group_planes],
|
|
dim=0,
|
|
)
|
|
qr = torch.einsum('bgci,cij->bgij', q, q_embedding)
|
|
kr = torch.einsum('bgci,cij->bgij', k, k_embedding).transpose(2, 3)
|
|
qk = torch.einsum('bgci, bgcj->bgij', q, k)
|
|
stacked_similarity = torch.cat([qk, qr, kr], dim=1)
|
|
stacked_similarity = (
|
|
self.bn_similarity(stacked_similarity)
|
|
.view(N * W, 3, self.groups, H, H)
|
|
.sum(dim=1)
|
|
)
|
|
|
|
# (N, groups, H, H, W)
|
|
similarity = F.softmax(stacked_similarity, dim=3)
|
|
sv = torch.einsum('bgij,bgcj->bgci', similarity, v)
|
|
sve = torch.einsum('bgij,cij->bgci', similarity, v_embedding)
|
|
stacked_output = torch.cat([sv, sve], dim=-1).view(
|
|
N * W, self.out_planes * 2, H
|
|
)
|
|
output = (
|
|
self.bn_output(stacked_output)
|
|
.view(N, W, self.out_planes, 2, H)
|
|
.sum(dim=-2)
|
|
)
|
|
|
|
if self.width:
|
|
output = output.permute(0, 2, 1, 3)
|
|
else:
|
|
output = output.permute(0, 2, 3, 1)
|
|
|
|
if self.stride > 1:
|
|
output = self.pooling(output)
|
|
|
|
return output
|
|
|
|
def reset_parameters(self):
|
|
self.qkv_transform.weight.data.normal_(
|
|
0, math.sqrt(1.0 / self.in_planes)
|
|
)
|
|
nn.init.normal_(self.relative, 0.0, math.sqrt(1.0 / self.group_planes))
|