Hướng Dẫn Triển Khai Word Embeddings Trong PyTorch

Hướng dẫn này trình bày cách triển khai và huấn luyện word embeddings trong PyTorch bằng nn.Embedding. Để bắt đầu, chúng ta cần định nghĩa mô hình cùng với các tham số cần thiết, như kích thước embedding và số lượng từ trong từ điển. Các bước bao gồm định nghĩa mô hình, huấn luyện với phương pháp Skip-gram, nơi mỗi từ trong ngữ cảnh được sử dụng để dự đoán các từ xung quanh, nhằm tối ưu hóa quan hệ ý nghĩa giữa chúng. Sau khi hoàn tất quá trình huấn luyện, chúng ta cần lưu trọng số embedding vào một tệp để có thể tải lại và sử dụng cho các nhiệm vụ khác trong tương lai, như phân loại văn bản hoặc tìm kiếm.

Trước khi bắt đầu, bạn có thể chạy code trên Colab và đọc phần hướng dẫn ở dưới.

1. Triển khai Word Embeddings trong PyTorch

PyTorch cung cấp nn.Embedding để tạo word embeddings.

import torch
import torch.nn as nn

# Định nghĩa kích thước từ vựng và số chiều của embedding
vocab_size = 10  # Kích thước từ vựng (số lượng từ)
embedding_dim = 5  # Số chiều của vector embedding

# Tạo một lớp embedding
embedding_layer = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)

# Ví dụ đầu vào (chỉ mục từ)
word_indices = torch.tensor([1, 3, 5, 7])  # Các từ được biểu diễn dưới dạng chỉ mục

# Lấy embeddings cho các từ đầu vào
word_embeddings = embedding_layer(word_indices)
print(word_embeddings)

Giải thích:

• nn.Embedding(num_embeddings, embedding_dim): Tạo ma trận embedding có kích thước [vocab_size, embedding_dim].
• word_indices: Các chỉ mục của từ trong tập từ vựng.
• Kết quả là một tensor có dạng [số lượng từ, embedding_dim].

2. Huấn luyện Word Embeddings trong PyTorch

Việc huấn luyện embeddings thường được thực hiện bằng mô hình Skip-gram hoặc CBOW (được sử dụng trong Word2Vec), hai phương pháp nổi bật trong xử lý ngôn ngữ tự nhiên.

Chuẩn bị dữ liệu huấn luyện

import torch
import torch.nn as nn
import torch.optim as optim

# Dữ liệu mẫu (cặp từ trung tâm và từ ngữ cảnh)
data = [(0, 1), (1, 2), (2, 3), (3, 4)]  # (từ trung tâm, từ ngữ cảnh)
vocab_size = 5  # Số lượng từ duy nhất
embedding_dim = 3  # Kích thước vector embedding

# Định nghĩa mô hình
class WordEmbeddingModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        super(WordEmbeddingModel, self).__init__()
        self.embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.linear = nn.Linear(embedding_dim, vocab_size)
    
    def forward(self, center_word):
        embed = self.embeddings(center_word)
        output = self.linear(embed)
        return output

# Khởi tạo mô hình
model = WordEmbeddingModel(vocab_size, embedding_dim)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Vòng lặp huấn luyện
for epoch in range(100):
    total_loss = 0
    for center, context in data:
        center_tensor = torch.tensor([center], dtype=torch.long)
        target_tensor = torch.tensor([context], dtype=torch.long)

        optimizer.zero_grad()
        output = model(center_tensor)
        loss = criterion(output, target_tensor)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
    
    if epoch % 20 == 0:
        print(f"Epoch {epoch}, Loss: {total_loss}")

Giải thích:

• Mô hình học embedding bằng cách dự đoán từ ngữ cảnh dựa trên từ trung tâm.
• nn.CrossEntropyLoss() được sử dụng để tối ưu hóa mô hình.
• SGD cập nhật embeddings dựa trên độ lỗi.

3. Lưu Mô Hình Embedding Đã Học

Sau khi huấn luyện, ta lưu lại lớp embedding để sử dụng sau này.

# Lưu chỉ phần trọng số của lớp embedding
torch.save(model.embeddings.state_dict(), "word_embeddings.pth")

# Lưu toàn bộ mô hình (tùy chọn)
torch.save(model.state_dict(), "embedding_model.pth")

Giải thích:

• torch.save(model.embeddings.state_dict(), "word_embeddings.pth"): Lưu chỉ trọng số của lớp embedding.
• torch.save(model.state_dict(), "embedding_model.pth"): Lưu toàn bộ mô hình (bao gồm cả lớp embedding và mạng tuyến tính).

4. Tải Lại Embeddings Đã Lưu

Để sử dụng lại embeddings đã huấn luyện:

# Khởi tạo lại mô hình
loaded_model = WordEmbeddingModel(vocab_size, embedding_dim)

# Tải trọng số của embedding đã lưu
loaded_model.embeddings.load_state_dict(torch.load("word_embeddings.pth"))

# Ví dụ sử dụng
word_idx = torch.tensor([2])  # Chỉ mục của từ cần lấy embedding
print(loaded_model.embeddings(word_idx))

Giải thích:

• WordEmbeddingModel(vocab_size, embedding_dim): Tạo lại cấu trúc mô hình.
• load_state_dict(torch.load("word_embeddings.pth")): Tải lại trọng số của embeddings đã lưu.

Kết quả:

tensor([[ 0.9381,  0.8774, -0.5396]], grad_fn=<EmbeddingBackward0>)

Tóm tắt

• Triển khai embeddings với nn.Embedding.
• Huấn luyện word embeddings với Skip-gram.
• Lưu embeddings đã học để sử dụng lại.
• Tải lại embeddings vào mô hình mới.