Класс nn.Embedding

Класс nn.Embedding представляет собой модуль, который хранит таблицу обучаемых векторов (эмбеддингов) фиксированного размера. Он используется для преобразования дискретных индексов (например, слов или категорий) в плотные непрерывные векторы. Первым параметром передается количество записей в словаре (num_embeddings), вторым - размерность каждого вектора (embedding_dim). Дополнительно можно указать параметр padding_idx для индекса, который будет игнорироваться при обновлении, или max_norm для нормирования векторов.

Синтаксис

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

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

Создадим слой эмбеддингов для словаря из 10 слов, каждый вектор размерности 3, и получим представление для индексов 2 и 5:

import torch
import torch.nn as nn

embedding = nn.Embedding(num_embeddings=10, embedding_dim=3)
torch.manual_seed(0)
indices = torch.tensor([2, 5])
res = embedding(indices)
print(res)

Результат выполнения кода:

tensor([
	[-0.0879, -0.7160, -0.2657],
	[-0.7109, -0.5042, -0.5213],
], grad_fn=<EmbeddingBackward0>)

Пример с параметром padding_idx

Создадим эмбеддинг, где индекс 0 будет игнорироваться при обучении (вектор для этого индекса останется нулевым):

import torch
import torch.nn as nn

embedding = nn.Embedding(num_embeddings=5, embedding_dim=4, padding_idx=0)
torch.manual_seed(0)
indices = torch.tensor([0, 1, 2])
res = embedding(indices)
print(res)

Результат выполнения кода:

tensor([
	[ 0.0000,  0.0000,  0.0000,  0.0000],
	[-0.0879, -0.7160, -0.2657, -0.3729],
	[-0.3805, -0.5520, -0.3736,  0.5322],
], grad_fn=<EmbeddingBackward0>)

Пример с нормированием векторов

Используем параметр max_norm для ограничения L2-нормы векторов значением 1.0:

import torch
import torch.nn as nn

embedding = nn.Embedding(num_embeddings=5, embedding_dim=3, max_norm=1.0)
torch.manual_seed(0)
indices = torch.tensor([1, 3])
res = embedding(indices)
print(res)
print(torch.norm(res, dim=1))

Результат выполнения кода:

tensor([
	[-0.1104, -0.8991, -0.3337],
	[ 0.0058, -0.1792, -0.4909],
], grad_fn=<EmbeddingBackward0>)
tensor([0.9659, 0.5227], grad_fn=<NormBackward1>)

Пример работы со слоем в модели

Создадим простую модель, которая использует эмбеддинги для обработки последовательности индексов:

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.embedding = nn.Embedding(10, 5)
		self.fc = nn.Linear(5, 2)
	
	def forward(self, x):
		x = self.embedding(x)
		x = x.mean(dim=1)
		return self.fc(x)

torch.manual_seed(0)
model = SimpleModel()
indices = torch.tensor([[1, 2, 3], [4, 5, 6]])
res = model(indices)
print(res)

Результат выполнения кода:

tensor([
	[-0.5461,  0.2154],
	[ 0.5602, -0.4681],
], grad_fn=<AddmmBackward0>)

Смотрите также

класс nn.EmbeddingBag,
который агрегирует эмбеддинги по сумме или среднему
класс nn.Linear,
который применяет линейное преобразование к входным данным
функцию nn.Dropout,
которая применяется для регуляризации после эмбеддингов
слой nn.LayerNorm,
который нормализует выходы эмбеддингов для стабильности обучения