added padding_mask entry to decoder and encoder
This commit is contained in:
parent
9c83d9fa71
commit
0f243eaac2
@ -36,11 +36,11 @@ class Decoder(nn.Module):
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
def forward(self, x, k_x, v_x, attention_mask) -> torch.Tensor: # k_x = v_x . While x_q = x
|
def forward(self, x, k_x, v_x, padding_mask = None) -> torch.Tensor: # k_x = v_x . While x_q = x
|
||||||
|
|
||||||
# 1) Masked Attention
|
# 1) Masked Attention
|
||||||
MASKED_ATTENTION = self.__masked_attention(
|
MASKED_ATTENTION = self.__masked_attention(
|
||||||
x, x, x, attention_mask=attention_mask
|
x, x, x, key_padding_mask=padding_mask
|
||||||
)
|
)
|
||||||
|
|
||||||
# 2) Dropout
|
# 2) Dropout
|
||||||
@ -57,7 +57,7 @@ class Decoder(nn.Module):
|
|||||||
x = self.__layer_norm_1(x)
|
x = self.__layer_norm_1(x)
|
||||||
|
|
||||||
# 5) Encoder–decoder (cross) attention
|
# 5) Encoder–decoder (cross) attention
|
||||||
CROSS_ATTENTION = self.__cross_attention(x, k_x, v_x)
|
CROSS_ATTENTION = self.__cross_attention(x, k_x, v_x key_padding_mask=padding_mask)
|
||||||
|
|
||||||
# 6) Dropout
|
# 6) Dropout
|
||||||
DROPPED_CROSS_ATTENTION = self.__dropout(CROSS_ATTENTION)
|
DROPPED_CROSS_ATTENTION = self.__dropout(CROSS_ATTENTION)
|
||||||
|
|||||||
@ -31,12 +31,12 @@ class Encoder(
|
|||||||
self.__dropout = nn.Dropout(0.1) # ...
|
self.__dropout = nn.Dropout(0.1) # ...
|
||||||
pass
|
pass
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x, padding_mask = None):
|
||||||
# -> ATTENTION -> dropout -> add and normalize -> FF -> dropout -> add and normalize ->
|
# -> ATTENTION -> dropout -> add and normalize -> FF -> dropout -> add and normalize ->
|
||||||
# Attention with Residual Connection [ input + self-attention]
|
# Attention with Residual Connection [ input + self-attention]
|
||||||
|
|
||||||
# 1) Multi Head Attention
|
# 1) Multi Head Attention
|
||||||
ATTENTION = self.__attention(x, x, x)
|
ATTENTION = self.__attention(x, x, x,key_padding_mask= padding_mask)
|
||||||
|
|
||||||
# 2) Dropout
|
# 2) Dropout
|
||||||
DROPPED_ATTENTION = self.__dropout(ATTENTION)
|
DROPPED_ATTENTION = self.__dropout(ATTENTION)
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user