PyTorch torch.nn.MultiheadAttention función

PyTorch torch.nn 参考手册Manual de referencia de PyTorch torch.nn

torch.nn.MultiheadAttentionEs el módulo de mecanismo de atención multi-cabeza en PyTorch.

Es un componente central de la arquitectura Transformer, que permite al modelo prestar atención simultáneamente a información de diferentes subespacios de representación en diferentes posiciones.

Definición de la función

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

Descripción de parámetros:

  • embed_dim(int): Dimensión de la incrustación de entrada.
  • num_heads(int): Número de cabezas de atención.
  • dropout(float): Probabilidad de dropout. Por defecto es 0.
  • kdim(int): Dimensión del vector de clave. Por defecto es None (igual que embed_dim).
  • vdim(int): Dimensión del vector de valor. Por defecto es None (igual que embed_dim).
  • batch_first(bool): Si es True, la primera dimensión de entrada y salida es batch. Por defecto es True.

Ejemplo de uso

Ejemplo 1: Uso básico

Ejemplo

import torch
import torch.nn as nn

# Crear atención multi-cabeza: 512 dimensiones, 8 cabezas
mha = nn.MultiheadAttention(embed_dim=512, num_heads=8)

# Entrada: batch=4, longitud de secuencia=100, dimensión=512
query = torch.randn(4, 100, 512)
key = torch.randn(4, 100, 512)
value = torch.randn(4, 100, 512)

# Propagación hacia adelante
output, attn_weight = mha(query, key, value)

print("Forma de Query:", query.shape)
print("Forma de Output:", output.shape)
print("Forma de los pesos de atención:", attn_weight.shape)

Ejemplo 2: Autoatención

Ejemplo

import torch
import torch.nn as nn

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

# La misma entrada como Q, K, V (autoatención)
x = torch.randn(2, 50, 256)

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

print("Forma de entrada:", x.shape)
print("Forma de salida:", output.shape)
print("Forma de los pesos de atención:", weights.shape)

Ejemplo 3: Atención con máscara

Ejemplo

import torch
import torch.nn as nn

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

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

# Crear máscara triangular superior (para el decodificador)
mask = torch.triu(torch.ones(20, 20), diagonal=1).bool()

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

print("Forma de entrada:", x.shape)
print("Forma de Output:", output.shape)
print("Forma de Mask:", mask.shape)

Ejemplo 4: Capa completa de codificador Transformer

Ejemplo

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):
        # Autoatención + residual
        attn_out, _ = self.self_attn(x, x, x)
        x = self.norm1(x + attn_out)

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

        return x

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

print("Forma de entrada:", x.shape)
print("Forma de salida:", output.shape)

Ejemplo 5: Ver los pesos de atención

Ejemplo

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

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

# Secuencia de video corta
x = torch.randn(1, 5, 64)

_, attn = mha(x, x, x)
attn = attn.squeeze(0)  # Quitar batch

print("Pesos de atención de la primera cabeza (primeras 3 posiciones):")
print(attn[0, :3, :3].tolist())

print("nVisualización - atención de la posición 0 hacia todas las posiciones:")
print(np.array2string(attn[0, 0].numpy(), precision=2))

Preguntas frecuentes

Q1: ¿Cómo elegir num_heads?

embed_dim debe ser divisible por num_heads. Valores comunes: 8, 12, 16.

Q2: ¿Por qué se necesitan las tres matrices Q, K, V?

Permite que el modelo aprenda diferentes proyecciones, mejorando la capacidad de expresión.

Q3: ¿Qué es key_padding_mask?

Se utiliza para enmascarar las posiciones de padding, evitando que la atención se calcule sobre el padding.

Escenarios de uso

  • Transformer: Codificador y decodificador
  • Modelo de autoatención: BERT、GPT
  • Modelado de secuencias: Reemplazo de RNN

Consejo: cuando batch_first=True, la forma de entrada es (batch, seq, embed_dim).

PyTorch torch.nn 参考手册Manual de referencia de PyTorch torch.nn

Otras extensiones