Ich habe mir ein PyTorch Modell für die Textklassifizierung erstellt.
Genauer möchte ich zwischen Programmiersprachen unterscheiden können.
Leider habe ich noch ein Problem bei der Ausgabegrösse. Ich möchte dass ich am Schluss nur die 34 Klassen pro Sample bekomme. Diese sind One-Hot kodiert also zum Beispiel: Klasse 3 = [0, 0, 1].
Die Eingabe ist immer [BATCH_SIZE, 1000]
Die Ausgabe ist [BATCH_SIZE, 1000, 34]
Ich möchte [BATCH_SIZE, 34]
Code:
Code: Alles auswählen
import torch
import torch.nn as nn
EMBEDDING_DIM = 512
HIDDEN_DIM = 256
NUM_LAYERS = 3
BIDIRECTIONAL = True
DROPOUT = 0.2
NUM_HEADS = 8
OUTPUT_DIM =34
BATCH_SIZE = 16
DEVICE = "cuda"
class LSTMWithMultiHeadAttention(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, num_layers, dropout, num_heads, bidirectional):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim)
self.lstm = nn.LSTM(embedding_dim, hidden_dim, num_layers=num_layers, bidirectional=bidirectional, batch_first=True)
self.multihead_attention = nn.MultiheadAttention(hidden_dim * 2, num_heads, dropout=dropout)
self.final = nn.Linear(hidden_dim * 2, output_dim)
def forward(self, x):
embedded = self.embedding(x)
lstm_output, (hidden, cell) = self.lstm(embedded)
attention_output, _ = self.multihead_attention(lstm_output, lstm_output, lstm_output)
return self.final(attention_output)
model = LSTMWithMultiHeadAttention(256 + 1, EMBEDDING_DIM, HIDDEN_DIM, OUTPUT_DIM, NUM_LAYERS, DROPOUT, NUM_HEADS, BIDIRECTIONAL).cuda()
print(model(torch.tensor([3]*1000).cuda().unsqueeze(0)).size())Falls ihr noch Verbesserungsvorschläge fürs Modell oder die Parameter habe könnt ihr die gerne einbringen.
