Phân Mảnh & Nhúng Patch Không Gian Cho Vision Transformer (ViT)

Mục tiêu kiến thức

• Vision Transformer (ViT - Dosovitskiy et al., 2020) xử lý ảnh tương tự như câu văn trong NLP bằng cách chia ảnh thành chuỗi các mảnh (patches).
• Quy trình nhúng Patch:
• Ảnh đầu vào X ∈ R^{B × C × H × W}.
• Chia thành N = H · WP² mảnh không gian có kích thước P × P.
• Chiếu tuyến tính từng mảnh vào không gian ẩn D chiều (bằng phép Conv2D(C, D, kernel=P, stride=P)).
• Trải phẳng thành chuỗi vector X_patches ∈ R^{B × N × D}.
• Ghép thêm một token học được đại diện toàn cục [CLS] vào đầu chuỗi: X_tokens ∈ R^{B × (N + 1) × D}.
• Cộng vector mã hóa vị trí học được E_pos ∈ R^{1 × (N + 1) × D}.

Yêu cầu

Xây dựng lớp PatchEmbedding(torch.nn.Module):

class PatchEmbedding(nn.Module):
    def __init__(self, img_size: int = 224, patch_size: int = 16, in_channels: int = 3, embed_dim: int = 768):
        pass

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        pass

• Đầu ra forward(x): Tensor shape (B, N + 1, embed_dim) trong đó N = (img_size // patch_size)².

Ràng buộc

• Thời gian chạy tối đa: 3000ms.
• Giới hạn bộ nhớ: 512MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.