File size: 2,413 Bytes
607293d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import torch
import torch.nn as nn
class SimpleRNN(nn.Module):
"""
Simple RNN model for text classification.
Architecture: Embedding -> RNN -> Hidden Layers -> Output
"""
def __init__(
self,
vocab_size: int,
embedding_dim: int = 128,
hidden_dim: int = 256,
num_layers: int = 2,
num_classes: int = 2,
dropout: float = 0.3,
num_hidden_nodes: int = 128
):
"""
Args:
vocab_size: Size of vocabulary
embedding_dim: Dimension of word embeddings
hidden_dim: Hidden dimension of RNN
num_layers: Number of RNN layers
num_classes: Number of output classes
dropout: Dropout rate
num_hidden_nodes: Number of nodes in hidden layer
"""
super(SimpleRNN, self).__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
self.rnn = nn.RNN(
input_size=embedding_dim,
hidden_size=hidden_dim,
num_layers=num_layers,
batch_first=True,
dropout=dropout if num_layers > 1 else 0
)
# Hidden layer (flexible nodes)
self.hidden_layer = nn.Sequential(
nn.Linear(hidden_dim, num_hidden_nodes),
nn.ReLU(),
nn.Dropout(dropout)
)
# Output layer
self.output_layer = nn.Linear(num_hidden_nodes, num_classes)
def forward(self, x):
"""
Forward pass.
Args:
x: Input tensor of shape (batch_size, seq_length)
Returns:
Output tensor of shape (batch_size, num_classes)
"""
# Embedding
embedded = self.embedding(x) # (batch_size, seq_length, embedding_dim)
# RNN
rnn_out, hidden = self.rnn(embedded) # rnn_out: (batch_size, seq_length, hidden_dim)
# Use the last output of the sequence
# For simple RNN, this is the final hidden state after processing all tokens
last_output = rnn_out[:, -1, :] # (batch_size, hidden_dim)
# Hidden layer
hidden_out = self.hidden_layer(last_output) # (batch_size, num_hidden_nodes)
# Output layer
output = self.output_layer(hidden_out) # (batch_size, num_classes)
return output
|