Класс Transformer

Класс Transformer из модуля torch.nn реализует архитектуру трансформера, предложенную в статье "Attention Is All You Need". Он состоит из стека слоёв кодировщика и стека слоёв декодировщика. Основными параметрами являются количество голов внимания nhead, размерность модели d_model, количество слоёв кодировщика num_encoder_layers и декодировщика num_decoder_layers. Класс принимает входные последовательности для кодировщика и декодировщика, а также опциональные маски для управления вниманием.

Синтаксис

torch.nn.Transformer(
	d_model=512,
	nhead=8,
	num_encoder_layers=6,
	num_decoder_layers=6,
	dim_feedforward=2048,
	dropout=0.1,
	activation='relu',
	custom_encoder=None,
	custom_decoder=None,
	layer_norm_eps=1e-5,
	batch_first=False,
	norm_first=False,
	bias=True,
	device=None,
	dtype=None
)

Основные параметры

Класс принимает следующие ключевые параметры:

d_model - размерность модели (по умолчанию 512);
nhead - количество голов внимания (по умолчанию 8);
num_encoder_layers - количество слоёв в кодировщике (по умолчанию 6);
num_decoder_layers - количество слоёв в декодировщике (по умолчанию 6);
dim_feedforward - размерность скрытого слоя в полносвязной сети (по умолчанию 2048);
dropout - вероятность дропаута (по умолчанию 0.1);
activation - функция активации ('relu' или 'gelu', по умолчанию 'relu');
batch_first - если True, входной тензор имеет форму (batch, seq, feature) (по умолчанию False).

Пример

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

import torch

transformer = torch.nn.Transformer(
	d_model=512,
	nhead=8,
	num_encoder_layers=6,
	num_decoder_layers=6
)

src = torch.rand((10, 32, 512))
tgt = torch.rand((20, 32, 512))

out = transformer(src, tgt)
print(out.shape)

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

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

Пример

Используем трансформер с масками для авторегрессивной генерации:

import torch

torch.manual_seed(0)

transformer = torch.nn.Transformer(
	d_model=128,
	nhead=4,
	num_encoder_layers=2,
	num_decoder_layers=2,
	batch_first=True
)

src = torch.rand((2, 10, 128))
tgt = torch.rand((2, 8, 128))

src_mask = torch.zeros((10, 10))
tgt_mask = torch.nn.Transformer.generate_square_subsequent_mask(8)
memory_mask = torch.zeros((8, 10))

out = transformer(
	src,
	tgt,
	src_mask=src_mask,
	tgt_mask=tgt_mask,
	memory_mask=memory_mask
)

print(out.shape)

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

torch.Size([2, 8, 128])

Пример

Создадим трансформер с использованием слоя нормализации перед вниманием:

import torch

transformer = torch.nn.Transformer(
	d_model=256,
	nhead=4,
	num_encoder_layers=3,
	num_decoder_layers=3,
	norm_first=True,
	activation='gelu'
)

src = torch.rand((15, 4, 256))
tgt = torch.rand((12, 4, 256))

out = transformer(src, tgt)

print(out.shape)

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

torch.Size([12, 4, 256])

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

класс TransformerEncoder,
который реализует только кодировщик трансформера
класс TransformerEncoderLayer,
который реализует один слой кодировщика
класс TransformerDecoder,
который реализует только декодировщик трансформера
класс MultiheadAttention,
который реализует механизм многоголового внимания