在 PyTorch 中為嵌入層分配自定義權重

Question

PyTorch 的nn.Embedding是否支持僅為特定值手動設置嵌入權重？

我知道我可以像這樣設置整個嵌入層的權重 -

emb_layer = nn.Embedding(num_embeddings, embedding_dim)
emb_layer.weights = torch.nn.Parameter(torch.from_numpy(weight_matrix))

但是，PyTorch 是否提供任何簡潔/有效的方法來設置僅為一個特定值的嵌入權重？

像emb_layer.set_weight(5) = torch.tensor([...])這樣的東西來手動設置值“5”的嵌入？

Answer 1

是的。 您可以運行emb_layer.weight.shape來查看權重的形狀，然后您可以像這樣訪問和更改單個權重，例如：

with torch.no_grad():
  emb_layer.weight[idx_1,idx_2] = some_value

我在這里使用兩個索引，因為嵌入層是二維的。 一些層，如線性層，只需要一個索引。