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),會把詞向量也加載在一起。