nn.Embedding

在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=)
'''

你可能感兴趣的:(自然语言处理,大模型,embedding)