Pytorch Embedding的本质理解：查找表

在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电子书。