仅冻结 torch.nn.Embedding 对象的某些行

Question

我是 Pytorch 的新手，我正在尝试在嵌入上实现一种 "post-training" 过程。

我有一个包含一组项目的词汇表，我已经为每个项目学习了一个向量。我将学习到的向量保存在 nn.Embedding 对象中。我现在想做的是在不更新已经学过的向量的情况下向词汇表中添加一个新项目。新项目的嵌入将随机初始化，然后在保持所有其他嵌入冻结的情况下进行训练。

我知道为了防止 nn.Embedding 被训练，我需要将其 requires_grad 变量设置为 False。我还发现与我的相似。最佳答案建议

要么将冻结的向量和要训练的向量存储在不同的nn.Embedding对象中，前者用requires_grad = False，后者用requires_grad = True
或者将冻结的向量和新向量存储在同一个 nn.Embedding 对象中，计算所有向量的梯度，但仅在向量的维度上下降新物品。然而，这会导致相关的性能下降（当然，我想避免这种情况）。

我的问题是我确实需要将新项目的矢量存储在与旧项目的冻结矢量相同的 nn.Embedding 对象中。此约束的原因如下：在使用项目（旧的和新的）的嵌入构建我的损失函数时，我需要根据项目的 id 查找向量，并且出于性能原因我需要使用 Python 切片。换句话说，给定项目 ID 列表 item_ids，我需要做类似 vecs = embedding[item_ids] 的事情。如果我对旧项目和新项目使用两个不同的 nn.Embedding 项目，我将需要使用带有 if-else 条件的显式 for 循环，这会导致更差的性能。

有什么办法可以做到这一点吗？

Answer 1

如果你看看前向传递中 nn.Embedding it uses the functional form of embedding 的实现。因此，我认为您可以实现一个自定义模块来执行如下操作：

import torch
from torch.nn.parameter import Parameter
import torch.nn.functional as F

weights_freeze = torch.rand(10, 5)  # Don't make parameter
weights_train = Parameter(torch.rand(2, 5))
weights = torch.cat((weights_freeze, weights_train), 0)

idx = torch.tensor([[11, 1, 3]])
lookup = F.embedding(idx, weights)

# Desired result
print(lookup)
lookup.sum().backward()
# 11 corresponds to idx 1 in weights_train so this has grad
print(weights_train.grad)

仅冻结 torch.nn.Embedding 对象的某些行

Freeze only some lines of a torch.nn.Embedding object

python

pytorch

gradient-descent

embedding

word-embedding