使用索引的二维张量访问 3D 张量（图像）

Question

使用以下表示图像的 3D 张量

img.shape=[H,W,F]

和表示该 img 索引的张量

indices.shape=[N,2]

例如，如果

indices = [[0,1],[5,3],...]] 我想创建一个新的形状张量new.shape=[N,F]，目前

new[k] == img[indices[k][0],indices[k][1]] 为了解决这个问题，我将两个张量展平：

    idx_flattened = idx_flattened [:,0] * (idx_flattened [:,1].max()+1) + idx_flattened[:,1]

    img = img .reshape(-1,F)

    new = img[idx_flattened ]

但我确信有更好的方法:)

这是一个完整的最小示例：

img = torch.arange(8*10*3).reshape(8,10,3)

indices = torch.tensor([[0,0],[3,0],[1,2]])

new = img[indices] <- This does not work

new = [[  0,   1,   2],[ 90,  91,  92],[ 36,  37,  38]]

想法？

Answer 1

切片会起作用img[indices[:,0], indices[:,1]]tensor([[ 0,&nbsp; 1,&nbsp; 2],&nbsp; &nbsp; &nbsp; &nbsp; [90, 91, 92],&nbsp; &nbsp; &nbsp; &nbsp; [36, 37, 38]])