pytorch中nn.Embedding理解

給輸入的詞建立詞向量

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

word={'hello':0,'world':1}
print(word)
embeds = nn.Embedding(2, 6)
hello_idx = torch.LongTensor([word['hello']])
hello_idx = Variable(hello_idx)
hello_embed = embeds(hello_idx)
print(hello_embed)
print(hello_embed.shape)

輸出：

{'hello': 0, 'world': 1}
tensor([[-1.0645, -1.1681,  0.6972, -0.5969, -0.4473, -1.0108]],
       grad_fn=<EmbeddingBackward>)
torch.Size([1, 6])

首先：

word={'hello':0,'world':1}

是為了讓每個數字來表示一個單詞。
然後，

embeds = nn.Embedding(2, 6)

表示hello 和world兩個詞，6表示6個維度。

hello_idx = torch.LongTensor([word['hello']])
hello_idx = Variable(hello_idx)

這兩行程式碼表示得到一個variable，值是hello的index 既0。
最後print出來的就是hello的詞向量。