pytorch的Embedding使用

webbery 2021-08-06 原文

torch.nn.Embedding存储的是形如num_embeddings*embedding_dim的矩阵，以词向量为例，num_embeddings表示词向量的个数，embedding_dim表示词向量的维度。

初始化：

它提供了从已知Tensor进行初始化的方法：nn.Embedding.from_pretrained

配合torch.from_numpy可以直接把numpy的array直接转换到Embedding

nn.Embedding.from_pretrained(torch.from_numpy(words_vector.wv.vectors))

取值：
在使用的时候，传入的是索引值张量，取出对应索引的词向量，如下所示

embedding = nn.Embedding(10, 3)
print(embedding.weight)
# 根据索引取embedding中的词向量
input = torch.LongTensor([[0,2,4,5],[4,3,2,0]])
embedding(input)

输出如下：

Parameter containing:
tensor([[-0.3226,  0.4114,  1.0047],
        [ 0.9196, -1.3295, -1.2954],
        [ 1.3443, -0.3448,  0.0851],
        [-0.2293, -1.3142, -1.0111],
        [-0.9291,  1.2002, -1.6681],
        [ 0.5507,  0.2129,  0.7609],
        [-0.3079, -1.5352, -0.0675],
        [ 0.8036, -0.2572,  0.4783],
        [-1.2597, -0.1978, -1.1519],
        [-0.7035, -0.0925,  0.1286]], requires_grad=True)

tensor([[[-0.3226,  0.4114,  1.0047],
         [ 1.3443, -0.3448,  0.0851],
         [-0.9291,  1.2002, -1.6681],
         [ 0.5507,  0.2129,  0.7609]],

        [[-0.9291,  1.2002, -1.6681],
         [-0.2293, -1.3142, -1.0111],
         [ 1.3443, -0.3448,  0.0851],
         [-0.3226,  0.4114,  1.0047]]], grad_fn=<EmbeddingBackward>)

对于通过gensim构造的词向量，如fastText，可以通过词来获取索引：