Fonction PyTorch torch.nn.MultiheadAttention

PyTorch torch.nn 参考手册Manuel de référence torch.nn de PyTorch

torch.nn.MultiheadAttentionIl s'agit d'un module de mécanisme d'attention multi-têtes dans PyTorch.

Il s'agit d'un composant central de l'architecture Transformer, permettant au modèle de prêter simultanément attention à des informations provenant de différents sous-espaces de représentation à différentes positions.

Définition de fonction

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

Description des paramètres :

  • embed_dim(int) : Dimension de l'embedding d'entrée.
  • num_heads(int) : Nombre de têtes d'attention.
  • dropout(float) : Probabilité de dropout. La valeur par défaut est 0.
  • kdim(int) : Dimension des vecteurs de clé. La valeur par défaut est None (identique à embed_dim).
  • vdim(int) : Dimension des vecteurs de valeur. La valeur par défaut est None (identique à embed_dim).
  • batch_first(bool) : Si True, la première dimension de l'entrée et de la sortie est le batch. La valeur par défaut est True.

Exemples d'utilisation

Exemple 1 : Utilisation de base

Exemple

import torch
import torch.nn as nn

# Créer une attention multi-têtes : 512 dimensions, 8 têtes
mha = nn.MultiheadAttention(embed_dim=512, num_heads=8)

# Entrée : batch=4, longueur de séquence=100, dimension=512
query = torch.randn(4, 100, 512)
key = torch.randn(4, 100, 512)
value = torch.randn(4, 100, 512)

# Passe avant
output, attn_weight = mha(query, key, value)

print("Forme de query :", query.shape)
print("Forme de sortie :", output.shape)
print("Forme des poids d'attention :", attn_weight.shape)

Exemple 2 : Auto-attention

Exemple

import torch
import torch.nn as nn

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

# Utiliser la même entrée que Q, K, V (auto-attention)
x = torch.randn(2, 50, 256)

# auto-attention : q=k=v=x
output, weights = mha(x, x, x)

print("Forme d'entrée :", x.shape)
print("Forme de sortie :", output.shape)
print("Forme des poids d'attention :", weights.shape)

Exemple 3 : Attention avec masque

Exemple

import torch
import torch.nn as nn

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

# Entrée
x = torch.randn(1, 20, 128)

# Créer un masque triangulaire supérieur (pour le décodeur)
mask = torch.triu(torch.ones(20, 20), diagonal=1).bool()

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

print("Forme d'entrée :", x.shape)
print("Forme de sortie :", output.shape)
print("Forme du masque :", mask.shape)

Exemple 4 : Couche d'encodeur Transformer complète

Exemple

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):
        # Auto-attention + résiduel
        attn_out, _ = self.self_attn(x, x, x)
        x = self.norm1(x + attn_out)

        # FFN + résidu
        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("Forme d'entrée :", x.shape)
print("Forme de sortie :", output.shape)

Exemple 5 : Visualisation des poids d’attention

Exemple

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

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

# Séquence vidéo courte
x = torch.randn(1, 5, 64)

_, attn = mha(x, x, x)
attn = attn.squeeze(0)  # Supprimer la dimension de batch

print("Poids d’attention de la première tête (3 premières positions) :")
print(attn[0, :3, :3].tolist())

print("nVisualisation - attention de la position 0 vers toutes les positions :")
print(np.array2string(attn[0, 0].numpy(), precision=2))

Questions fréquentes

Q1 : Comment choisir num_heads ?

embed_dim doit être divisible par num_heads. Valeurs courantes : 8, 12, 16.

Q2 : Pourquoi les trois matrices Q, K et V sont-elles nécessaires ?

Cela permet au modèle d'apprendre différentes projections et d'améliorer la puissance expressive.

Q3 : Qu'est-ce que key_padding_mask ?

Il est utilisé pour masquer les positions de padding, empêchant le calcul de l'attention sur le padding.

Cas d'utilisation

  • Transformer: Encodeur et Décodeur
  • Modèle d'auto-attention: BERT、GPT
  • Modélisation de séquences: Alternative aux RNN

Astuce : Lorsque batch_first=True, la forme d'entrée est (batch, seq, embed_dim).

PyTorch torch.nn 参考手册Manuel de référence torch.nn de PyTorch

Autres extensions