Función PyTorch torch.nn.TransformerEncoder

PyTorch torch.nn 参考手册Manual de referencia de PyTorch torch.nn

torch.nn.TransformerEncoderEs el módulo codificador Transformer en PyTorch.

Se compone de múltiples TransformerEncoderLayer y se utiliza para procesar secuencias de entrada.

Definición de la función

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

Parámetros

  • encoder_layer: capa del codificador
  • num_layers: número de capas

Ejemplo de uso

Ejemplo 1: Uso básico

Ejemplo

import torch
import torch.nn as nn

# Capa de codificador única
encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)

# Codificador completo
transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)

# Entrada
x = torch.randn(32, 100, 512)  # batch, seq, d_model
output = transformer_encoder(x)

print("Forma de entrada:", x.shape)
print("Forma de salida:", output.shape)

Ejemplo 2: Clasificación de texto

Ejemplo

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):
        # Codificación posicional simple
        x = self.embedding(x) + self.pos_encoder[:, :x.size(1), :]
        x = self.encoder(x)
        # Tomar el primer token
        return self.fc(x[:, 0, :])

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

print("Forma de salida:", output.shape)

Escenarios de uso

  • BERT: Base del codificador
  • Clasificación de texto
  • Extracción de características

PyTorch torch.nn 参考手册Manual de referencia de PyTorch torch.nn

Otras extensiones