介绍

torch.nn.Embedding 是 PyTorch 中用于实现 嵌入层（Embedding Layer）的模块，核心作用是将 离散的整数索引（比如单词的编号 0、1、2…）映射为连续的低维稠密向量（比如 128 维的向量）。

简单来说：

• 输入：一组整数（比如 [1, 5, 3]，代表 3 个单词的索引）
• 输出：一组对应的稠密向量（比如 3 个 128 维的向量）
• 本质：一个可训练的查找表（lookup table），表中的每一行对应一个索引的嵌入向量，训练过程中这些向量会不断优化。

应用场景

• 自然语言处理（NLP）：将单词/字符的索引映射为词向量（如 Transformer、LSTM 的输入层）；
• 推荐系统：将用户 ID、物品 ID 映射为稠密向量，用于后续的相似度计算；
• 计算机视觉：将类别索引（如标签编号）映射为向量，辅助分类任务。

语法

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

关键参数：

属性

torch.nn.Embedding 作为 nn.Module 的子类，除了继承的通用属性（如 training、_device 等），最核心的是和嵌入向量直接相关的属性，以下是新手必须掌握的关键属性：

只读属性 vs 可修改属性

• num_embeddings、embedding_dim 是 只读属性（初始化后无法修改），如果要调整词典大小/嵌入维度，只能重新创建 Embedding 层。
• padding_idx、max_norm 可以修改（如 embedding.padding_idx = 1），修改后会影响后续的前向传播。
• weight 是可修改的，既可以直接改值，也可以修改其 requires_grad 状态。

填充索引的权重特性

如果设置了 padding_idx，该索引对应的权重在反向传播时梯度会被置 0（即不会更新），保证 padding 向量始终是初始值（默认全 0）。

# 验证 padding_idx=0 的权重梯度
input_idx = torch.tensor([0, 1, 2])
output = embedding(input_idx)
output.sum().backward()  # 反向传播
print("\npadding_idx=0 的权重梯度:", embedding.weight.grad[0])  # 输出：全 0
print("非 padding 索引 1 的权重梯度:", embedding.weight.grad[1])  # 输出：非 0

例子

创建、常用属性

import torch.nn as nn
# 创建 Embedding 层
# 词典大小：4（索引 0-4），嵌入维度：3（每个索引映射为 3 维向量），即：权重是4行3列
embedding = nn.Embedding(num_embeddings=4, embedding_dim=3)

print(embedding)
print("嵌入词典的大小：", embedding.num_embeddings)  # 4
print("每个嵌入向量的维度：", embedding.embedding_dim)  # 3

print("\n=== 权重矩阵（weight） ===")
print("权重weight：", embedding.weight)  # 4行3列的张量
# weight 是 nn.Parameter 类型，本质是张量，形状 [num_embeddings, embedding_dim]
print("weight 形状:", embedding.weight.shape)  # 输出：torch.Size([10, 4])
print("weight 数据类型:", embedding.weight.dtype)  # 输出：torch.float32
# 查看索引 1 对应的嵌入向量（权重矩阵的第 1 行）
print("索引 1 对应的嵌入向量:", embedding.weight[1])

执行结果：

Embedding(4, 3)
嵌入词典的大小： 4
每个嵌入向量的维度： 3

=== 权重矩阵（weight） ===
权重weight： Parameter containing:
tensor([[ 0.5295, -1.9354,  0.3722],
        [ 0.6719, -1.4125, -0.0068],
        [-0.3018, -1.1436, -0.7975],
        [ 0.3823, -1.0013,  0.4396]], requires_grad=True)
weight 形状: torch.Size([4, 3])
weight 数据类型: torch.float32
索引 1 对应的嵌入向量: tensor([ 0.6719, -1.4125, -0.0068], grad_fn=<SelectBackward0>)

创建并使用 Embedding 层

import torch
import torch.nn as nn

# 1. 定义 Embedding 层
# 词典大小：10（索引 0-9），嵌入维度：4（每个索引映射为 4 维向量）
embedding = nn.Embedding(num_embeddings=10, embedding_dim=4)

# 2. 准备输入（整数索引，形状：[批量大小, 序列长度]）
# 比如批量大小为 2，每个序列有 3 个索引
input_indices = torch.tensor([[1, 3, 5], [2, 4, 6]])

# 3. 前向传播（获取嵌入向量）
embedded_output = embedding(input_indices)

# 输出结果说明
print("输入索引形状：", input_indices.shape)  # torch.Size([2, 3])
print("嵌入输出形状：", embedded_output.shape)  # torch.Size([2, 3, 4])
print("索引 1 对应的嵌入向量：")
print(embedded_output[0, 0])  # 第一行第一个索引（1）的 4 维向量

带 padding_idx 的 Embedding 层

# 指定 padding_idx=0，索引 0 的向量训练时不会更新
embedding = nn.Embedding(num_embeddings=10, embedding_dim=4, padding_idx=0)

# 输入包含 padding 索引 0
input_indices = torch.tensor([[0, 1, 3], [0, 0, 2]])
embedded_output = embedding(input_indices)

# 验证：padding_idx=0 的向量梯度会被置 0（训练时不更新）
embedded_output.sum().backward()  # 模拟反向传播
print("嵌入层权重的梯度：")
print(embedding.weight.grad)
# 输出中，索引 0 对应的行全为 0，其他索引的行有梯度值

特性与注意事项

1. 可训练性：Embedding 层的权重（embedding.weight）是可训练的参数，会在反向传播中更新，这也是“词向量学习”的核心。
2. 输入要求：输入必须是 整数类型的张量（torch.long/torch.int64），不能是浮点型。
3. 预训练词向量加载：实际应用中，可加载预训练的词向量（如 Word2Vec、GloVe）初始化 Embedding 层：

   # 假设 pretrained_embeddings 是形状为 [10, 4] 的预训练向量张量
   pretrained_embeddings = torch.randn(10, 4)
   embedding = nn.Embedding(10, 4)
   embedding.weight.data = pretrained_embeddings  # 加载预训练权重
   # 如果不想更新预训练向量，设置 requires_grad=False
   embedding.weight.requires_grad = False
   

4. 批量处理：支持批量输入，输入形状一般为 [batch_size, seq_len]，输出形状为 [batch_size, seq_len, embedding_dim]。

总结

1. torch.nn.Embedding 是将 整数索引 映射为 低维稠密向量 的可训练层，核心是一个查找表；
2. 核心参数是 num_embeddings（词典大小）和 embedding_dim（向量维度），padding_idx 是 NLP 中常用的实用参数；
3. 它是 NLP 模型的基础组件，支持加载预训练向量，也可通过训练学习专属的嵌入表示。

原文出处：/show_1GW2aoJ5nKFn.html