Recurrent Neural Network

import torch
import torch.nn as nn

class SimpleRNNCell(nn.Module):
    """
    Một ô RNN thuần tuân đúng công thức Elman:
        h_t = tanh(W_xh * x_t + W_hh * h_{t-1} + b_h)
    Dùng nn.Parameter để PyTorch tự tính gradient qua BPTT.
    """
    def __init__(self, input_dim: int, hidden_dim: int):
        super().__init__()
        # Khởi tạo theo Xavier để tránh vanishing/exploding gradient
        # ngay từ bước đầu. nn.Linear chỉ là tiện nghi cho W·x + b.
        self.W_xh = nn.Linear(input_dim, hidden_dim, bias=True)
        self.W_hh = nn.Linear(hidden_dim, hidden_dim, bias=False)

    def forward(self, x_t: torch.Tensor, h_prev: torch.Tensor) -> torch.Tensor:
        # x_t: (batch, input_dim); h_prev: (batch, hidden_dim)
        return torch.tanh(self.W_xh(x_t) + self.W_hh(h_prev))


class SimpleRNN(nn.Module):
    """
    Quấn nhiều bước thời gian lại: nhận chuỗi (batch, T, input_dim),
    trả về toàn bộ hidden states (batch, T, hidden_dim) và h_T.
    """
    def __init__(self, input_dim: int, hidden_dim: int):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.cell = SimpleRNNCell(input_dim, hidden_dim)

    def forward(self, x: torch.Tensor, h0: torch.Tensor | None = None):
        batch, T, _ = x.shape
        if h0 is None:
            h0 = x.new_zeros(batch, self.hidden_dim)

        h = h0
        hs = []
        # Vòng lặp tuần tự theo thời gian — đây chính là lý do RNN
        # KHÔNG song song hóa được giữa các bước (t và t+1).
        for t in range(T):
            h = self.cell(x[:, t, :], h)
            hs.append(h)

        # hs: danh sách T tensor kích thước (batch, hidden_dim)
        return torch.stack(hs, dim=1), h


# Thực tế: dùng nn.RNN của PyTorch sẽ nhanh hơn (CuDNN tối ưu).
# Ví dụ sử dụng:
if __name__ == "__main__":
    rnn = SimpleRNN(input_dim=50, hidden_dim=128)
    # Batch 4, chuỗi 20 bước, mỗi bước embedding 50 chiều
    dummy = torch.randn(4, 20, 50)
    out, h_final = rnn(dummy)
    print(out.shape, h_final.shape)  # (4, 20, 128) (4, 128)

import torch
import torch.nn as nn

# Giả sử bạn có sẵn mô hình và data loader
model = nn.RNN(input_size=50, hidden_size=128, batch_first=True)
head = nn.Linear(128, vocab_size)  # vocab_size tùy bài toán
opt = torch.optim.Adam(list(model.parameters()) + list(head.parameters()), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(num_epochs):
    for batch_x, batch_y in train_loader:
        # batch_x: (B, T, 50); batch_y: (B, T) — chỉ số token đích
        out, _ = model(batch_x)          # (B, T, 128)
        logits = head(out)               # (B, T, V)
        loss = loss_fn(logits.reshape(-1, vocab_size), batch_y.reshape(-1))

        opt.zero_grad()
        loss.backward()

        # QUAN TRỌNG: cắt gradient để tránh exploding gradient
        # Với RNN thuần, norm gradient rất dễ nổ > 10^3 khi chuỗi dài.
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)

        opt.step()