在PyTorch中,针对词向量有一个专门的层nn.Embedding,用来实现词与词向量的映射。
nn.Embedding具有一个权重(.weight),形状是(vocab_size, embedding_dim)
Embedding层的输入形状是batch_size × seq_len(seq_len是输入序列长度, token个数),输出层的形状是batch_size × seq_len × embedding_dim。(embedding_dim等于Transformer模型维度d_model)
输入必须是LongTensor,FloatTensor需通过tensor.long()方法转成LongTensor。
Embedding的权重是可以训练的,既可以采用随机初始化,也可以采用预训练好的词向量初始化。
nn.Embedding(vocab_size, emb_size)
# vocab_size:词典大小(不是每个batch size的长度,而是数据集词库的大小)
# emb_size:每个词需要嵌入多少维来表示
import torch
from torch import nn
# 构造一个词典大小为10,每个词嵌入维度为3的嵌入权重矩阵
# embedding参数是随机初始化的,可以网络训练,可以用训练好的词向量初始化
embedding = nn.Embedding(10, 3)
print("embedding.weight:\n",embedding.weight)
'''
Parameter containing:
tensor([[-1.6941e-03, -5.0713e-01, -2.0273e+00],
[ 5.6720e-01, 2.2476e-01, -7.2537e-01],
[-1.1460e+00, -9.9649e-01, 1.3192e-01],
[-1.2278e-01, 1.2373e+00, 7.7873e-01],
[ 7.4545e-01, 1.0781e+00, 3.1572e-01],
[ 1.0267e+00, 9.3396e-01, -1.3193e+00],
[ 2.6370e-01, 7.7447e-01, -8.3768e-01],
[-1.3007e+00, -5.6405e-01, 4.6887e-01],
[-4.0872e-01, 7.9012e-02, 1.0753e-01],
[-7.9542e-02, -5.5294e-01, 1.3220e+00]], requires_grad=True)
'''
# 输入两个句子,每个句子包含4个token,输入的是toekn的序号id
# 输入的id大于9会报错,因为词典大小为10
input = torch.LongTensor([[1,2,4,5],[4,3,2,9]])
output = embedding(input)
print("output:\n", output)
'''
tensor([[[ 0.5672, 0.2248, -0.7254],
[-1.1460, -0.9965, 0.1319],
[ 0.7455, 1.0781, 0.3157],
[ 1.0267, 0.9340, -1.3193]],
[[ 0.7455, 1.0781, 0.3157],
[-0.1228, 1.2373, 0.7787],
[-1.1460, -0.9965, 0.1319],
[-0.0795, -0.5529, 1.3220]]], grad_fn=)
'''