PyTorch中的nn.Embedding应用详解

PyTorch

---

---

前言

在深度学习中，词嵌入（Word Embedding）是一种常见的技术，用于将离散的词汇或符号映射到连续的向量空间。这种映射使得相似的词汇在向量空间中具有相似的向量表示，从而可以捕捉词汇之间的语义关系。在PyTorch中，nn.Embedding模块提供了一种简单而高效的方式来实现词嵌入。

一、nn.Embedding的基本原理

nn.Embedding是一个存储固定大小的词典的嵌入向量的查找表。给定一个编号，嵌入层能够返回该编号对应的嵌入向量。这些嵌入向量反映了各个编号代表的符号之间的语义关系。在输入一个编号列表时，nn.Embedding会输出对应的符号嵌入向量列表。

在内部，nn.Embedding实际上是一个参数化的查找表，其中每一行都对应一个符号的嵌入向量。这些嵌入向量在训练过程中通过反向传播算法进行更新，以优化模型的性能。因此，nn.Embedding不仅可以用于降低数据的维度，减少计算和存储开销，还可以通过训练学习输入数据中的语义或结构信息。

二、nn.Embedding的实际应用

简单的例子

import torch
from torch.nn import Embeddingclass Model(torch.nn.Module):def __init__(self):super(Model, self).__init__()self.emb = Embedding(5, 3)def forward(self,vec):input = torch.tensor([0, 1, 2, 3, 4])emb_vec1 = self.emb(input)# print(emb_vec1)  ### 输出对同一组词汇的编码output = torch.einsum('ik, kj -> ij', emb_vec1, vec)return output
def simple_train():model = Model()vec = torch.randn((3, 1))label = torch.Tensor(5, 1).fill_(3)loss_fun = torch.nn.MSELoss()opt = torch.optim.SGD(model.parameters(), lr=0.015)print('初始化emebding参数权重:\n',model.emb.weight)for iter_num in range(100):output = model(vec)loss = loss_fun(output, label)opt.zero_grad()loss.backward(retain_graph=True)opt.step()# print('第{}次迭代emebding参数权重{}:\n'.format(iter_num, model.emb.weight))print('训练后emebding参数权重:\n',model.emb.weight)torch.save(model.state_dict(),'./embeding.pth')return modeldef simple_test():model = Model()ckpt = torch.load('./embeding.pth')model.load_state_dict(ckpt)model=model.eval()vec = torch.randn((3, 1))print('加载emebding参数权重:\n', model.emb.weight)for iter_num in range(100):output = model(vec)print('n次预测后emebding参数权重:\n', model.emb.weight)if __name__ == '__main__':simple_train()  # 训练与保存权重simple_test()

训练代码

import torch
from torch.nn import Embeddingclass Model(torch.nn.Module):def __init__(self):super(Model, self).__init__()self.emb = Embedding(5, 10)def forward(self,vec):input = torch.tensor([0, 1, 2, 3, 4])emb_vec1 = self.emb(input)print(emb_vec1)  ### 输出对同一组词汇的编码output = torch.einsum('ik, kj -> ij', emb_vec1, vec)print(output)return outputdef simple_train():model = Model()vec = torch.randn((10, 1))label = torch.Tensor(5, 1).fill_(3)print(label)loss_fun = torch.nn.MSELoss()opt = torch.optim.SGD(model.parameters(), lr=0.015)for iter_num in range(1):output = model(vec)loss = loss_fun(output, label)print('iter:%d loss:%.2f' % (iter_num, loss))opt.zero_grad()loss.backward(retain_graph=True)opt.step()if __name__ == '__main__':simple_train()

自然语言处理任务

在自然语言处理任务中，词嵌入是一种非常有用的技术。通过将每个单词表示为一个实数向量，我们可以将高维的词汇空间映射到一个低维的连续向量空间。这有助于提高模型的泛化能力和计算效率。例如，在文本分类任务中，我们可以使用nn.Embedding将文本中的每个单词转换为嵌入向量，然后将这些向量输入到神经网络中进行分类。

以下是一个简单的示例代码，演示了如何在PyTorch中使用nn.Embedding进行文本分类：

import torch
import torch.nn as nn
# 定义词嵌入层，词典大小为10000，嵌入向量维度为128
embedding = nn.Embedding(num_embeddings=10000, embedding_dim=128)
# 假设我们有一个包含5个单词的文本，每个单词的编号分别为1, 2, 3, 4, 5
input_ids = torch.tensor([1, 2, 3, 4, 5], dtype=torch.long)
# 通过词嵌入层将单词编号转换为嵌入向量
embedded = embedding(input_ids)
# 输出嵌入向量的形状：(5, 128)
print(embedded.shape)
# 定义神经网络模型
class TextClassifier(nn.Module):def __init__(self, embedding_dim, hidden_dim, num_classes):super(TextClassifier, self).__init__()self.embedding = embeddingself.fc1 = nn.Linear(embedding_dim, hidden_dim)self.fc2 = nn.Linear(hidden_dim, num_classes)def forward(self, input_ids):embedded = self.embedding(input_ids)# 对嵌入向量进行平均池化，得到一个固定长度的向量表示整个文本pooled = embedded.mean(dim=0)# 通过全连接层进行分类logits = self.fc2(self.fc1(pooled))return logits
# 实例化模型并进行训练...

上述代码中，我们首先定义了一个词嵌入层embedding，词典大小为10000，嵌入向量维度为128。然后，我们创建了一个包含5个单词的文本，每个单词的编号分别为1到5。通过调用embedding(input_ids)，我们将单词编号转换为嵌入向量。最后，我们定义了一个文本分类器模型TextClassifier，其中包含了词嵌入层、全连接层等组件。在模型的前向传播过程中，我们首先对嵌入向量进行平均池化，得到一个固定长度的向量表示整个文本，然后通过全连接层进行分类。

除了自然语言处理任务外，nn.Embedding还可以用于图像处理任务。例如，在卷积神经网络（CNN）中，嵌入层可以将图像的像素值映射到一个高维的空间，从而更好地捕捉图像中的复杂特征和结构。这有助于提高模型的性能和泛化能力。

需要注意的是，在图像处理任务中，我们通常使用卷积层（nn.Conv2d）或像素嵌入层（nn.PixelEmbed）等模块来处理图像数据，而不是直接使用nn.Embedding。