Внимание по головам в PyTorch

Класс MultiheadAttention считает взвешенную смесь векторов последовательности. При создании задают длину вектора на входе и сколько голов участвует в расчёте.

По умолчанию слой ждёт пакет последовательностей в порядке длина, затем пакет, затем размер вектора. Создадим слой, подадим тензор в такой форме и выведем форму ответа:

import torch
import torch.nn as nn

layer = nn.MultiheadAttention(embed_dim=8, num_heads=2)
seq = torch.randn(5, 3, 8)
output, weights = layer(seq, seq, seq)
print(output.shape) # выведет torch.Size([5, 3, 8])

Длину ряда, размер пакета и ширину вектора слой сохраняет на выходе. Если удобнее держать пакет первым измерением, при создании включают режим batch_first и подают тензор иначе:

import torch
import torch.nn as nn

layer = nn.MultiheadAttention(
	embed_dim=8,
	num_heads=2,
	batch_first=True,
)
seq = torch.randn(3, 5, 8)
output, weights = layer(seq, seq, seq)
print(output.shape) # выведет torch.Size([3, 5, 8])

Соберите слой с шириной вектора 12 и 3 головами. Подайте ряд длиной 4, пакет 2 и ту же ширину и выведите форму результата.

Возьмите ширину 16 и 4 головы, включите режим с пакетом первым. Передайте тензор размером 2 на 6 на 16 и выведите форму ответа.

Опишите слой на 10 чисел в векторе и 2 головы. Прогоните ряд из 7 шагов и пакет 5 и выведите форму выхода.