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

Question

我是 Pytorch 的新手，我正在尝试在嵌入上实现一种“训练后”程序。

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

我知道为了防止 nn.Embedding 被训练，我需要设置False它的requires_grad变量。我还发现了与我类似的另一个问题。最佳答案建议

1. 要么存储冻结向量和要在不同的 nn.Embedding 对象中训练的向量，前者使用requires_grad = False，后者使用requires_grad = True

2. 或者将冻结的向量和新的向量存储在同一个 nn.Embedding 对象中，计算所有向量的梯度，但仅在新项的向量的维度上进行降序。然而，这会导致性能的相关下降（当然，我想避免这种情况）。

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

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

Answer 1

如果你看一下nn.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)