Fonction torch.nn.EmbeddingBag de PyTorch

PyTorch torch.nn 参考手册Manuel de référence torch.nn de PyTorch

torch.nn.EmbeddingBagIl s'agit d'un module d'embedding bag dans PyTorch.

Il agrège plusieurs vecteurs d'embedding en un seul vecteur, couramment utilisé pour la classification de texte et le traitement rapide.

Définition de fonction

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)

Paramètres :

  • mode: Méthode d'agrégation, facultative'mean'、'sum'、'max'

Exemples d'utilisation

Exemple 1 : Utilisation de base

Exemple

import torch
import torch.nn as nn

# Sac d'embedding
ebag = nn.EmbeddingBag(1000, 128, mode='mean')

# Indices de mots
indices = torch.tensor([[1, 2, 3], [4, 5]])

# Décalages (indiquant les limites de chaque phrase)
offsets = torch.tensor([0, 3])

output = ebag(indices, offsets)

print("Forme d'entrée :", indices.shape)
print("Forme de sortie :", output.shape)

Exemple 2 : Classification rapide de texte

Exemple

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)

# Simuler des données
text = torch.randint(0, 10000, (100,))
offset = torch.tensor([0, 30, 60, 100])

output = model(text, offset)
print("Forme de sortie :", output.shape)

Exemple 3 : Différentes méthodes d'agrégation

Exemple

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)

Cas d'utilisation

  • Classification de texte: FastText
  • Modèle de sac de mots: Agrégation rapide
  • Texte à grande échelle: Traitement efficace

Remarque : Le paramètre offsets doit être fourni pour indiquer les limites de phrase.

PyTorch torch.nn 参考手册Manuel de référence torch.nn de PyTorch

Autres extensions