Función PyTorch torch.nn.TransformerEncoderLayer

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

nn.TransformerEncoderLayerEs la estructura de una sola capa del codificador Transformer.

Contiene autoatención y una red de avance (feed-forward); es la unidad básica que constituye el codificador completo.

Definición de la función

torch.nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='gelu', batch_first=True)

Parámetros

  • d_model: Dimensión del modelo
  • nhead: Número de cabezas de atención
  • dim_feedforward: Dimensión de la capa oculta de FFN

Ejemplo de uso

Ejemplo 1: Uso básico

Ejemplo

import torch
import torch.nn as nn

# Capa del codificador
layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)

# Entrada
x = torch.randn(32, 100, 512)
output = layer(x)

print("Entrada:", x.shape, "-> Salida:", output.shape)

Ejemplo 2: Función de activación personalizada

Ejemplo

import torch
import torch.nn as nn

# Usando activación ReLU
layer_relu = nn.TransformerEncoderLayer(d_model=256, nhead=4, activation='relu', batch_first=True)

x = torch.randn(8, 50, 256)
out = layer_relu(x)

print("Salida de activación ReLU:", out.shape)

Ejemplo 3: Apilamiento de múltiples capas

Ejemplo

import torch
import torch.nn as nn

# Apilando múltiples capas
layer = nn.TransformerEncoderLayer(256, 4, batch_first=True)
encoder = nn.Sequential(
    layer,
    nn.TransformerEncoderLayer(256, 4, batch_first=True),
    nn.TransformerEncoderLayer(256, 4, batch_first=True)
)

x = torch.randn(4, 30, 256)
out = encoder(x)

print("Salida del codificador de 3 capas:", out.shape)

Casos de uso

  • Construcción del codificador
  • Modelos como BERT
  • Comprensión de texto

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

Otras extensiones