nn.Embedding 是 PyTorch 中处理离散符号的核心模块，它将离散的整数索引映射到连续的稠密向量空间中，为自然语言处理、推荐系统等任务提供了基础的表示学习能力。通过可学习的查找表机制，nn.Embedding实现了符号数据的分布式表示，使得模型能够捕捉词汇间的语义关系和语法特性。

一、nn.Embedding介绍

1.1 核心作用

• 映射功能：将离散的整数索引（如词汇ID）转换为连续的向量表示
• 表示学习：通过训练学习每个符号的分布式表示
• 维度转换：从一维的索引空间到高维的连续向量空间

1.2 基本结构

输入：单词索引 [2, 0, 3]  ← 表示 ["apple", "the", "banana"]

       ↓
+---------------------------+
|     嵌入层权重矩阵         |
|  (形状: vocab_size × embedding_dim) |
|                           |
| 行0: the    → [0.1, 0.2, 0.3]  |
| 行1: cat    → [0.4, 0.5, 0.6]  |
| 行2: apple  → [0.7, 0.8, 0.9]  | ← 查找第2行
| 行3: banana → [1.0, 1.1, 1.2]  | ← 查找第3行
| 行4: dog    → [1.3, 1.4, 1.5]  |
| ...        → ...              |
+---------------------------+
       ↓
输出：[[0.7, 0.8, 0.9],  ← apple的嵌入
       [0.1, 0.2, 0.3],  ← the的嵌入
       [1.0, 1.1, 1.2]]  ← banana的嵌入

内部机制：

• 查找表：本质是一个可学习的权重矩阵，形状为 (num_embeddings, embedding_dim)
• 索引查询：通过整数索引直接查找对应的行向量
• 梯度传播：在训练过程中通过反向传播更新权重矩阵

1.3 主要参数详解

参数	类型	默认值	描述
num_embeddings	int	必须指定，应设置为大于或等于输入中可能出现的最大索引值	词汇表大小，即最大索引值+1，表示嵌入矩阵的行数
embedding_dim	int	必须指定	每个嵌入向量的维度，维度越高表示能力越强，但需要更多数据和计算资源
padding_idx	int	None	指定填充符号的索引，对应的向量初始为0且在训练中通常不更新。
max_norm	float	None	对嵌入向量进行范数约束，超过则重新归一化
norm_type	float	2.0	计算范数时的p范数类型，默认为L2范数
scale_grad_by_freq	bool	False	根据小批量中单词频率缩放梯度
sparse	bool	False	是否使用稀疏梯度更新，适用于特定优化器

1.4 输入输出维度

输入维度

nn.Embedding 接收形状为 (*) 的 LongTensor，其中 * 表示任意形状：

• 最常见：(batch_size, sequence_length)
• 单样本：(sequence_length,)
• 3D输入：(batch_size, seq_len1, seq_len2)

输入要求：

• 数据类型必须是 torch.long
• 所有索引值应在 [0, num_embeddings-1] 范围内
• 超出范围的索引会引发错误

输出维度

输出形状为 (*, embedding_dim)，在输入形状基础上增加一个维度：

• 输入：(batch_size, sequence_length)
• 输出：(batch_size, sequence_length, embedding_dim)

维度示例表

输入形状	输出形状	应用场景
$(N,L)$	$(N,L,D)$	批次序列处理（NLP、语音）
$(L,)$	$(L,D)$	单序列处理
$(N,H,W)$	$(N,H,W,D)$	图像位置编码
$(N,)$	$(N,D)$	批次分类标签嵌入
$(N,L1,L2)$	$(N,L1,L2,D)$	层次序列处理

其中：

• $N$：批次大小
• $L,L1,L2$：序列长度
• $D$：embedding_dim
• $H,W$：高度和宽度

二、nn.Embedding 使用示例

2.1 最简单的嵌入层创建

import torch
import torch.nn as nn

# 创建嵌入层：词汇表大小为10，每个词用3维向量表示
embedding = nn.Embedding(num_embeddings=10, embedding_dim=3)

# 查看嵌入层权重（随机初始化）
print("嵌入层权重矩阵形状:", embedding.weight.shape)  # torch.Size([10, 3])
print("权重矩阵:\n", embedding.weight)

嵌入层权重矩阵形状: torch.Size([10, 3])
权重矩阵:
 Parameter containing:
tensor([[ 0.9731, -1.2337, -1.8971],
        [ 0.0623, -0.1567,  0.3442],
        [-0.0225, -0.1363,  0.1420],
        [-0.1450, -0.3207,  0.6558],
        [ 1.1185, -1.9263, -1.3967],
        [-1.6611,  0.5666, -0.7993],
        [-1.6170, -0.2339, -1.2161],
        [-0.2325, -0.4294, -0.6188],
        [ 0.5447, -0.9562,  1.5816],
        [ 0.8761,  1.2522, -0.8475]], requires_grad=True)

2.2 处理批量句子

import torch
import torch.nn as nn

# 创建嵌入层：词汇表大小为10，每个词用5维向量表示
embedding = nn.Embedding(num_embeddings=10, embedding_dim=5)

# 批量句子处理示例
sentences = [
    "I love NLP",  # 3个词
    "Natural language",  # 2个词
    "Deep learning is fun"  # 4个词
]

# 构建词汇表
all_words = set()  # 使用集合来存储所有唯一的词
for sentence in sentences:
    all_words.update(sentence.split())  # 将每个句子中的词分割并添加到集合中

print('all_words:', all_words)  # 打印所有唯一的词
vocab = {'<PAD>': 0, '<UNK>': 1}  # 初始化词汇表，包含填充和未知标记
for i, word in enumerate(sorted(all_words), 2):  # 从2开始为每个唯一词分配索引
    vocab[word] = i  # 将词和其对应的索引添加到词汇表中

print('vocab:', vocab)  # 打印构建的词汇表

# 创建嵌入层
vocab_size = len(vocab)  # 计算词汇表的大小
embedding_layer = nn.Embedding(vocab_size, 16, padding_idx=0)  # 创建嵌入层，设置填充索引为0

# 将句子转换为索引序列（填充到相同长度）
max_len = max(len(s.split()) for s in sentences)  # 计算所有句子的最大长度
batch_indices = []  # 初始化一个列表，用于存储每个句子的索引序列

for sentence in sentences:
    words = sentence.split()  # 将句子分割成单词
    # 将每个单词转换为其对应的索引，如果单词不在词汇表中，则使用'<UNK>'的索引
    indices = [vocab.get(word, vocab['<UNK>']) for word in words]
    # 填充到最大长度，使用'<PAD>'的索引
    indices += [vocab['<PAD>']] * (max_len - len(words))
    batch_indices.append(indices)  # 将索引序列添加到批次索引列表中

# 将索引序列转换为张量
batch_tensor = torch.tensor(batch_indices)
print(f"批量句子索引:\n{batch_tensor}")  # 打印批量句子的索引
print(f"输入形状: {batch_tensor.shape}")  # 打印输入的形状

# 获取嵌入
batch_embeddings = embedding_layer(batch_tensor)  # 使用嵌入层获取批量的嵌入表示
print(f"批量嵌入形状: {batch_embeddings.shape}")  # 打印批量嵌入的形状

all_words: {'is', 'I', 'love', 'Deep', 'learning', 'fun', 'Natural', 'NLP', 'language'}
vocab: {'<PAD>': 0, '<UNK>': 1, 'Deep': 2, 'I': 3, 'NLP': 4, 'Natural': 5, 'fun': 6, 'is': 7, 'language': 8, 'learning': 9, 'love': 10}
批量句子索引:
tensor([[ 3, 10,  4,  0],
        [ 5,  8,  0,  0],
        [ 2,  9,  7,  6]])
输入形状: torch.Size([3, 4])
批量嵌入形状: torch.Size([3, 4, 16])

2.3 简单的文本分类模型

import matplotlib
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from sklearn.manifold import TSNE

matplotlib.rcParams['axes.unicode_minus'] = False
matplotlib.rcParams['font.family'] = 'Kaiti SC'
words = [
    "快乐", "愉快", "兴奋", "欢快", "高兴", "满足", "满意", "激动", "欣喜", "乐观",  # 正面
    "悲伤", "生气", "无聊", "沮丧", "失望", "烦恼", "恼火", "抑郁", "痛苦", "阴郁",  # 负面
    "中立", "还好", "不错", "一般", "平庸", "无所谓", "马马虎虎", "公平", "过得去", "可接受"  # 中性
]

labels = [
    2, 2, 2, 2, 2, 2, 2, 2, 2, 2,  # 正面
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0,  # 负面
    1, 1, 1, 1, 1, 1, 1, 1, 1, 1  # 中性
]

# 构建词汇表
vocab = {word: idx for idx, word in enumerate(words)}

# 创建嵌入层
embedding_dim = 4
embedding_layer = nn.Embedding(len(vocab), embedding_dim)

# 将单词转换为索引
word_indices = torch.tensor([vocab[word] for word in words])


# 构建简单的情感分类模型
class SimpleClassifier(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(SimpleClassifier, self).__init__()
        self.fc = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        return self.fc(x)


# 初始化模型
model = SimpleClassifier(input_dim=embedding_dim, output_dim=3)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(
    list(embedding_layer.parameters()) + list(model.parameters()),
    lr=0.01
)

# 训练模型
num_epochs = 100
for epoch in range(num_epochs):
    optimizer.zero_grad()
    word_embeddings = embedding_layer(word_indices)
    outputs = model(word_embeddings)
    loss = criterion(outputs, torch.tensor(labels))
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 20 == 0:
        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

# 获取最终嵌入向量
with torch.no_grad():
    final_embeddings = embedding_layer(word_indices)

# 使用t-SNE降维到2D
tsne = TSNE(n_components=2, random_state=42, perplexity=3, init='pca', learning_rate=200)
embeddings_2d = tsne.fit_transform(final_embeddings.numpy())

# 设置颜色映射
sentiment_colors = {0: 'red', 1: 'blue', 2: 'green'}
sentiment_names = {0: '负面', 1: '中性', 2: '正面'}
colors = [sentiment_colors[label] for label in labels]

# 绘制t-SNE降维图
plt.figure(figsize=(12, 9))
plt.scatter(embeddings_2d[:, 0], embeddings_2d[:, 1], c=colors, s=150, alpha=0.8, edgecolors='black')

# 添加单词标签
for i, word in enumerate(words):
    plt.annotate(word,
                 (embeddings_2d[i, 0], embeddings_2d[i, 1]),
                 fontsize=12, ha='center', va='bottom', fontweight='bold')

plt.title('Word Embeddings Visualization with t-SNE')
plt.xlabel('t-SNE Dimension 1')
plt.ylabel('t-SNE Dimension 2')
plt.grid(alpha=0.3)

# 添加图例
legend_elements = [plt.Line2D([0], [0], marker='o', color='w',
                              markerfacecolor=sentiment_colors[i],
                              markersize=10, label=sentiment_names[i])
                   for i in [0, 1, 2]]
plt.legend(handles=legend_elements, title="情感标签")

plt.tight_layout()
plt.show()