pytorch中的nn.Embedding

直接看代碼：

import torch
import torch.nn as nn
embedding=nn.Embedding(10,3)
input=torch.LongTensor([[1,2,4,5],[4,3,2,9]])
embedding(input)

tensor([[[ 0.8052, -0.1044, -0.6971],
         [ 1.3792, -0.1265, -1.1444],
         [ 1.4152, -0.1551, -1.2433],
         [ 0.7060, -1.0585,  0.5130]],

        [[ 1.4152, -0.1551, -1.2433],
         [-0.9881, -0.1601,  0.6339],
         [ 1.3792, -0.1265, -1.1444],
         [-1.1703,  1.8496,  0.8113]]], grad_fn=<EmbeddingBackward>)

第一個參數是字的總數，第二個參數是字的向量表示的維度。

我們的輸入input是兩個句子，每個句子都是由四個字組成的，使用每個字的索引來表示，於是使用nn.Embedding對輸入進行編碼，每個字都會編碼成長度為3的向量。

再看看下個例子：

embedding = nn.Embedding(10, 3, padding_idx=0)
input=torch.LongTensor([[0,2,0,5]])
embedding(input)

tensor([[[ 0.0000,  0.0000,  0.0000],
         [ 0.0829,  1.4141,  0.0277],
         [ 0.0000,  0.0000,  0.0000],
         [ 0.1337, -1.1472,  0.2182]]], grad_fn=<EmbeddingBackward>)

transformer中的字的編碼就可以這么表示：

class Embeddings(nn.Module):
  def __init__(self,d_model,vocab):
    #d_model=512, vocab=當前語言的詞表大小
    super(Embeddings,self).__init__()
    self.lut=nn.Embedding(vocab,d_model) 
    # one-hot轉詞嵌入，這里有一個待訓練的矩陣E，大小是vocab*d_model
    self.d_model=d_model # 512
  def forward(self,x): 
     # x ~ (batch.size, sequence.length, one-hot), 
     #one-hot大小=vocab，當前語言的詞表大小
     return self.lut(x)*math.sqrt(self.d_model) 
     # 得到的10*512詞嵌入矩陣，主動乘以sqrt(512)=22.6，
     #這里我做了一些對比，感覺這個乘以sqrt(512)沒啥用… 求反駁。
     #這里的輸出的tensor大小類似於(batch.size, sequence.length, 512)

 

參考：

https://zhuanlan.zhihu.com/p/107889011

https://blog.csdn.net/qq_38883844/article/details/104331382