PyTorch中如何用两个索引向量从矩阵中提取对应元素组成新向量

你可以直接使用PyTorch的**高级索引（Advanced Indexing）**特性实现需求，不需要写任何循环。

方法1：直接双张量索引（最简便） #

PyTorch原生支持传入两个形状相同的1D张量分别作为行、列索引，会逐位置提取对应元素，返回的张量形状和输入索引张量完全一致：

import torch

# 你的输入张量
x = torch.tensor([ 0,  1,  0,  0,  2,  2,  3,  3,  0,  0,  2,  2,  4,  4,  0,  5,  0,  4,
         6,  7,  6,  8,  6,  9,  6, 10,  5, 11,  4, 11,  8, 12, 10, 12,  9, 12,
        13, 14, 15, 16, 17, 14, 17, 18, 19, 16, 19, 18])
y = torch.tensor([ 1,  0,  2,  3,  0,  3,  2,  0,  2,  4,  0,  4,  2,  0,  5,  0,  4,  0,
         7,  6,  8,  6,  9,  6, 10,  6, 11,  5, 11,  4, 12,  8, 12, 10, 12,  9,
        14, 13, 16, 15, 14, 17, 18, 17, 16, 19, 18, 19])
mtx = torch.rand(20,20)

# 核心实现代码，仅需一行
result = mtx[x, y]

你可以自行验证：result[0]的值和mtx[0,1]完全一致，result[1]和mtx[1,0]完全一致，完全符合需求。

方法2：使用torch.gather（适用于更复杂的索引场景） #

如果后续有维度更复杂的索引需求，也可以用torch.gather接口实现，效果和上面的方法完全一致：

result = torch.gather(mtx, 1, y.unsqueeze(1))[x].squeeze()

注：常规场景下优先使用第一种方法，代码更简洁，运算效率没有差异。

内容的提问来源于stack exchange，提问作者Penguin