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

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

torch.nn.TransformerЭто полная модель Transformer в PyTorch.

Он содержит энкодер и декодер и может использоваться для задач последовательность-в-последовательность.

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

torch.nn.Transformer(d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dim_feedforward=2048, dropout=0.1, activation='gelu', batch_first=True)

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

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

Пример

import torch
import torch.nn as nn

# Создать Transformer
transformer = nn.Transformer(d_model=512, nhead=8, batch_first=True)

# Вход энкодера
src = torch.randn(10, 32, 512)  # (seq, batch, d_model)
# Вход декодера
tgt = torch.randn(20, 32, 512)

output = transformer(src, tgt)

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

Пример 2: Простая модель перевода

Пример

import torch
import torch.nn as nn

class TransformerMT(nn.Module):
    def __init__(self, vocab_size, d_model=256, nhead=4):
        super(TransformerMT, self).__init__()
        self.d_model = d_model
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer = nn.Transformer(d_model, nhead, batch_first=True)
        self.fc = nn.Linear(d_model, vocab_size)

    def forward(self, src, tgt):
        src = self.embedding(src) * (self.d_model ** 0.5)
        tgt = self.embedding(tgt) * (self.d_model ** 0.5)
        out = self.transformer(src, tgt)
        return self.fc(out)

model = TransformerMT(10000)
src = torch.randint(0, 10000, (32, 50))
tgt = torch.randint(0, 10000, (32, 40))
output = model(src, tgt)

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

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

  • Машинный перевод
  • Генерация текста
  • Последовательность-в-последовательность

Примечание: Когда batch_first=True, форма входных данных — (batch, seq, d_model).

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

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