时间戳

首页 > 解决方案 > 如何在 PyTorch 中使用嵌入层作为线性层？

pytorch - 如何在 PyTorch 中使用嵌入层作为线性层？

问题描述

我目前正在对Transformer 论文进行个人重新实现，并且有一个问题。

在“3.4 Embeddings and Softmax”部分的第 5 页上，它指出：

在我们的模型中，我们在两个嵌入层和 pre-softmax 线性变换之间共享相同的权重矩阵。

我目前已经实现了我的模型，只对源张量和目标张量使用一个嵌入层，但我想知道是否有一种方法可以将嵌入层的权重用作线性层。我目前所做的事情是这样的：

output = previous_layer(previous_input)
final_output = torch.matmul(output, embedding_layer.embedding.weight.transpose(1, 0))

我在矩阵乘法之前转置了权重矩阵，因为它的形状(vocab_size, embedding_dim)和形状output是(batch_size, seq_len, embedding_dim)。这是将嵌入层用作线性层的正确方法吗？如果没有，我想要一些关于我应该做什么的提示。

谢谢。

编辑

来自答案链接的权重共享的特定行。

标签： pytorch

解决方案

您可以定义一个nn.Linear层并通过从以下复制权重来替换其权重nn.Embedding：

trg_emb = nn.Embedding(trg_enc_dim, embedding_dim)
src_emb = nn.Embedding(src_enc_dim, embedding_dim)

trg_projection = nn.Linear(embedding_dim, trg_enc_dim, bias=False)
trg_projection.weight = trg_emb.weight # copy to nn.Linear

您还可以复制到源嵌入，以便您的两个嵌入层也共享相同的权重。

src_emb.weight = trg_emb.weight

这意味着源嵌入最终将具有与目标嵌入相同的大小，即trg_enc_dim x embedding_dim。

这是一个可能的 PyTorch实现以获取灵感。

推荐阅读