pytorch的Embedding層

class torch.nn.Embedding(num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=2, scale_grad_by_freq=False, sparse=False)

num_embeddings (int) - 嵌入字典的大小

embedding_dim (int) - 每個嵌入向量的大小

padding_idx (int, optional) - 如果提供的話,輸出遇到此下標時用零填充

max_norm (float, optional) - 如果提供的話,會重新歸一化詞嵌入,使它們的範數小於提供的值

norm_type (float, optional) - 對於max_norm選項計算p範數時的p

scale_grad_by_freq (boolean, optional) - 如果提供的話,會根據字典中單詞頻率縮放梯度

weight weight (Tensor) -形狀爲(num_embeddings, embedding_dim)的模塊中可學習的權值

輸入: LongTensor (N, W), N = mini-batch, W = 每個mini-batch中提取的下標數
輸出: (N, W, embedding_dim)

# -*- encoding: utf-8 -*-
import torch
from torch import nn

# padding_idx:指定輸出中,如果遇到這個padding_idx則用0填充
# ebedding層的參數是默認是N(0,1)初始化的,均值爲0,方差爲1
embedding = nn.Embedding(10, 4, padding_idx=9)

# embedding的輸入要求必需是torch.long類型的,[batch_size, seq_len]
inputs = torch.tensor([[1, 2, 0, 9, 8]], dtype=torch.long)

# 輸出是[batch_size, seq_len, embedding_size]
outputs = embedding(inputs)
print(outputs)

# tensor([[[ 0.5328, -0.3801,  0.1498,  0.0932],
#          [-0.4715,  0.0643, -0.2485, -0.7041],
#          [ 0.0345,  0.1276,  1.9376, -0.5003],
#          [ 0.0000,  0.0000,  0.0000,  0.0000],
#          [-0.9140, -0.3401,  1.4773, -0.5758]]], grad_fn=<EmbeddingBackward>)

如果使用預訓練的詞向量的話,WordEembedding層沒有必要訓練的,

# emb_weight是預訓練好的詞向量
self.word_embedding = Embedding(self.vocab_size, self.embedding_dim,
                                            padding_idx=padding_idx, _weight=emb_weight)
                                       
self.word_embedding.weight.requires_grad = False

當保存模型時,torch.save(model.state_dict(), save_path),是會把詞向量一起保存了,
當重新使用加載時,saved_state_dict = torch.load(save_path),
model.load(saved_state_dict, strict=True),會把詞向量也加載在一起。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章