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

Question

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

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

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

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

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

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

有什么办法可以做到吗？

Answer 1

如果您查看nn.Embedding 的实现，它在前向传递中使用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)