def __init__(
self,
embed_dim,
dropout,
att_type,
att_heads,
att_mid_dim,
att_mid_drop,
bifeat_emb_act,
bifeat_emb_drop,
ff_dropout,
last_layer = False
):
super(DecoderLayer, self).__init__()
self.last_layer = last_layer
self.word_attn = LowRank(
embed_dim = embed_dim,
att_type = att_type,
att_heads = att_heads,
att_mid_dim = att_mid_dim,
att_mid_drop = att_mid_drop)
self.word_dropout = nn.Dropout(dropout)
self.cross_att = LowRank(
embed_dim = embed_dim,
att_type = att_type,
att_heads = att_heads,
att_mid_dim = att_mid_dim,
att_mid_drop = att_mid_drop)
self.cross_dropout = nn.Dropout(dropout)
self.layer_norm_cross = torch.nn.LayerNorm(embed_dim)
if self.last_layer == False:
self.bifeat_emb = nn.Sequential(
nn.Linear(2 * embed_dim, embed_dim),
utils.activation(bifeat_emb_act),
nn.Dropout(bifeat_emb_drop)
)
self.layer_norm_x = torch.nn.LayerNorm(embed_dim)
self.ff_layer = blocks.create(
'FeedForward',
embed_dim = embed_dim,
ffn_embed_dim = embed_dim * 4,
relu_dropout = ff_dropout,
dropout = ff_dropout)
self.layer_norm_gx = torch.nn.LayerNorm(embed_dim)
def __init__(
self,
embed_dim,
dropout,
att_type,
att_heads,
att_mid_dim,
att_mid_drop,
bifeat_emb_act,
bifeat_emb_drop,
ff_dropout
):
super(EncoderLayer, self).__init__()
self.encoder_attn = LowRank(
embed_dim = embed_dim,
att_type = att_type,
att_heads = att_heads,
att_mid_dim = att_mid_dim,
att_mid_drop = att_mid_drop)
self.dropout = nn.Dropout(dropout)
self.bifeat_emb = nn.Sequential(
nn.Linear(2 * embed_dim, embed_dim),
utils.activation(bifeat_emb_act),
nn.Dropout(bifeat_emb_drop)
)
self.layer_norm = torch.nn.LayerNorm(embed_dim)
self.ff_layer = blocks.create(
'FeedForward',
embed_dim = embed_dim,
ffn_embed_dim = embed_dim * 4,
relu_dropout = ff_dropout,
dropout = ff_dropout)
def __init__(self, embed_dim, att_type, att_heads, att_mid_dim,
att_mid_drop, dropout):
super(LowRankBilinearLayer, self).__init__()
self.encoder_attn = LowRank(embed_dim=embed_dim,
att_type=att_type,
att_heads=att_heads,
att_mid_dim=att_mid_dim,
att_mid_drop=att_mid_drop)
self.dropout = nn.Dropout(dropout) if dropout > 0 else None