Gộp mọi kỹ năng: một lớp causal self-attention nhiều đầu — giống hệt khối trong GPT. So sánh với bản NumPy trong khóa Nhập môn AI: logic y hệt, chỉ khác có thêm chiều batch và chạy được trên GPU.
import math, torch, torch.nn as nn
class CausalSelfAttention(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
self.h = n_heads
self.qkv = nn.Linear(d_model, 3 * d_model) # tính Q, K, V một lần
self.proj = nn.Linear(d_model, d_model)
def forward(self, x): # x: (B, T, C)
B, T, C = x.shape
q, k, v = self.qkv(x).split(C, dim=-1)
q, k, v = (t.view(B, T, self.h, C // self.h).transpose(1, 2) for t in (q, k, v))
att = q @ k.transpose(-2, -1) / math.sqrt(C // self.h) # (B, h, T, T)
mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), 1)
att = att.masked_fill(mask, float('-inf')).softmax(dim=-1)
y = (att @ v).transpose(1, 2).contiguous().view(B, T, C) # ghép đầu
return self.proj(y)Trong thực tế dùng F.scaled_dot_product_attention(q, k, v, is_causal=True) — cùng kết quả nhưng nhanh hơn nhiều (FlashAttention). Tự viết một lần để hiểu, sau đó dùng hàm có sẵn.
?Trong code trên, att có shape gì?
Luyện tập ngay
Viết code và chấm tự động — có trợ lý AI gợi ý khi bí.