当前位置：首页 > news >正文

pytorch 中是如何实现embeding 的

news 2025/9/9 5:41:49

嵌入层（Embedding）说明

一个简单的查找表，用于存储固定字典和固定大小的嵌入向量。

该模块通常用于存储词嵌入，并通过索引检索对应的词嵌入。模块的输入是一组索引，输出则是对应的词嵌入向量。

参数（Args）

num_embeddings (int)：嵌入字典的大小（即离散特征的总类别数，如词汇表大小）。
embedding_dim (int)：每个嵌入向量的维度。
padding_idx (int, 可选)：若指定该参数，则索引为padding_idx的嵌入向量不会对梯度产生贡献；因此，训练过程中该索引对应的嵌入向量不会被更新，即始终保持为固定的 “填充向量”。对于新创建的 Embedding 层，padding_idx对应的嵌入向量默认全为 0，但也可修改为其他值作为填充向量。
max_norm (float, 可选)：若指定该参数，则所有范数（norm）大于max_norm的嵌入向量，都会被重新归一化，使其范数等于max_norm。
norm_type (float, 可选)：用于计算max_norm的 p 范数（p-norm）中的 “p 值”，默认值为2（即 L2 范数）。
scale_grad_by_freq (bool, 可选)：若指定该参数为True，则梯度会按小批量（mini-batch）中单词频率的倒数进行缩放。默认值为False。
sparse (bool, 可选)：若为True，则相对于weight矩阵的梯度会是一个稀疏张量（sparse tensor）。关于稀疏梯度的更多细节，可参考下方的 “注意事项（Note）”。

属性（Attributes）

weight (Tensor)：模块的可学习权重，形状为(num_embeddings, embedding_dim)，初始值服从标准正态分布（ $N (0, 1)$ ）。

形状（Shape）

输入（Input）：形状为(*)的整数张量（IntTensor）或长整数张量（LongTensor），其中*表示任意形状，张量内元素为待检索的索引。
输出（Output）：形状为(*, H)的张量，其中*与输入形状保持一致，H等于embedding_dim（即嵌入向量维度）。

注意事项（.. note::）

需注意，目前仅有少数优化器支持稀疏梯度（sparse gradients），具体包括：
- optim.SGD（支持 CPU 和 CUDA 设备）
- optim.SparseAdam（支持 CPU 和 CUDA 设备）
- optim.Adagrad（仅支持 CPU 设备）
当max_norm不为None时，Embedding 层的前向传播（forward）方法会原地（in-place）修改weight张量。由于用于梯度计算的张量不能被原地修改，因此当max_norm不为None时，若要在调用 Embedding 层的前向方法前，对Embedding.weight执行可微操作（differentiable operation），需先对Embedding.weight进行克隆（clone）。示例如下：

n, d, m = 3, 5, 7
embedding = nn.Embedding(n, d, max_norm=True)
W = torch.randn((m, d), requires_grad=True)
idx = torch.tensor([1, 2])
a = embedding.weight.clone() @ W.t()  # 需克隆weight，否则该操作无法计算梯度
b = embedding(idx) @ W.t()  # 原地修改weight
out = (a.unsqueeze(0) + b.unsqueeze(1))
loss = out.sigmoid().prod()
loss.backward()

示例（Examples::）

>>> # 一个包含10个3维嵌入向量的Embedding模块
>>> embedding = nn.Embedding(10, 3)
>>> # 一个批量（batch）：2个样本，每个样本包含4个索引
>>> input = torch.LongTensor([[1, 2, 4, 5], [4, 3, 2, 9]])
>>> # xdoctest: +IGNORE_WANT("non-deterministic")（注：用于忽略非确定性输出的测试检查）
>>> embedding(input)
tensor([[[-0.0251, -1.6902,  0.7172],[-0.6431,  0.0748,  0.6969],[ 1.4970,  1.3448, -0.9685],[-0.3677, -2.7265, -0.1685]],[[ 1.4970,  1.3448, -0.9685],[ 0.4362, -0.4004,  0.9400],[-0.6431,  0.0748,  0.6969],[ 0.9124, -2.3616,  1.1151]]])>>> # 带有padding_idx的示例
>>> embedding = nn.Embedding(10, 3, padding_idx=0)
>>> input = torch.LongTensor([[0, 2, 0, 5]])
>>> embedding(input)
tensor([[[ 0.0000,  0.0000,  0.0000],  # padding_idx=0对应的向量全为0[ 0.1535, -2.0309,  0.9315],[ 0.0000,  0.0000,  0.0000],  # 再次出现padding_idx=0，向量仍为0[-0.1655,  0.9897,  0.0635]]])>>> # 修改“填充向量（pad vector）”的示例
>>> padding_idx = 0
>>> embedding = nn.Embedding(3, 3, padding_idx=padding_idx)
>>> embedding.weight  # 初始时，padding_idx=0对应的向量全为0
Parameter containing:
tensor([[ 0.0000,  0.0000,  0.0000],[-0.7895, -0.7089, -0.0364],[ 0.6778,  0.5803,  0.2678]], requires_grad=True)
>>> with torch.no_grad():  # 禁用梯度计算，避免修改操作影响训练
...     embedding.weight[padding_idx] = torch.ones(3)  # 将填充向量改为全1
>>> embedding.weight  # 修改后，padding_idx=0对应的向量变为全1
Parameter containing:
tensor([[ 1.0000,  1.0000,  1.0000],[-0.7895, -0.7089, -0.0364],[ 0.6778,  0.5803,  0.2678]], requires_grad=True)