Spaces:
Build error
Build error
File size: 4,086 Bytes
cb433d6 |
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 |
import torch
import torch.nn as nn
from .transformer import PositionalEncoding
class Attention(nn.Module):
def __init__(self, in_channels=512, max_length=25, n_feature=256):
super().__init__()
self.max_length = max_length
self.f0_embedding = nn.Embedding(max_length, in_channels)
self.w0 = nn.Linear(max_length, n_feature)
self.wv = nn.Linear(in_channels, in_channels)
self.we = nn.Linear(in_channels, max_length)
self.active = nn.Tanh()
self.softmax = nn.Softmax(dim=2)
def forward(self, enc_output):
enc_output = enc_output.permute(0, 2, 3, 1).flatten(1, 2)
reading_order = torch.arange(self.max_length, dtype=torch.long, device=enc_output.device)
reading_order = reading_order.unsqueeze(0).expand(enc_output.size(0), -1) # (S,) -> (B, S)
reading_order_embed = self.f0_embedding(reading_order) # b,25,512
t = self.w0(reading_order_embed.permute(0, 2, 1)) # b,512,256
t = self.active(t.permute(0, 2, 1) + self.wv(enc_output)) # b,256,512
attn = self.we(t) # b,256,25
attn = self.softmax(attn.permute(0, 2, 1)) # b,25,256
g_output = torch.bmm(attn, enc_output) # b,25,512
return g_output, attn.view(*attn.shape[:2], 8, 32)
def encoder_layer(in_c, out_c, k=3, s=2, p=1):
return nn.Sequential(nn.Conv2d(in_c, out_c, k, s, p),
nn.BatchNorm2d(out_c),
nn.ReLU(True))
def decoder_layer(in_c, out_c, k=3, s=1, p=1, mode='nearest', scale_factor=None, size=None):
align_corners = None if mode=='nearest' else True
return nn.Sequential(nn.Upsample(size=size, scale_factor=scale_factor,
mode=mode, align_corners=align_corners),
nn.Conv2d(in_c, out_c, k, s, p),
nn.BatchNorm2d(out_c),
nn.ReLU(True))
class PositionAttention(nn.Module):
def __init__(self, max_length, in_channels=512, num_channels=64,
h=8, w=32, mode='nearest', **kwargs):
super().__init__()
self.max_length = max_length
self.k_encoder = nn.Sequential(
encoder_layer(in_channels, num_channels, s=(1, 2)),
encoder_layer(num_channels, num_channels, s=(2, 2)),
encoder_layer(num_channels, num_channels, s=(2, 2)),
encoder_layer(num_channels, num_channels, s=(2, 2))
)
self.k_decoder = nn.Sequential(
decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
decoder_layer(num_channels, in_channels, size=(h, w), mode=mode)
)
self.pos_encoder = PositionalEncoding(in_channels, dropout=0, max_len=max_length)
self.project = nn.Linear(in_channels, in_channels)
def forward(self, x):
N, E, H, W = x.size()
k, v = x, x # (N, E, H, W)
# calculate key vector
features = []
for i in range(0, len(self.k_encoder)):
k = self.k_encoder[i](k)
features.append(k)
for i in range(0, len(self.k_decoder) - 1):
k = self.k_decoder[i](k)
k = k + features[len(self.k_decoder) - 2 - i]
k = self.k_decoder[-1](k)
# calculate query vector
# TODO q=f(q,k)
zeros = x.new_zeros((self.max_length, N, E)) # (T, N, E)
q = self.pos_encoder(zeros) # (T, N, E)
q = q.permute(1, 0, 2) # (N, T, E)
q = self.project(q) # (N, T, E)
# calculate attention
attn_scores = torch.bmm(q, k.flatten(2, 3)) # (N, T, (H*W))
attn_scores = attn_scores / (E ** 0.5)
attn_scores = torch.softmax(attn_scores, dim=-1)
v = v.permute(0, 2, 3, 1).view(N, -1, E) # (N, (H*W), E)
attn_vecs = torch.bmm(attn_scores, v) # (N, T, E)
return attn_vecs, attn_scores.view(N, -1, H, W)
|