132 lines
4.8 KiB
Python
132 lines
4.8 KiB
Python
import math
|
|
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from einops import rearrange
|
|
|
|
from .transformer_modules import ScaledDotProductAttention
|
|
|
|
|
|
class MultiHeadAttention(nn.Module):
|
|
def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1, granular=True):
|
|
super().__init__()
|
|
self.n_head = n_head
|
|
self.d_k = d_k
|
|
self.d_v = d_v
|
|
self.query_embedding = nn.Conv2d(512, 512, kernel_size=1, padding=0)
|
|
self.value_embedding = nn.Conv2d(512, 512, kernel_size=1, padding=0)
|
|
self.key_embedding = nn.Conv2d(512, 512, kernel_size=1, padding=0)
|
|
|
|
self.granular = granular
|
|
|
|
# option b
|
|
self.w_qs_h = nn.Linear(d_model, n_head * d_k, bias=False)
|
|
self.w_ks_h = nn.Linear(d_model, n_head * d_k, bias=False)
|
|
self.w_vs_h = nn.Linear(d_model, n_head * d_v, bias=False)
|
|
self.fc_h = nn.Linear(n_head * d_v, d_model, bias=False)
|
|
|
|
self.attention_h = ScaledDotProductAttention(temperature=d_k**0.5)
|
|
|
|
self.w_qs_w = nn.Linear(d_model, n_head * d_k, bias=False)
|
|
self.w_ks_w = nn.Linear(d_model, n_head * d_k, bias=False)
|
|
self.w_vs_w = nn.Linear(d_model, n_head * d_v, bias=False)
|
|
self.fc_w = nn.Linear(n_head * d_v, d_model, bias=False)
|
|
|
|
self.attention_w = ScaledDotProductAttention(temperature=d_k**0.5)
|
|
|
|
self.dropout1 = nn.Dropout(dropout)
|
|
self.dropout2 = nn.Dropout(dropout)
|
|
self.layer_norm1 = nn.LayerNorm(d_model, eps=1e-6)
|
|
self.layer_norm2 = nn.LayerNorm(d_model, eps=1e-6)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
|
|
def forward(self, q, k, v, mask=None):
|
|
residual = q
|
|
bt, l, _ = q.size() # noqa: E741
|
|
|
|
q = rearrange(q, "b (m1 m2) c -> b c m1 m2", m1=int(math.sqrt(l)))
|
|
k = rearrange(k, "b (m1 m2) c -> b c m1 m2", m1=int(math.sqrt(l)))
|
|
v = rearrange(v, "b (m1 m2) c -> b c m1 m2", m1=int(math.sqrt(l)))
|
|
|
|
q = self.query_embedding(q)
|
|
k_prev = self.key_embedding(k)
|
|
v = self.value_embedding(v)
|
|
|
|
q = rearrange(q, "b c m1 m2 -> (b m1) m2 c")
|
|
k = rearrange(k_prev, "b c m1 m2 -> (b m1) m2 c")
|
|
v = rearrange(v, "b c m1 m2 -> (b m1) m2 c")
|
|
|
|
d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
|
|
q = self.w_qs_h(q)
|
|
k = self.w_ks_h(k)
|
|
v = self.w_vs_h(v)
|
|
|
|
q = rearrange(q, "b len (head d_k) -> b len head d_k", d_k=d_k, head=n_head)
|
|
k = rearrange(k, "b len (head d_k) -> b len head d_k", d_k=d_k, head=n_head)
|
|
v = rearrange(v, "b len (head d_v) -> b len head d_v", d_v=d_v, head=n_head)
|
|
|
|
sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)
|
|
|
|
q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
|
|
|
|
if mask is not None:
|
|
mask = mask.unsqueeze(1)
|
|
|
|
q, attn = self.attention_h(q, k, v, mask=mask)
|
|
|
|
q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
|
|
q = self.dropout1(self.fc_h(q))
|
|
q = rearrange(q, "(bt h) w c ->bt (h w) c", h=int(math.sqrt(l)))
|
|
q_mid = self.layer_norm1(q)
|
|
|
|
q = rearrange(q_mid, "bt (h w) c ->(bt w) h c", h=int(math.sqrt(l)))
|
|
if self.granular:
|
|
k = rearrange(q_mid,'bt (h w) c ->(bt w) h c',h =int(math.sqrt(l)))
|
|
else:
|
|
k = rearrange(k_prev, "b c m1 m2 -> (b m2) m1 c")
|
|
v = rearrange(q_mid, "bt (h w) c ->(bt w) h c", h=int(math.sqrt(l)))
|
|
|
|
d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
|
|
q = self.w_qs_w(q)
|
|
k = self.w_ks_w(k)
|
|
v = self.w_vs_w(v)
|
|
|
|
q = rearrange(q, "b len (head d_k) -> b len head d_k", d_k=d_k, head=n_head)
|
|
k = rearrange(k, "b len (head d_k) -> b len head d_k", d_k=d_k, head=n_head)
|
|
v = rearrange(v, "b len (head d_v) -> b len head d_v", d_v=d_v, head=n_head)
|
|
|
|
sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1) # noqa: F841
|
|
|
|
q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
|
|
|
|
if mask is not None:
|
|
mask = mask.unsqueeze(1)
|
|
|
|
q, attn = self.attention_w(q, k, v, mask=mask)
|
|
q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
|
|
q = self.dropout2(self.fc_w(q))
|
|
q = rearrange(q, "(bt w) h c -> bt (h w) c", w=int(math.sqrt(l)))
|
|
q += residual
|
|
q = self.layer_norm2(q)
|
|
|
|
return q, attn
|
|
|
|
|
|
class PositionwiseFeedForward(nn.Module):
|
|
def __init__(self, d_in, d_hid, dropout=0.1):
|
|
super().__init__()
|
|
self.w_1 = nn.Linear(d_in, d_hid)
|
|
self.w_2 = nn.Linear(d_hid, d_in)
|
|
self.layer_norm = nn.LayerNorm(d_in, eps=1e-6)
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
def forward(self, x):
|
|
residual = x
|
|
x = self.w_2(F.relu(self.w_1(x)))
|
|
x = self.dropout(x)
|
|
x += residual
|
|
|
|
x = self.layer_norm(x)
|
|
|
|
return x
|