Define the linear transformation layer

2024-10-06 11:34:31 +08:00
parent ee739c3c81
commit f5e307d3f8
1 changed files with 6 additions and 1 deletions
@@ -9,10 +9,15 @@ class MultiHeadAttentionLayer(nn.Module):
        self.head_dim = embed_size // 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__":
    embed_size = 512
    num_heads = 8
    mha_layer = MultiHeadAttentionLayer(embed_size, num_heads)
-    print("Model initialized successfully.")
+    print("Linear layers for Q, K, V initialized.")