机器学习·注意力的时代

import torchimport torch.nn as nnimport torch.nn.functional as Fimport matplotlibmatplotlib.use("Agg")import matplotlib.pyplot as plt# ========== 1. 缩放点积注意力 ==========class ScaledDotProductAttention(nn.Module):    def __init__(self, d_k):        super().__init__()        self.d_k = d_k    def forward(self, Q, K, V, mask=None):        # Q,K,V: (batch, n_heads, seq_len, d_k)        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)        if mask is not None:            scores = scores.masked_fill(mask == 0, -1e9)        attn = F.softmax(scores, dim=-1)        out = torch.matmul(attn, V)        return out, attn# ========== 2. 多头注意力 ==========class MultiHeadAttention(nn.Module):    def __init__(self, d_model, n_heads):        super().__init__()        assert d_model % n_heads == 0        self.d_k = d_model // n_heads        self.n_heads = n_heads        self.W_q = nn.Linear(d_model, d_model)        self.W_k = nn.Linear(d_model, d_model)        self.W_v = nn.Linear(d_model, d_model)        self.W_o = nn.Linear(d_model, d_model)        self.attn = ScaledDotProductAttention(self.d_k)    def forward(self, x, mask=None):        B, L, D = x.size()        Q = self.W_q(x).view(B, L, self.n_heads, self.d_k).transpose(1, 2)        K = self.W_k(x).view(B, L, self.n_heads, self.d_k).transpose(1, 2)        V = self.W_v(x).view(B, L, self.n_heads, self.d_k).transpose(1, 2)        out, attn = self.attn(Q, K, V, mask)        out = out.transpose(1, 2).contiguous().view(B, L, D)        out = self.W_o(out)        return out, attn# ========== 3. 正弦位置编码 ==========class PositionalEncoding(nn.Module):    def __init__(self, d_model, max_len=512):        super().__init__()        pe = torch.zeros(max_len, d_model)        pos = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)        div = torch.exp(torch.arange(0, d_model, 2).float() *                        (-torch.log(torch.tensor(10000.0)) / d_model))        pe[:, 0::2] = torch.sin(pos * div)        pe[:, 1::2] = torch.cos(pos * div)        self.register_buffer("pe", pe.unsqueeze(0))    def forward(self, x):        return x + self.pe[:, :x.size(1)]# ========== 4. Transformer 编码器块 ==========class TransformerEncoderLayer(nn.Module):    def __init__(self, d_model, n_heads, d_ff):        super().__init__()        self.attn = MultiHeadAttention(d_model, n_heads)        self.ffn = nn.Sequential(            nn.Linear(d_model, d_ff),            nn.ReLU(),            nn.Linear(d_ff, d_model),        )        self.norm1 = nn.LayerNorm(d_model)        self.norm2 = nn.LayerNorm(d_model)    def forward(self, x):        a, attn = self.attn(x)        x = self.norm1(x + a)        x = self.norm2(x + self.ffn(x))        return x, attnclass TinyTransformer(nn.Module):    def __init__(self, vocab_size, d_model, n_heads, d_ff, n_layers, max_len=64):        super().__init__()        self.embed = nn.Embedding(vocab_size, d_model)        self.pos = PositionalEncoding(d_model, max_len)        self.layers = nn.ModuleList(            [TransformerEncoderLayer(d_model, n_heads, d_ff) for _ in range(n_layers)])        self.pool = nn.AdaptiveAvgPool1d(1)        self.head = nn.Linear(d_model, 2)    def forward(self, x):        x = self.embed(x)        x = self.pos(x)        attns = []        for layer in self.layers:            x, attn = layer(x)            attns.append(attn)        x = x.transpose(1, 2)        x = self.pool(x).squeeze(-1)        return self.head(x), attns# ========== 5. 合成数据：判断序列中 1 是否过半 ==========def make_data(n, L, seed=0):    g = torch.Generator().manual_seed(seed)    x = torch.randint(0, 2, (n, L), generator=g)    y = (x.sum(dim=1) > L // 2).long()    return x, y# ========== 6. 训练 ==========DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")L, VOCAB, D_MODEL, N_HEADS, D_FF, N_LAYERS, EPOCHS, LR, BATCH = 16, 2, 32, 4, 64, 2, 12, 1e-3, 64model = TinyTransformer(VOCAB, D_MODEL, N_HEADS, D_FF, N_LAYERS, max_len=L).to(DEVICE)crit = nn.CrossEntropyLoss()opt = torch.optim.Adam(model.parameters(), lr=LR)Xtr, ytr = make_data(2000, L, seed=1)Xte, yte = make_data(500, L, seed=2)for epoch in range(1, EPOCHS + 1):    model.train()    perm = torch.randperm(Xtr.size(0))    run_loss, correct, total = 0.0, 0, 0    for i in range(0, Xtr.size(0), BATCH):        idx = perm[i:i + BATCH]        xb, yb = Xtr[idx].to(DEVICE), ytr[idx].to(DEVICE)        opt.zero_grad()        logits, _ = model(xb)        loss = crit(logits, yb)        loss.backward()        opt.step()        run_loss += loss.item() * xb.size(0)        correct += logits.argmax(1).eq(yb).sum().item()        total += yb.size(0)    # 测试    model.eval()    with torch.no_grad():        logits, _ = model(Xte.to(DEVICE))        te_acc = 100. * logits.argmax(1).eq(yte.to(DEVICE)).sum().item() / yte.size(0)    print(f"Epoch {epoch:02d} | Loss {run_loss/total:.4f} | "          f"Train Acc {100.*correct/total:.2f}% | Test Acc {te_acc:.2f}%")# ========== 7. 可视化注意力热力图 ==========model.eval()sample = Xte[:1].to(DEVICE)with torch.no_grad():    _, attns = model(sample)last_layer = attns[-1][0, 0].detach().cpu().numpy()  # 第 0 个样本、第 0 个头fig, ax = plt.subplots(figsize=(6, 5))im = ax.imshow(last_layer, cmap="viridis")ax.set_title("Self-Attention Map (head 0, last layer)")ax.set_xlabel("Key position")ax.set_ylabel("Query position")fig.colorbar(im, ax=ax)plt.savefig("attention_map.png", dpi=120)plt.show()print("注意力热力图已保存至 attention_map.png")