python - 仅卡住 torch.nn.Embedding 对象的某些行

Question

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

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

我知道为了防止 nn.Embedding 被训练，我需要将其 requires_grad 变量设置为 False 。我还发现了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)