File size: 5,306 Bytes
25986db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import trunc_normal_
from lib.models.layers.rpe import generate_2d_concatenated_self_attention_relative_positional_encoding_index
class Attention(nn.Module):
def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.,
rpe=False, z_size=7, x_size=14):
super().__init__()
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = head_dim ** -0.5
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
self.rpe =rpe
if self.rpe:
relative_position_index = \
generate_2d_concatenated_self_attention_relative_positional_encoding_index([z_size, z_size],
[x_size, x_size])
self.register_buffer("relative_position_index", relative_position_index)
# define a parameter table of relative position bias
self.relative_position_bias_table = nn.Parameter(torch.empty((num_heads,
relative_position_index.max() + 1)))
trunc_normal_(self.relative_position_bias_table, std=0.02)
def forward(self, x, mask=None, return_attention=False):
# x: B, N, C
# mask: [B, N, ] torch.bool
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.unbind(0) # make torchscript happy (cannot use tensor as tuple)
attn = (q @ k.transpose(-2, -1)) * self.scale
if self.rpe:
relative_position_bias = self.relative_position_bias_table[:, self.relative_position_index].unsqueeze(0)
attn += relative_position_bias
if mask is not None:
attn = attn.masked_fill(mask.unsqueeze(1).unsqueeze(2), float('-inf'),)
split_attn = False
len_t = 49
if split_attn:
attn_t = attn[..., :len_t].softmax(dim=-1)
attn_s = attn[..., len_t:].softmax(dim=-1)
attn = torch.cat([attn_t, attn_s], dim=-1)
else:
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)
if return_attention:
return x, attn
else:
return x
class Attention_talking_head(nn.Module):
# taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
# with slight modifications to add Talking Heads Attention (https://arxiv.org/pdf/2003.02436v1.pdf)
def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.,
rpe=True, z_size=7, x_size=14):
super().__init__()
self.num_heads = num_heads
head_dim = dim // num_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(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_l = nn.Linear(num_heads, num_heads)
self.proj_w = nn.Linear(num_heads, num_heads)
self.proj_drop = nn.Dropout(proj_drop)
self.rpe = rpe
if self.rpe:
relative_position_index = \
generate_2d_concatenated_self_attention_relative_positional_encoding_index([z_size, z_size],
[x_size, x_size])
self.register_buffer("relative_position_index", relative_position_index)
# define a parameter table of relative position bias
self.relative_position_bias_table = nn.Parameter(torch.empty((num_heads,
relative_position_index.max() + 1)))
trunc_normal_(self.relative_position_bias_table, std=0.02)
def forward(self, x, mask=None):
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] * self.scale, qkv[1], qkv[2]
attn = (q @ k.transpose(-2, -1))
if self.rpe:
relative_position_bias = self.relative_position_bias_table[:, self.relative_position_index].unsqueeze(0)
attn += relative_position_bias
if mask is not None:
attn = attn.masked_fill(mask.unsqueeze(1).unsqueeze(2),
float('-inf'),)
attn = self.proj_l(attn.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
attn = attn.softmax(dim=-1)
attn = self.proj_w(attn.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
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 |