def __init__(self, args):
super().__init__()
self.embed_dim = args.decoder_embed_dim
self.dropout = args.dropout
self.relu_dropout = args.relu_dropout
self.self_attn = MultiheadAttention(self.embed_dim,
args.decoder_attention_heads,
args.attention_dropout)
self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
self.encoder_attn = MultiheadAttention(self.embed_dim,
args.decoder_attention_heads,
args.attention_dropout)
self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
self.fc1 = nn.Linear(self.embed_dim, args.decoder_ffn_embed_dim)
self.fc2 = nn.Linear(args.decoder_ffn_embed_dim, self.embed_dim)
self.final_layer_norm = nn.LayerNorm(self.embed_dim)
self.reset_parameters()
self.v_proj = nn.Linear(2048, self.embed_dim)
nn.init.kaiming_normal_(self.v_proj.weight)
nn.init.constant_(self.v_proj.bias, 0.)
self.encoder_attn_v = MultiheadAttention(self.embed_dim,
args.decoder_attention_heads,
args.attention_dropout)
self.ha_attn = HierarchicalAttention(self.embed_dim, self.embed_dim)
self.vt_layer_norm = nn.LayerNorm(self.embed_dim)
def __init__(self, args):
super().__init__()
self.embed_dim = args.decoder_embed_dim
self.dropout = args.dropout
self.relu_dropout = args.relu_dropout
self.self_attn = MultiheadAttention(self.embed_dim,
args.decoder_attention_heads,
args.attention_dropout)
self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
self.encoder_attn = MultiheadAttention(self.embed_dim,
args.decoder_attention_heads,
args.attention_dropout)
self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
self.fc1 = nn.Linear(self.embed_dim, args.decoder_ffn_embed_dim)
self.fc2 = nn.Linear(args.decoder_ffn_embed_dim, self.embed_dim)
self.final_layer_norm = nn.LayerNorm(self.embed_dim)
def __init__(self, args):
super().__init__()
self.embed_dim = args.encoder_embed_dim
self.dropout = args.dropout
self.relu_dropout = args.relu_dropout
self.self_attn = MultiheadAttention(self.embed_dim,
args.encoder_attention_heads,
args.attention_dropout)
self.fc1 = nn.Linear(self.embed_dim, args.encoder_ffn_embed_dim)
self.fc2 = nn.Linear(args.encoder_ffn_embed_dim, self.embed_dim)
self.layer_norms = nn.ModuleList(
[nn.LayerNorm(self.embed_dim) for _ in range(2)])
self.reset_parameters()
class TransformerDecoderLayer(nn.Module):
def __init__(self, args):
super().__init__()
self.embed_dim = args.decoder_embed_dim
self.dropout = args.dropout
self.relu_dropout = args.relu_dropout
self.self_attn = MultiheadAttention(self.embed_dim,
args.decoder_attention_heads,
args.attention_dropout)
self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
self.encoder_attn = MultiheadAttention(self.embed_dim,
args.decoder_attention_heads,
args.attention_dropout)
self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
self.fc1 = nn.Linear(self.embed_dim, args.decoder_ffn_embed_dim)
self.fc2 = nn.Linear(args.decoder_ffn_embed_dim, self.embed_dim)
self.final_layer_norm = nn.LayerNorm(self.embed_dim)
self.reset_parameters()
def reset_parameters(self):
nn.init.kaiming_normal_(self.fc1.weight)
nn.init.kaiming_normal_(self.fc2.weight)
nn.init.constant_(self.fc1.bias, 0.)
nn.init.constant_(self.fc2.bias, 0.)
def forward(
self,
x,
encoder_out,
encoder_padding_mask,
incremental_state,
prev_self_attn_state=None,
prev_attn_state=None,
self_attn_mask=None,
self_attn_padding_mask=None,
):
residual = x
x = self.self_attn_layer_norm(x)
if prev_self_attn_state is not None:
if incremental_state is None:
incremental_state = {}
prev_key, prev_value = prev_self_attn_state
saved_state = {'prev_key': prev_key, 'prev_value': prev_value}
self.self_attn._set_input_buffer(incremental_state, saved_state)
x, _ = self.self_attn(query=x,
key=x,
value=x,
key_padding_mask=self_attn_padding_mask,
incremental_state=incremental_state,
need_weights=False,
attn_mask=self_attn_mask)
x = F.dropout(x, p=self.dropout, training=self.training)
x = residual + x
attn = None
residual = x
x = self.encoder_attn_layer_norm(x)
if prev_attn_state is not None:
if incremental_state is None:
incremental_state = {}
prev_key, prev_value = prev_attn_state
saved_state = {'prev_key': prev_key, 'prev_value': prev_value}
self.encoder_attn._set_input_buffer(incremental_state, saved_state)
x, attn = self.encoder_attn(query=x,
key=encoder_out,
value=encoder_out,
key_padding_mask=encoder_padding_mask,
incremental_state=incremental_state,
static_kv=True,
need_weights=(not self.training))
x = F.dropout(x, p=self.dropout, training=self.training)
x = residual + x
residual = x
x = self.final_layer_norm(x)
x = F.relu(self.fc1(x))
x = F.dropout(x, p=self.relu_dropout, training=self.training)
x = self.fc2(x)
x = F.dropout(x, p=self.dropout, training=self.training)
x = residual + x
x = self.final_layer_norm(x)
return x, attn