PoliBa-DeepLearning/NanoSocrates
2
0
Fork 0
Code Issues 3 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

update to batch attention mask

Browse Source
This commit is contained in:
GassiGiuseppe
2025-10-06 13:03:03 +02:00
parent 87409fecd5
commit 948c3fd7ac
3 changed files with 19 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
Project_Model/Libs/Transformer/Classes/Decoder.py
View File

@@ -2,8 +2,9 @@ import torch
import torch.nn as nn
from .FeedForwardNetwork import FeedForwardNetwork
from .TorchMultiHeadAttention import TorchMultiHeadAttention as MultiHeadAttention
from ..Utils.attention_mask import get_attention_mask
from ..Utils.attention_mask import get_causal_attention_mask
# B, L(T), E_D
class Decoder(nn.Module):
@@ -17,7 +18,7 @@ class Decoder(nn.Module):
super().__init__()
self.__masked_attention = MultiHeadAttention(
embedding_dimension, number_of_attention_heads, dropout=0.1, attention_mask=get_attention_mask(embedding_dimension)
embedding_dimension, number_of_attention_heads, dropout=0.1
)
self.__layer_norm_1 = nn.LayerNorm(embedding_dimension)
@@ -38,9 +39,12 @@ class Decoder(nn.Module):
def forward(self, x, k_x, v_x, padding_mask = None) -> torch.Tensor: # k_x = v_x . While x_q = x
# build of attention mask
attention_mask = get_causal_attention_mask(x.size(1))
# 1) Masked Attention
MASKED_ATTENTION = self.__masked_attention(
x, x, x, key_padding_mask=padding_mask
x, x, x, key_padding_mask=padding_mask, attn_mask=attention_mask
)
# 2) Dropout
@@ -57,7 +61,7 @@ class Decoder(nn.Module):
x = self.__layer_norm_1(x)
# 5) Encoder–decoder (cross) attention
CROSS_ATTENTION = self.__cross_attention(x, k_x, v_x key_padding_mask=padding_mask)
CROSS_ATTENTION = self.__cross_attention(x, k_x, v_x, key_padding_mask=padding_mask)
# 6) Dropout
DROPPED_CROSS_ATTENTION = self.__dropout(CROSS_ATTENTION)