Spaces:
Build error
Build error
File size: 6,433 Bytes
711b041 |
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 |
import torch.nn as nn
from .trm import *
class _MultiHeadAttention(nn.Module):
def __init__(self, d_k, d_v, d_model, n_heads, dropout):
super(_MultiHeadAttention, self).__init__()
self.d_k = d_k
self.d_v = d_v
self.d_model = d_model
self.n_heads = n_heads
self.w_q = Linear(d_model, d_k * n_heads)
self.w_k = Linear(d_model, d_k * n_heads)
self.w_v = Linear(d_model, d_v * n_heads)
def forward(self, q, k, v):
# q: [b_size x len_q x d_model]
# k: [b_size x len_k x d_model]
# v: [b_size x len_k x d_model]
b_size = q.size(0)
# q_s: [b_size x n_heads x len_q x d_k]
# k_s: [b_size x n_heads x len_k x d_k]
# v_s: [b_size x n_heads x len_k x d_v]
q_s = self.w_q(q).view(b_size, -1, self.n_heads, self.d_k).transpose(1, 2)
k_s = self.w_k(k).view(b_size, -1, self.n_heads, self.d_k).transpose(1, 2)
v_s = self.w_v(v).view(b_size, -1, self.n_heads, self.d_v).transpose(1, 2)
return q_s, k_s, v_s
class PoswiseFeedForwardNet(nn.Module):
def __init__(self, d_model, d_ff, dropout=0.1):
super(PoswiseFeedForwardNet, self).__init__()
self.relu = nn.ReLU()
self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
self.dropout = nn.Dropout(dropout)
self.layer_norm = LayerNormalization(d_model)
def forward(self, inputs):
# inputs: [b_size x len_q x d_model]
residual = inputs
output = self.relu(self.conv1(inputs.transpose(1, 2)))
# outputs: [b_size x len_q x d_model]
output = self.conv2(output).transpose(1, 2)
output = self.dropout(output)
return self.layer_norm(residual + output)
class MultiHeadAttention(nn.Module):
def __init__(self, d_k, d_v, n_heads, dropout, d_model, visual_len, sen_len, fea_v, fea_s, pos):
super(MultiHeadAttention, self).__init__()
self.n_heads = n_heads
self.multihead_attn_v = _MultiHeadAttention(d_k, d_v, d_model, n_heads, dropout)
self.multihead_attn_s = _MultiHeadAttention(d_k, d_v, d_model, n_heads, dropout)
self.pos_emb_v = PosEncoding(visual_len * 10, d_model)
self.pos_emb_s = PosEncoding(sen_len * 10, d_model)
self.linear_v = nn.Linear(in_features=fea_v, out_features=d_model)
self.linear_s = nn.Linear(in_features=fea_s, out_features=d_model)
self.proj_v = Linear(n_heads * d_v, d_model)
self.proj_s = Linear(n_heads * d_v, d_model)
self.d_v = d_v
self.dropout = nn.Dropout(dropout)
self.layer_norm_v = LayerNormalization(d_model)
self.layer_norm_s = LayerNormalization(d_model)
self.attention = ScaledDotProductAttention(d_k, dropout)
self.pos = pos
def forward(self, v, s, v_len, s_len):
b_size = v.size(0)
# q: [b_size x len_q x d_model]
# k: [b_size x len_k x d_model]
# v: [b_size x len_v x d_model] note (len_k == len_v)
v, s = self.linear_v(v), self.linear_s(s)
if self.pos:
pos_v, pos_s = self.pos_emb_v(v_len), self.pos_emb_s(s_len)
residual_v, residual_s = v + pos_v, s + pos_s
else:
residual_v, residual_s = v, s
# context: a tensor of shape [b_size x len_q x n_heads * d_v]
q_v, k_v, v_v = self.multihead_attn_v(v, v, v)
q_s, k_s, v_s = self.multihead_attn_s(s, s, s)
context_v, attn_v = self.attention(q_v, k_s, v_s)
context_s, attn_s = self.attention(q_s, k_v, v_v)
context_v = context_v.transpose(1, 2).contiguous().view(b_size, -1, self.n_heads * self.d_v)
context_s = context_s.transpose(1, 2).contiguous().view(b_size, -1, self.n_heads * self.d_v)
# project back to the residual size, outputs: [b_size x len_q x d_model]
output_v = self.dropout(self.proj_v(context_v))
output_s = self.dropout(self.proj_s(context_s))
return self.layer_norm_v(residual_v + output_v), self.layer_norm_s(residual_s + output_s)
class co_attention(nn.Module):
def __init__(self, d_k, d_v, n_heads, dropout, d_model, visual_len, sen_len, fea_v, fea_s, pos):
super(co_attention, self).__init__()
# self.layer_num = layer_num
# self.multi_head = MultiHeadAttention(d_k=d_k, d_v=d_v, n_heads=n_heads, dropout=dropout, d_model=d_model,
# visual_len=visual_len, sen_len=sen_len, fea_v=fea_v, fea_s=fea_s, pos=False)
# self.PoswiseFeedForwardNet_v = nn.ModuleList([PoswiseFeedForwardNet(d_model=d_model, d_ff=256)])
# self.PoswiseFeedForwardNet_s = nn.ModuleList([PoswiseFeedForwardNet(d_model=d_model, d_ff=256)])
# self.multi_head = nn.ModuleList([MultiHeadAttention(d_k=d_k, d_v=d_v, n_heads=n_heads, dropout=dropout, d_model=d_model,
# visual_len=visual_len, sen_len=sen_len, fea_v=fea_v, fea_s=fea_s, pos=False)])
# for i in range(1, layer_num):
# self.PoswiseFeedForwardNet_v.append(PoswiseFeedForwardNet(d_model=d_model, d_ff=256))
# self.PoswiseFeedForwardNet_s.append(PoswiseFeedForwardNet(d_model=d_model, d_ff=256))
# self.multi_head.append(MultiHeadAttention(d_k=d_k, d_v=d_v, n_heads=n_heads, dropout=dropout, d_model=d_model,
# visual_len=visual_len, sen_len=sen_len, fea_v=d_model, fea_s=d_model, pos=True))
self.multi_head = MultiHeadAttention(d_k=d_k, d_v=d_v, n_heads=n_heads, dropout=dropout, d_model=d_model,
visual_len=visual_len, sen_len=sen_len, fea_v=fea_v, fea_s=fea_s, pos=pos)
self.PoswiseFeedForwardNet_v = PoswiseFeedForwardNet(d_model=d_model, d_ff=128, dropout=dropout)
self.PoswiseFeedForwardNet_s = PoswiseFeedForwardNet(d_model=d_model, d_ff=128,dropout=dropout)
def forward(self, v, s, v_len, s_len):
# for i in range(self.layer_num):
# v, s = self.multi_head[i](v, s, v_len, s_len)
# v = self.PoswiseFeedForwardNet_v[i](v)
# s = self.PoswiseFeedForwardNet_s[i](s)
v, s = self.multi_head(v, s, v_len, s_len)
v = self.PoswiseFeedForwardNet_v(v)
s = self.PoswiseFeedForwardNet_s(s)
return v, s
|