深入理解 PyTorch 的 nn.Embedding：词向量映射及变量 weight 的更新机制

前言

在深度学习领域，特别是在自然语言处理（NLP）中，nn.Embedding 是一个非常重要的模块，用于将离散的词汇（如单词或标记）映射为连续的向量表示。本文详细讲解了 nn.Embedding 的使用方法、其权重是否会在模型更新过程中被更新的问题，以及如何控制这些权重是否参与训练。

一、直接使用 nn.Embedding 获得变量

1、典型场景

这种用法通常用于处理离散的词汇表（如单词、token等），将这些离散的 token 映射为连续的向量表示。例如，在 NLP 任务中，输入是一批句子或标记序列，每个标记都有一个唯一的索引（ID）。通过 nn.Embedding，可以将这些索引映射为对应的词向量。

2、示例代码：

import torch
import torch.nn as nn

# 假设词汇表大小为 10，每个词嵌入维度为 5
vocab_size = 10
embedding_dim = 5

# 创建 Embedding 层
embedding_layer = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)

# 输入是一个批次的 token 索引
input_tokens = torch.tensor([2, 3, 5])  # 示例输入索引
embedded_vectors = embedding_layer(input_tokens)  # 获取词向量

print(embedded_vectors)

3、特点

• nn.Embedding 是一个可训练的参数层。
• 输入是离散的 token 索引，输出是对应的连续向量表示。
• 这种用法适用于需要批量处理 token 的场景，比如文本分类、机器翻译等任务。

二、使用 iou_token = nn.Embedding(1, transformer_dim) 并访问 iou_token.weight

1、典型场景

这种用法通常用于定义一些特殊的、全局共享的向量，而不是处理整个词汇表中的 token。常见的例子包括在目标检测任务中，定义一个可学习的 “特殊 token” 来表示某些特定的对象或区域（如 IoU 预测中的 token）。

2、示例代码：

import torch
import torch.nn as nn

# 定义一个特殊的 token，维度为 transformer_dim
transformer_dim = 64
iou_token = nn.Embedding(num_embeddings=1, embedding_dim=transformer_dim)

# 访问这个特殊 token 的权重
special_token_vector = iou_token.weight  # 形状为 [1, transformer_dim]

print("Special Token Vector:", special_token_vector)

3、特点

• iou_token 是一个 nn.Embedding 实例，但它的词汇表大小为 1（即只有一个 token）。
• iou_token.weight 是这个特殊 token 的实际值，形状为 [1, embedding_dim]。
• 这种用法适用于需要定义一个可学习的、全局共享的向量的场景，而不是处理多个离散 token。

三、第一种方法在模型更新中会更新其值吗？

1、默认行为

默认情况下，nn.Embedding 的权重（即词向量）是模型的可训练参数，默认情况下会被优化器更新。

2、示例代码：

import torch
import torch.nn as nn
import torch.optim as optim

# 创建一个 Embedding 层
vocab_size = 10
embedding_dim = 5
embedding_layer = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)

# 定义输入和目标
input_tokens = torch.tensor([2, 3, 5])  # 输入 token 索引
target = torch.randn(3, embedding_dim)  # 假设的目标向量

# 定义优化器
optimizer = optim.SGD(embedding_layer.parameters(), lr=0.01)

# 前向传播
embedded_vectors = embedding_layer(input_tokens)

# 计算损失
loss_fn = nn.MSELoss()
loss = loss_fn(embedded_vectors, target)

# 反向传播和更新
optimizer.zero_grad()
loss.backward()
optimizer.step()

# 查看更新后的权重
print("Updated Embedding Weights:", embedding_layer.weight)

3、控制权重更新的方法

有时我们希望固定某些权重，不让它们参与训练。这可以通过以下方式实现：

方法 1：设置 requires_grad = False

将 embedding_layer.weight.requires_grad 设置为 False，可以阻止这些权重被更新。

embedding_layer.weight.requires_grad = False

方法 2：加载预训练权重并冻结

如果我们使用预训练的词向量（如 GloVe 或 Word2Vec），可以选择加载这些权重并冻结它们。

# 加载预训练权重
pretrained_weights = torch.load('glove_embeddings.pth')

# 创建 Embedding 层并加载权重
embedding_layer = nn.Embedding.from_pretrained(pretrained_weights, freeze=True)

方法 3：在优化器中排除某些参数

我们可以在定义优化器时，排除某些参数，从而避免更新它们。

# 排除 embedding_layer 的权重
optimizer = optim.SGD(
    [param for param in model.parameters() if param is not embedding_layer.weight],
    lr=0.01
)

四、总结

• 默认情况下，nn.Embedding 的权重是可训练的，会在每次反向传播后被更新。
• 如果需要固定权重，可以通过设置 requires_grad = False、使用 from_pretrained 并设置 freeze=True 或在优化器中排除这些参数来实现。
• 选择是否更新权重取决于任务需求：如果你希望模型从头学习词向量（如随机初始化的场景），让权重可训练；如果你使用预训练的词向量并希望保持它们不变，则固定权重。