使用2d索引张量访问3D张量(图像)

时间:2020-08-22 06:02:29

标签: python pytorch

使用以下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]]

想法?

1 个答案:

答案 0 :(得分:1)

切片会起作用

img[indices[:,0], indices[:,1]]
tensor([[ 0,  1,  2],
        [90, 91, 92],
        [36, 37, 38]])