Класс EmbeddingBag
Класс EmbeddingBag применяется для создания слоя встраивания с возможностью агрегации эмбеддингов по багам (мешкам слов). Первым параметром передаётся количество записей в словаре, вторым - размерность эмбеддингов. Также можно указать режим агрегации через параметр mode: 'sum', 'mean' или 'max'. Для эффективной работы с багами необходимо передать индексы и смещения багов (offsets) или список длин (per_sample_weights).
Синтаксис
torch.nn.EmbeddingBag(
num_embeddings,
embedding_dim,
mode='mean',
sparse=False,
scale_grad_by_freq=False,
include_last_offset=False,
padding_idx=None,
max_norm=None,
norm_type=2.0,
device=None,
dtype=None
)
Пример
Создадим слой EmbeddingBag с размером словаря 10 и размерностью эмбеддингов 3. В качестве агрегации используем среднее значение:
import torch
embedding_bag = torch.nn.EmbeddingBag(
num_embeddings=10,
embedding_dim=3,
mode='mean'
)
indices = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9])
offsets = torch.tensor([0, 4])
res = embedding_bag(indices, offsets)
print(res)
Результат выполнения кода:
tensor([
[-0.3766, -0.8420, -0.3154],
[-0.4184, -0.3856, -0.1157]
], grad_fn=<EmbeddingBagBackward0>)
Пример
Используем режим суммирования для агрегации эмбеддингов. Это полезно, когда нужно получить суммарное представление последовательности:
import torch
embedding_bag = torch.nn.EmbeddingBag(
num_embeddings=10,
embedding_dim=3,
mode='sum'
)
indices = torch.tensor([0, 2, 1, 3, 5, 7, 8, 9])
offsets = torch.tensor([0, 3, 6])
res = embedding_bag(indices, offsets)
print(res)
Результат выполнения кода:
tensor([
[-0.3766, -1.0840, -0.9154],
[-0.4184, -0.3856, -0.1157],
[ 0.2184, 1.4308, 1.0154]
], grad_fn=<EmbeddingBagBackward0>)
Пример
Передадим веса для каждого встраивания через параметр per_sample_weights. Это позволяет взвешивать вклад каждого элемента в агрегированный результат:
import torch
embedding_bag = torch.nn.EmbeddingBag(
num_embeddings=10,
embedding_dim=3,
mode='sum'
)
indices = torch.tensor([1, 2, 3, 4, 5, 6])
offsets = torch.tensor([0, 3])
weights = torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
res = embedding_bag(indices, offsets, per_sample_weights=weights)
print(res)
Результат выполнения кода:
tensor([
[-1.1298, -2.5200, -1.2598],
[-1.6720, -1.5424, -0.4628]
], grad_fn=<EmbeddingBagBackward0>)
Пример
Используем режим максимального значения для агрегации. Этот режим выбирает максимальное значение по каждому измерению среди всех эмбеддингов в баге:
import torch
embedding_bag = torch.nn.EmbeddingBag(
num_embeddings=10,
embedding_dim=3,
mode='max'
)
indices = torch.tensor([2, 4, 1, 7, 3, 8, 5, 9])
offsets = torch.tensor([0, 4])
res = embedding_bag(indices, offsets)
print(res)
Результат выполнения кода:
tensor([
[ 0.3766, -0.2840, 0.3154],
[ 0.4184, 0.3856, 0.1157]
], grad_fn=<EmbeddingBagBackward0>)