pytorch 对索引子集进行矢量化答案

【问题标题】：pytorch vectorizing over index subsetpytorch 对索引子集进行矢量化
【发布时间】：2021-11-02 01:46:12
【问题描述】：

我想将一个涉及子集化到另一个张量的操作向量化。我不确定这是否可以完成，或者在这种情况下是否需要 for 循环。一个带有 for 循环的简化示例是：

for i in range(1,result.shape[0]):
        result[i,:] = other_tensor[ list_of_subset_indices[i], :].prod(axis=0)

“list_of_subset_indices”对象是用于子集张量的索引列表列表。有什么办法可以去掉 i 上的循环吗？

【问题讨论】：

【解决方案1】：

torch.index_select(other_tensor, 0, list_of_subset_indices) 做你想做的事吗？

【讨论】：

这不起作用，因为 torch.index_select 需要 list_of_subset_indices 成为张量。在我的情况下，它的列表不一定都是相同的长度。例如list_of_subset_indices = [[1,2,3],[2,4],[1,5]]

【解决方案2】：

你的操作已经被向量化了，你可以使用一个范围来索引你的赋值：

r = range(1, len(result))
result[r] = other_tensor[list_of_subset_indices[r]].prod(0)

两个侧面问题：

您确定1 在您的范围内，它不会超过index=0？
你的输出是重复数据，你确定它是你想要的二维张量吗？

在您的代码中，other_tensor[ list_of_subset_indices[i], :].prod(axis=0) 是一个单一的标量，您将其复制到 result 的所有行中。

【讨论】：

我的list_of_subset_indices 对象是一个列表列表，不接受范围作为输入。对于附带问题，(1) 是的，第一种情况 (index=0) 不同，(2) other_tensor[ list_of_subset_indices[i], :].prod(axis=0) 最终成为 1 阶张量，而不是标量。产品仅在第一个轴上。