pytorch中的可学习查找表实现之nn.Embedding

假设我们需要一个查找表(Lookup Table),我们可以根据索引数字快速定位查找表中某个具体位置并读取出来。最简单的方法,可以通过一个二维数组或者二维list来实现。但如果我希望查找表的值可以通过梯度反向传播来修改,那么就需要用到nn.Embedding来实现了。

其实,我们需要用反向传播来修正表值的场景还是很多的,比如我们想存储数据的通用特征时,这个通用特征就可以用nn.Embedding来表示,常见于现在的各种codebook的trick。闲话不多说,我们来看栗子:

python 复制代码
import torch
from torch import nn

table = nn.Embedding(10, 3)
print(table.weight)
idx = torch.LongTensor([[1]])
b = table(idx)
print(b)

'''
output
Parameter containing:
tensor([[-0.2317, -0.9679, -1.9324],
        [ 0.2473,  1.1043, -0.7218],
        [ 0.5425, -0.3109, -0.1330],
        [-1.4006, -0.0675,  0.1376],
        [-0.1995,  0.7168,  0.5692],
        [-1.3572, -0.6407, -0.0128],
        [-0.0773,  1.1928, -1.0836],
        [ 0.1721, -0.9232, -0.4059],
        [ 1.6108, -0.4640,  0.3535],
        [ 0.6975,  1.6554, -0.2217]], requires_grad=True)
tensor([[[ 0.2473,  1.1043, -0.7218]]], grad_fn=<EmbeddingBackward0>)
'''

这段代码实际上就实现了一个查找表的功能,索引值为[[1]](注意有两个中括弧),返回值为对应的表值。我们还可以批量查找表值:

python 复制代码
import torch
from torch import nn

table = nn.Embedding(10, 3)
print(table)
print(table.weight)

indices = torch.LongTensor([[1,2,4,5],[4,3,2,9]])
print(indices)

out = table(indices)
print(out)
print(out.shape)

通过输入索引张量来获取表值:[2,4] -> [2,4,3],请注意这个shape变化,即对应位置的索引获得对应位置的表值。

参考:https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html

本人亲自整理,有问题可留言交流~

相关推荐
2401_878624794 小时前
pytorch 自动微分
人工智能·pytorch·python·机器学习
水龙吟啸5 小时前
从零开始搭建深度学习大厦系列-2.卷积神经网络基础(5-9)
人工智能·pytorch·深度学习·cnn·mxnet
慕婉03078 小时前
深度学习中的常见损失函数详解及PyTorch实现
人工智能·pytorch·深度学习
聚客AI9 小时前
搜索引擎vs向量数据库:LangChain混合检索架构实战解析
人工智能·pytorch·语言模型·自然语言处理·数据分析·gpt-3·文心一言
咸鱼鲸11 小时前
【PyTorch】PyTorch中数据准备工作(AI生成)
人工智能·pytorch·python
羊八井12 小时前
使用 Earth2Studio 和 AI 模型进行全球天气预测:太阳辐照
pytorch·python·nvidia
向左转, 向右走ˉ12 小时前
PyTorch随机擦除:提升模型抗遮挡能力
人工智能·pytorch·python·深度学习
HuashuiMu花水木1 天前
PyTorch笔记3----------统计学相关函数
人工智能·pytorch·笔记
AndrewHZ1 天前
【图像处理基石】如何检测到画面中的ppt并对其进行增强?
图像处理·人工智能·pytorch·opencv·目标检测·计算机视觉·图像增强
九章云极AladdinEdu1 天前
冷冻电镜重构的GPU加速破局:从Relion到CryoSPARC的并行重构算法
人工智能·pytorch·深度学习·机器学习·自然语言处理·架构·gpu算力