火炬中的单次多维索引 - 也许使用 index_select 或收集？

Posted 2023-03-29

【中文标题】火炬中的单次多维索引 - 也许使用 index_select 或收集？

【英文标题】：Single shot multi-dimension indexing in torch - perhaps with index_select or gather?

【发布时间】：2020-11-13 20:14:03

【问题描述】：

我正在根据其对应数据对矩阵进行多索引重新排列。现在，我用一对 index_select 调用来做这件事，但这是非常低效的内存（n^2 就内存使用而言），并且在计算效率方面也不是很理想。有什么方法可以将我的操作简化为一个 .gather 或 .index_select 调用？

我本质上想要做的是，当给定一个形状为 (I,J,K) 的源数组和一个形状为 (I,J,2) 的索引数组时，产生一个满足条件的结果：

result[i][j][:] = source[idx[i][j][0]] [idx[i][j][1]] [:]

这是一个可运行的玩具示例，说明我现在是如何做事的：

source = torch.tensor([[1,2,3], [4,5,6], [7,8,9], [10,11,12]])
indices = torch.tensor([[[2,2],[3,1],[0,2]],[[0,2],[0,1],[0,2]],[[0,2],[0,1],[0,2]],[[0,2],[0,1],[0,2]]])

ax1 = torch.index_select(source,0,indices[:,:,0].flatten())
ax2 = torch.index_select(ax1, 1, indices[:,:,1].flatten())

result = ax2.diagonal().reshape(indices.shape(0), indices.shape(1))

这种方法对我有用，只是因为我的图像相当小，因此即使存在对角化问题，它们也能放入内存中。无论如何，我正在生成大量不需要的数据。此外，如果 K 变大，那么这个问题会成倍地恶化。也许我只是在文档中遗漏了一些明显的东西，但我觉得这是其他人必须遇到的问题才能帮助我！

【参考方案1】：

你已经为integer array indexing 准备好了索引，所以我们可以简单地做

result = source[indices[..., 0], indices[..., 1], ...]