Función PyTorch torch.nn.EmbeddingBag

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

torch.nn.EmbeddingBagEs el módulo de bolsa de embeddings en PyTorch.

Agrega múltiples vectores de embedding en un solo vector, comúnmente utilizado para clasificación de texto y procesamiento rápido.

Definición de la función

torch.nn.EmbeddingBag(num_embeddings, embedding_dim, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, mode='mean', sparse=False, include_last_offset=False)

Parámetros:

  • mode: método de agregación, opcional'mean'、'sum'、'max'

Ejemplos de uso

Ejemplo 1: Uso básico

Ejemplo

import torch
import torch.nn as nn

# Bolsa de embeddings
ebag = nn.EmbeddingBag(1000, 128, mode='mean')

# Índice de palabras
indices = torch.tensor([[1, 2, 3], [4, 5]])

# Desplazamientos (indican los límites de cada oración)
offsets = torch.tensor([0, 3])

output = ebag(indices, offsets)

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

Ejemplo 2: Clasificación rápida de texto

Ejemplo

import torch
import torch.nn as nn

class FastText(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_classes):
        super(FastText, self).__init__()
        self.embedding = nn.EmbeddingBag(vocab_size, embed_dim, mode='mean')
        self.fc = nn.Linear(embed_dim, num_classes)

    def forward(self, text, offset):
        x = self.embedding(text, offset)
        return self.fc(x)

model = FastText(10000, 128, 2)

# Datos simulados
text = torch.randint(0, 10000, (100,))
offset = torch.tensor([0, 30, 60, 100])

output = model(text, offset)
print(Forma de salida:, output.shape)

Ejemplo 3: Diferentes métodos de agregación

Ejemplo

import torch
import torch.nn as nn

for mode in ['mean', 'sum', 'max']:
    ebag = nn.EmbeddingBag(100, 32, mode=mode)
    indices = torch.tensor([1, 2, 3, 4])
    offsets = torch.tensor([0, 2])
    out = ebag(indices, offsets)
    print(f"{mode}:", out.shape)

Casos de uso

  • Clasificación de texto: FastText
  • Modelo de bolsa de palabras: agregación rápida
  • Textos a gran escala: procesamiento eficiente

Nota: se debe proporcionar el parámetro offsets para indicar los límites de las oraciones.

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

Otras extensiones