Класс TransformerDecoder

Класс TransformerDecoder реализует стек слоёв декодера, который используется в архитектуре Transformer. Он принимает на вход скрытые представления от энкодера и генерирует выходную последовательность. Конструктор класса первым параметром принимает экземпляр слоя TransformerDecoderLayer, вторым параметром - количество таких слоёв num_layers. Также можно передать параметр norm для нормализации выходных данных.

Синтаксис

torch.nn.TransformerDecoder(
	decoder_layer: TransformerDecoderLayer,
	num_layers: int,
	norm: Optional[Module] = None
)

Пример

Создадим простой декодер с одним слоем и применим его к входным данным:

import torch
import torch.nn as nn

# Создаём один слой декодера с размерностью модели 512 и 8 головами внимания
decoder_layer = nn.TransformerDecoderLayer(
	d_model=512,
	nhead=8,
	batch_first=True
)
# Создаём декодер из двух таких слоёв
decoder = nn.TransformerDecoder(
	decoder_layer,
	num_layers=2
)

# Генерируем случайные входные данные
tgt = torch.rand(10, 32, 512)  # (seq_len, batch, d_model)
memory = torch.rand(20, 32, 512)  # (seq_len, batch, d_model)

# Прямой проход через декодер
output = decoder(tgt, memory)
print(output.shape)

Результат выполнения кода:

torch.Size([10, 32, 512])

Пример

При работе с декодером важно использовать маски, чтобы скрыть будущие токены в целевой последовательности. Рассмотрим создание маски и передачу её в метод forward:

import torch
import torch.nn as nn

torch.manual_seed(0)

decoder_layer = nn.TransformerDecoderLayer(
	d_model=512,
	nhead=8,
	batch_first=True
)
decoder = nn.TransformerDecoder(
	decoder_layer,
	num_layers=1
)

tgt = torch.rand(2, 5, 512)  # (batch, seq_len, d_model)
memory = torch.rand(2, 7, 512)  # (batch, seq_len, d_model)

# Создаём маску для целевой последовательности (причинная маска)
tgt_mask = torch.triu(
	torch.ones(5, 5) * float('-inf'),
	diagonal=1
)

output = decoder(
	tgt,
	memory,
	tgt_mask=tgt_mask
)
print(output.shape)

Результат выполнения кода:

torch.Size([2, 5, 512])

Пример

Декодер поддерживает передачу ключевых и сто́имостных проекций отдельно от памяти энкодера. Это полезно в задачах, где требуется дополнительная информация, например, в моделях с кэшированием состояний:

import torch
import torch.nn as nn

torch.manual_seed(0)

decoder_layer = nn.TransformerDecoderLayer(
	d_model=512,
	nhead=8,
	batch_first=True
)
decoder = nn.TransformerDecoder(
	decoder_layer,
	num_layers=2
)

tgt = torch.rand(3, 4, 512)  # (batch, seq_len, d_model)
memory = torch.rand(3, 6, 512)  # (batch, seq_len, d_model)

# Отдельные проекции для ключей и значений
memory_key = torch.rand(3, 6, 512)
memory_value = torch.rand(3, 6, 512)

output = decoder(
	tgt,
	memory,
	memory_key=memory_key,
	memory_value=memory_value
)
print(output.shape)

Результат выполнения кода:

torch.Size([3, 4, 512])

Смотрите также

класс Transformer,
который объединяет энкодер и декодер в одной модели
класс TransformerEncoder,
который реализует стек слоёв энкодера
класс TransformerEncoderLayer,
который представляет один слой энкодера
класс TransformerDecoderLayer,
который представляет один слой декодера