向量化 Pytorch 张量索引操作答案

【问题标题】：Vectorize Pytorch Tensor Indexing Operation向量化 Pytorch 张量索引操作
【发布时间】：2020-12-06 16:56:45
【问题描述】：

我正在尝试对 PyTorch 中的操作进行矢量化处理，但我不知道该怎么做。这是现在使用 for 循环的代码。 'm' 是一个以 int 键和一维张量为值的字典。输出掩码是 2d。 L 是层数，这个循环可能是必要的。因此，我希望主要替换 2 个内部循环。我正在考虑以某种方式使用 torch.gather，但没有成功

for l in range(L):
    mask = torch.zeros((m[l].shape[0], m[l-1].shape[0]))
    for i in range(m[l].shape[0]):
        for j in range(m[l-1].shape[0]):
            mask[i,j] = R[m[l-1][j], m[l][i]]
    masks.append(mask)

如果有任何帮助，我将不胜感激！提前致谢。

【问题讨论】：

标签： python numpy pytorch vectorization

【解决方案1】：

我想我自己已经找到了答案。可以使用此处描述的 numpy 高级索引：https://docs.scipy.org/doc/numpy-1.13.0/reference/arrays.indexing.html

然后归结为以下内容

for l in range(L):
    mask = R[m[l-1], m[l][:, np.newaxis]]
    masks.append(mask)

np.newaxis 确保每行重复列索引。

【讨论】：