在Pytorch中torch.nn.Embedding层的作用很简单,它的本质就是一张查找表。如下代码所示:
import torch
import torch.nn as nn
torch.manual_seed(1)
embedding = nn.Embedding(5, 3)
print(embedding.weight)
打印出的结果为:
tensor([[ 0.6614, 0.2669, 0.0617],
[ 0.6213, -0.4519, -0.1661],
[-1.5228, 0.3817, -1.0276],
[-0.5631, -0.8923, -0.0583],
[-0.1955, -0.9656, 0.4224]], requires_grad=True)
我们可以查找第2行和第3行的数据,代码如下所示:
input = torch.LongTensor([1, 2])
result = embedding(input)
print(result)
结果输出为:
tensor([[ 0.6213, -0.4519, -0.1661],
[-1.5228, 0.3817, -1.0276]], grad_fn=<EmbeddingBackward0>)
当然,我们可以对这个查找表进行特殊操作,例如重置第一行为0,代码如下所示:
import torch
import torch.nn as nn
torch.manual_seed(1)
embedding = nn.Embedding(5, 3, padding_idx=0)
print(embedding.weight)
打印出的结果为:
tensor([[ 0.0000, 0.0000, 0.0000],
[ 0.6213, -0.4519, -0.1661],
[-1.5228, 0.3817, -1.0276],
[-0.5631, -0.8923, -0.0583],
[-0.1955, -0.9656, 0.4224]], requires_grad=True)
备注:本文参考自《PyTorch面试精华》,这是国内最好的AI电子书。