PyTorch中的nn.Embedding应用详解

PyTorch

文章目录

PyTorch
前言
一、nn.Embedding的基本原理
二、nn.Embedding的实际应用
- 简单的例子
- 自然语言处理任务

前言

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

一、nn.Embedding的基本原理

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

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

二、nn.Embedding的实际应用

简单的例子

python 复制代码

import torch
from torch.nn import Embedding

class 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 model

def 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()

训练代码

python 复制代码

import torch
from torch.nn import Embedding

class 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 output

def 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进行文本分类：

python 复制代码

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 = embedding
        self.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。