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),会把词向量也加载在一起。

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