panda/bettafish-company
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Define the linear transformation layer

Browse Source
This commit is contained in:
戒酒的李白
2024-10-06 11:34:31 +08:00
parent ee739c3c81
commit f5e307d3f8
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model_pro/MHA.py
+6 -1
View File
@@ -9,10 +9,15 @@ class MultiHeadAttentionLayer(nn.Module):
self.head_dim = embed_size // num_heads self.head_dim = embed_size // num_heads
assert (self.head_dim * num_heads == embed_size), "Embedding size needs to be divisible by num_heads" assert (self.head_dim * num_heads == embed_size), "Embedding size needs to be divisible by num_heads"
# Define linear layers for Q, K, V
self.q_linear = nn.Linear(embed_size, embed_size)
self.k_linear = nn.Linear(embed_size, embed_size)
self.v_linear = nn.Linear(embed_size, embed_size)
if __name__ == "__main__": if __name__ == "__main__":
embed_size = 512 embed_size = 512
num_heads = 8 num_heads = 8
mha_layer = MultiHeadAttentionLayer(embed_size, num_heads) mha_layer = MultiHeadAttentionLayer(embed_size, num_heads)
print("Model initialized successfully.") print("Linear layers for Q, K, V initialized.")