Función torch.nn.Embedding de PyTorch

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

torch.nn.EmbeddingEs el módulo de PyTorch utilizado para el embedding de palabras.

Mapea índices de vocabulario discretos a un espacio vectorial continuo; es una de las operaciones más básicas en el procesamiento del lenguaje natural.

Definición de la función

torch.nn.Embedding(num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False)

Descripción de parámetros:

  • num_embeddings(int): Tamaño del vocabulario, es decir, número de filas de la matriz de embedding.
  • embedding_dim(int): Dimensión de cada vector de embedding.
  • padding_idx(int): Especifica el índice de padding, cuyo vector de embedding es un vector cero. Por defecto es None.
  • max_norm(float): Si no es None, los vectores de embedding se normalizan a esta norma. Por defecto es None.
  • norm_type(float): Orden de la norma a calcular. Por defecto es 2.0.
  • scale_grad_by_freq(bool): Indica si se escalan los gradientes según la frecuencia de las palabras. Por defecto es False.
  • sparse(bool): Indica si la matriz de pesos es una matriz dispersa. Por defecto es False.

Atributos:

  • weight(Tensor): Pesos aprendibles con forma (num_embeddings, embedding_dim).

Ejemplos de uso

Ejemplo 1: Uso básico

Crear y usar embeddings de palabras:

Ejemplo

import torch
import torch.nn as nn

# Crear la capa de embedding: vocabulario 10000, dimensión de embedding 256
embedding = nn.Embedding(num_embeddings=10000, embedding_dim=256)

# Índices de palabras (comenzando desde 0)
# Forma: (batch, seq_len)
input_indices = torch.tensor([[12, 45, 678], [901, 23, 56]])

# Obtener los vectores de embedding mediante la tabla de consulta
output = embedding(input_indices)

print("Forma del índice de entrada:", input_indices.shape)
print("Forma del embedding de salida:", output.shape)  # (2, 3, 256)

# Ver la forma de la matriz de embedding
print("Forma de la matriz de embedding:", embedding.weight.shape)

Ejemplo 2: Usando padding_idx

Especificar el índice de padding:

Ejemplo

import torch
import torch.nn as nn

# Crear una capa de embedding con padding
embedding = nn.Embedding(num_embeddings=1000, embedding_dim=64, padding_idx=0)

# 0 se usa como padding
input_indices = torch.tensor([[1, 2, 3], [4, 0, 0]])  # La segunda oración tiene padding

output = embedding(input_indices)

print("Forma de entrada:", input_indices.shape)
print("Forma de salida:", output.shape)
print("Vector de embedding del padding:", output[1, 1].tolist())  # Todo ceros
print("Vector de embedding de no padding:", output[0, 0].tolist()[:5])  # Distinto de cero

Ejemplo 3: Vectores de palabras preentrenados

Cargar vectores de palabras preentrenados:

Ejemplo

import torch
import torch.nn as nn
import numpy as np

# Simular vectores de palabras preentrenados (en realidad se pueden usar GloVe, Word2Vec, etc.)
vocab_size = 1000
embedding_dim = 300

# Inicialización aleatoria (en realidad se deben cargar los vectores preentrenados)
pretrained_weights = np.random.randn(vocab_size, embedding_dim).astype('float32')

# Crear la capa de embedding
embedding = nn.Embedding(vocab_size, embedding_dim)

# Cargar los pesos preentrenados
embedding.weight.data = torch.from_numpy(pretrained_weights)

# Congelar la capa de embedding (no participa en el entrenamiento)
embedding.weight.requires_grad = False

print("Capa de embedding entrenable:", embedding.weight.requires_grad)
print("Forma de la matriz de embedding:", embedding.weight.shape)

Ejemplo 4: Limitar la norma del vector de embedding

Usar max_norm para limitar la norma del vector:

h2 class="example">Ejemplo
import torch
import torch.nn as nn

# Limitar la norma máxima del vector de embedding a 1.0
embedding = nn.Embedding(1000, 64, max_norm=1.0)

# Entrada
input_indices = torch.tensor([1, 2, 3])

# Norma de los pesos originales
original_norm = embedding.weight.data.norm(dim=1)[:3]
print("Norma de los pesos originales:", original_norm.tolist())

# Norma del vector después de la consulta de la tabla
output = embedding(input_indices)
output_norm = output.norm(dim=1)
print("Norma del vector de salida:", output_norm.tolist())

Ejemplo 5: Modelo completo de clasificación de texto

Clasificación de texto usando capa de embedding:

Ejemplo

import torch
import torch.nn as nn

class TextClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes=2):
        super(TextClassifier, self).__init__()
        # Capa de embedding de palabras
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        # LSTM
        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
        # Clasificador
        self.classifier = nn.Linear(hidden_dim, num_classes)

    def forward(self, x):
        # x: (batch, seq_len)
        embedded = self.embedding(x)  # (batch, seq_len, embed_dim)

        # LSTM toma la última salida
        _, (hidden, _) = self.lstm(embedded)
        hidden = hidden[-1]  # Estado oculto de la última capa

        # Clasificación
        logits = self.classifier(hidden)
        return logits

# Parámetros
VOCAB_SIZE = 10000
EMBED_DIM = 128
HIDDEN_DIM = 256

model = TextClassifier(VOCAB_SIZE, EMBED_DIM, HIDDEN_DIM)

# Entrada: batch=4, longitud de secuencia=50
input_ids = torch.randint(1, VOCAB_SIZE, (4, 50))
output = model(input_ids)

print("Estructura del modelo:")
print(model)
print("nForma de entrada:", input_ids.shape)
print("Forma de salida:", output.shape)

Diferencia entre Embedding y EmbeddingBag

Tipo Entrada Salida Escenario de aplicación
nn.Embedding Índice de palabra Secuencia de vectores de palabras Modelos secuenciales, LSTM, Transformer
nn.EmbeddingBag Índice de palabra + offset Vector agregado Clasificación de texto, procesamiento rápido

Preguntas frecuentes

Q1: ¿Cómo elegir la dimensión de embedding?

  • Conjunto de datos pequeño: 50-100 dimensiones
  • Conjunto de datos mediano: 100-300 dimensiones
  • Conjunto de datos grande: 300-500 dimensiones

Q2: ¿Cuál es la función de padding_idx?

Establece el vector de embedding del índice especificado como vector cero y no calcula su gradiente durante la propagación hacia atrás.

Q3: ¿Cuándo congelar la capa de embedding?

Al usar vectores de palabras preentrenados, normalmente se congela primero durante un tiempo de entrenamiento y luego se ajusta finamente (fine-tuning).

Escenarios de uso

nn.EmbeddingLos principales escenarios de aplicación incluyen:

  • Representación de vectores de palabras: convertir palabras en vectores densos
  • Clasificación de texto: como capa de entrada para modelos NLP
  • Modelos secuenciales: entrada de LSTM, GRU
  • Sistema de recomendación: representación de embeddings de usuarios y elementos

Nota: embedding.weight es un parámetro aprendible que se puede entrenar directamente en el optimizador.

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

Otras extensiones