Функция torch.nn.TransformerEncoder в PyTorch

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

torch.nn.TransformerEncoderЭто модуль энкодера Transformer в PyTorch.

Он состоит из нескольких слоёв TransformerEncoderLayer и используется для обработки входных последовательностей.

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

torch.nn.TransformerEncoder(encoder_layer, num_layers, norm=None)

Parameters

  • encoder_layer: Слой энкодера
  • num_layers: Количество слоёв

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

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

Пример

import torch
import torch.nn as nn

# Один слой энкодера
encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)

# Полный энкодер
transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)

# Ввод
x = torch.randn(32, 100, 512)  # batch, seq, d_model
output = transformer_encoder(x)

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

Пример 2: Классификация текста

Пример

import torch
import torch.nn as nn

class TransformerClassifier(nn.Module):
    def __init__(self, vocab_size, d_model=256, nhead=4, num_layers=4):
        super(TransformerClassifier, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = nn.Parameter(torch.randn(1, 512, d_model) * 0.1)

        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, batch_first=True)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)
        self.fc = nn.Linear(d_model, 10)

    def forward(self, x):
        # Простое позиционное кодирование
        x = self.embedding(x) + self.pos_encoder[:, :x.size(1), :]
        x = self.encoder(x)
        # Взять первый токен
        return self.fc(x[:, 0, :])

model = TransformerClassifier(10000)
x = torch.randint(0, 10000, (32, 100))
output = model(x)

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

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

  • BERT: основы энкодера
  • Классификация текста
  • Извлечение признаков

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

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