Функция torch.nn.EmbeddingBag в PyTorch

PyTorch torch.nn 参考手册Справочное руководство PyTorch torch.nn

torch.nn.EmbeddingBagЭто модуль embedding bag в PyTorch.

Он агрегирует несколько embedding-векторов в один вектор, часто используется для классификации текста и быстрой обработки.

Определение функции

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)

Параметры:

  • mode: Способ агрегации, необязательный'mean'、'sum'、'max'

Примеры использования

Пример 1: Базовое использование

Пример

import torch
import torch.nn as nn

# Embedding bag
ebag = nn.EmbeddingBag(1000, 128, mode='mean')

# Индексы слов
indices = torch.tensor([[1, 2, 3], [4, 5]])

# Смещения (указывающие границы каждого предложения)
offsets = torch.tensor([0, 3])

output = ebag(indices, offsets)

print("Форма входных данных:", indices.shape)
print("Форма выхода:", output.shape)

Пример 2: Быстрая классификация текста

Пример

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)

# Имитация данных
text = torch.randint(0, 10000, (100,))
offset = torch.tensor([0, 30, 60, 100])

output = model(text, offset)
print("Форма выхода:", output.shape)

Пример 3: Различные методы агрегации

Пример

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)

Сценарии использования

  • Классификация текста: FastText
  • Модель мешка слов: Быстрая агрегация
  • Крупномасштабный текст: Эффективная обработка

Примечание: необходимо указать параметр offsets, чтобы обозначить границы предложений.

PyTorch torch.nn 参考手册Справочное руководство PyTorch torch.nn

Другие расширения