Функция PyTorch torch.nn.MultiheadAttention

PyTorch torch.nn 参考手册Справочное руководство PyTorch torch.nn

torch.nn.MultiheadAttentionЭто модуль механизма многоголового внимания в PyTorch.

Это ключевой компонент архитектуры Transformer, позволяющий модели одновременно обрабатывать информацию из разных подпространств представлений в разных позициях.

Определение функции

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

Описание параметров:

  • embed_dim(int): Размерность входного эмбеддинга.
  • num_heads(int): Количество голов внимания.
  • dropout(float): Вероятность dropout. По умолчанию 0.
  • kdim(int): Размерность векторов ключей. По умолчанию None (совпадает с embed_dim).
  • vdim(int): Размерность векторов значений. По умолчанию None (совпадает с embed_dim).
  • batch_first(bool): Если True, первое измерение входа и выхода — batch. По умолчанию True.

Примеры использования

Пример 1: Базовое использование

Пример

import torch
import torch.nn as nn

# Создать многоголовое внимание: 512-мерное, 8 голов
mha = nn.MultiheadAttention(embed_dim=512, num_heads=8)

# Вход: batch=4, длина последовательности=100, размерность=512
query = torch.randn(4, 100, 512)
key = torch.randn(4, 100, 512)
value = torch.randn(4, 100, 512)

# Прямой проход
output, attn_weight = mha(query, key, value)

print("Query shape:", query.shape)
print("Форма выхода:", output.shape)
print("Attention weight shape:", attn_weight.shape)

Пример 2: Самовнимание

Пример

import torch
import torch.nn as nn

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

# Использовать тот же вход, что и Q, K, V (самовнимание)
x = torch.randn(2, 50, 256)

# самовнимание: q=k=v=x
output, weights = mha(x, x, x)

print("Форма входа:", x.shape)
print("Форма выхода:", output.shape)
print("Форма весов внимания:", weights.shape)

Пример 3: Внимание с маской

Пример

import torch
import torch.nn as nn

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

# Ввод
x = torch.randn(1, 20, 128)

# Создать верхнетреугольную маску (для декодера)
mask = torch.triu(torch.ones(20, 20), diagonal=1).bool()

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

print("Форма входа:", x.shape)
print("Форма выхода:", output.shape)
print("Форма маски:", mask.shape)

Пример 4: Полный слой энкодера Transformer

Пример

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):
        # Самовнимание + остаточная связь
        attn_out, _ = self.self_attn(x, x, x)
        x = self.norm1(x + attn_out)

        # FFN + остаточная связь
        ffn_out = self.ffn(x)
        x = self.norm2(x + ffn_out)

        return x

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

print("Форма входа:", x.shape)
print("Форма выхода:", output.shape)

Пример 5: Просмотр весов внимания

Пример

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

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

# Короткая видеопоследовательность
x = torch.randn(1, 5, 64)

_, attn = mha(x, x, x)
attn = attn.squeeze(0)  # Удалить размерность батча

print("Веса внимания первой головы (первые 3 позиции):")
print(attn[0, :3, :3].tolist())

print("nВизуализация - внимание позиции 0 ко всем позициям:")
print(np.array2string(attn[0, 0].numpy(), precision=2))

Часто задаваемые вопросы

Вопрос 1: Как выбрать num_heads?

embed_dim должен делиться на num_heads. Типичные значения: 8, 12, 16.

Вопрос 2: Зачем нужны три матрицы Q, K и V?

Это позволяет модели изучать разные проекции и повышать выразительную способность.

Вопрос 3: Что такое key_padding_mask?

Он используется для маскирования позиций паддинга, предотвращая вычисление внимания на паддинге.

Сценарии использования

  • Transformer: Энкодер и декодер
  • Модель самовнимания: BERT、GPT
  • Моделирование последовательностей: Альтернатива RNN

Совет: Когда batch_first=True, форма входных данных — (batch, seq, embed_dim).

PyTorch torch.nn 参考手册Справочное руководство PyTorch torch.nn

Другие расширения