PyTorch torch.nn.MultiheadAttention-Funktion

PyTorch torch.nn 参考手册PyTorch torch.nn Referenzhandbuch

torch.nn.MultiheadAttentionEs ist ein Modul für den Multi-Head-Attention-Mechanismus in PyTorch.

Es ist eine Kernkomponente der Transformer-Architektur, die es dem Modell ermöglicht, gleichzeitig Informationen aus verschiedenen Repräsentationsunterräumen an unterschiedlichen Positionen zu berücksichtigen.

Funktionsdefinition

torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, kdim=None, vdim=None, batch_first=True)

Parameterbeschreibung:

  • embed_dim(int): Dimension des Eingabe-Embeddings.
  • num_heads(int): Anzahl der Attention-Heads.
  • dropout(float): Dropout-Wahrscheinlichkeit. Standardwert ist 0.
  • kdim(int): Dimension der Schlüssel-Vektoren. Standardwert ist None (identisch mit embed_dim).
  • vdim(int): Dimension der Wert-Vektoren. Standardwert ist None (identisch mit embed_dim).
  • batch_first(bool): Wenn True, ist die erste Dimension von Eingabe und Ausgabe der Batch. Standard ist True.

Anwendungsbeispiele

Beispiel 1: Grundlegende Verwendung

Beispiel

import torch
import torch.nn as nn

# Multi-Head-Attention erstellen: 512-dimensional, 8 Köpfe
mha = nn.MultiheadAttention(embed_dim=512, num_heads=8)

# Eingabe: batch=4, Sequenzlänge=100, Dimension=512
query = torch.randn(4, 100, 512)
key = torch.randn(4, 100, 512)
value = torch.randn(4, 100, 512)

# Vorwärtsdurchlauf
output, attn_weight = mha(query, key, value)

print("Query-Form:", query.shape)
print("Ausgabeform:", output.shape)
print("Form der Attention-Gewichte:", attn_weight.shape)

Beispiel 2: Self-Attention

Beispiel

import torch
import torch.nn as nn

mha = nn.MultiheadAttention(embed_dim=256, num_heads=4)

# Dieselbe Eingabe wie Q, K, V verwenden (Self-Attention)
x = torch.randn(2, 50, 256)

# Selbstaufmerksamkeit: q=k=v=x
output, weights = mha(x, x, x)

print("Eingabeform:", x.shape)
print("Ausgabeform:", output.shape)
print("Form der Attention-Gewichte:", weights.shape)

Beispiel 3: Aufmerksamkeit mit Maske

Beispiel

import torch
import torch.nn as nn

mha = nn.MultiheadAttention(embed_dim=128, num_heads=4)

# Eingabe
x = torch.randn(1, 20, 128)

# Erstelle eine obere Dreiecksmaske (für Decoder)
mask = torch.triu(torch.ones(20, 20), diagonal=1).bool()

output, _ = mha(x, x, x, attn_mask=mask)

print("Eingabeform:", x.shape)
print("Ausgabeform:", output.shape)
print("Form der Maske:", mask.shape)

Beispiel 4: Vollständige Transformer-Encoder-Schicht

Beispiel

import torch
import torch.nn as nn

class TransformerLayer(nn.Module):
    def __init__(self, d_model, nhead):
        super(TransformerLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.ReLU(),
            nn.Linear(d_model * 4, d_model)
        )

    def forward(self, x):
        # Selbstaufmerksamkeit + Restverbindung
        attn_out, _ = self.self_attn(x, x, x)
        x = self.norm1(x + attn_out)

        # FFN + Restverbindung
        ffn_out = self.ffn(x)
        x = self.norm2(x + ffn_out)

        return x

# Test
layer = TransformerLayer(d_model=512, nhead=8)
x = torch.randn(4, 100, 512)
output = layer(x)

print("Eingabeform:", x.shape)
print("Ausgabeform:", output.shape)

Beispiel 5: Anzeigen der Attention-Gewichte

Beispiel

import torch
import torch.nn as nn
import numpy as np

mha = nn.MultiheadAttention(embed_dim=64, num_heads=2, batch_first=True)

# Kurze Videosequenz
x = torch.randn(1, 5, 64)

_, attn = mha(x, x, x)
attn = attn.squeeze(0)  # Batch-Dimension entfernen

print("Attention-Gewichte des ersten Heads (erste 3 Positionen):")
print(attn[0, :3, :3].tolist())

print("nVisualisierung - Attention von Position 0 zu allen Positionen:")
print(np.array2string(attn[0, 0].numpy(), precision=2))

Häufig gestellte Fragen

F1: Wie wählt man num_heads?

embed_dim muss durch num_heads teilbar sein. Übliche Werte: 8, 12, 16.

F2: Warum werden die drei Matrizen Q, K und V benötigt?

Es ermöglicht dem Modell, unterschiedliche Projektionen zu erlernen und die Ausdruckskraft zu erhöhen.

F3: Was ist key_padding_mask?

Es wird verwendet, um Padding-Positionen zu maskieren und zu verhindern, dass Attention auf Padding berechnet wird.

Anwendungsfälle

  • Transformer: Encoder und Decoder
  • Self-Attention-Modell: BERT、GPT
  • Sequenzmodellierung: Alternative zu RNN

Tipp: Wenn batch_first=True, ist die Eingabeform (batch, seq, embed_dim).

PyTorch torch.nn 参考手册PyTorch torch.nn Referenzhandbuch

Andere Erweiterungen